#!/usr/bin/env python3
"""Curate source prompts and build deterministic purpose-classifier splits."""

from __future__ import annotations

import argparse
import hashlib
import json
import sys
from pathlib import Path
from typing import Any, Sequence

from purpose_data import (
    HARD_SLICES,
    LABELS,
    CurationResult,
    DataError,
    SourceRecord,
    curate_records,
    distribution,
    exclude_reviewed_duplicates,
    file_sha256,
    jsonl_bytes,
    load_classifiable_fixtures,
    load_sources,
    prompt_hash,
    split_records,
    write_json,
    write_jsonl,
)
from review_contract import apply_completed_human_review


SCRIPT_DIR = Path(__file__).resolve().parent
REPOSITORY_ROOT = SCRIPT_DIR.parent.parent
DATA_DIR = SCRIPT_DIR / "data"
GENERATION_MANIFEST = DATA_DIR / "generation-manifest.json"
DEFAULT_FIXTURES = (
    REPOSITORY_ROOT
    / "Tests"
    / "NucleicCoreTests"
    / "Fixtures"
    / "purpose-prompts.json"
)
DEFAULT_OUTPUT_DIR = SCRIPT_DIR / ".artifacts" / "dataset-v1"
DEFAULT_FROZEN_TEST = DATA_DIR / "frozen-test-v1.jsonl"
DEFAULT_SPLIT_MANIFEST = DATA_DIR / "dataset-v1-manifest.json"
DEFAULT_CURATION_REVIEW = DATA_DIR / "curation-review-v1.json"
DATASET_VERSION = "purpose-dataset-v1"
DEFAULT_SEED = 0xC1A551F1


def _relative(path: Path) -> str:
    try:
        return str(path.resolve().relative_to(REPOSITORY_ROOT))
    except ValueError:
        return str(path.resolve())


def default_source_paths() -> list[Path]:
    try:
        value = json.loads(GENERATION_MANIFEST.read_text(encoding="utf-8"))
        names = value["canonicalFiles"]
    except (OSError, UnicodeError, json.JSONDecodeError, KeyError, TypeError) as exc:
        raise DataError(f"{GENERATION_MANIFEST}: cannot read canonicalFiles: {exc}") from exc
    if not isinstance(names, list) or not names or not all(
        isinstance(name, str) and name for name in names
    ):
        raise DataError(
            f"{GENERATION_MANIFEST}: canonicalFiles must be a non-empty string array"
        )
    return [(DATA_DIR / name).resolve() for name in names]


def _fixture_counts(path: Path) -> tuple[int, int]:
    try:
        value = json.loads(path.read_text(encoding="utf-8"))
    except (OSError, UnicodeError, json.JSONDecodeError) as exc:
        raise DataError(f"{path}: cannot read fixtures: {exc}") from exc
    if not isinstance(value, list):
        raise DataError(f"{path}: fixture root must be an array")
    classifiable = sum(
        isinstance(item, dict) and item.get("purpose") in LABELS for item in value
    )
    return len(value), classifiable


def _sha256_bytes(value: bytes) -> str:
    return hashlib.sha256(value).hexdigest()


def load_reviewed_semantic_exclusions(path: Path) -> list[dict[str, Any]]:
    try:
        value = json.loads(path.read_text(encoding="utf-8"))
        review = value["semanticDuplicateReview"]
        decisions = review["excluded"]
        status = review["status"]
    except (OSError, UnicodeError, json.JSONDecodeError, KeyError, TypeError) as exc:
        raise DataError(f"{path}: cannot read semantic exclusion review: {exc}") from exc
    if status != "complete":
        raise DataError(f"{path}: semantic duplicate review is not complete")
    if not isinstance(decisions, list):
        raise DataError(f"{path}: semantic excluded must be an array")
    return decisions


def _records(values: Sequence[SourceRecord]) -> list[dict[str, Any]]:
    return [record.value for record in values]


def _build_manifest(
    *,
    sources: Sequence[Path],
    source_record_count: int,
    curated_record_count: int,
    duplicate_counts: dict[str, int],
    splits: Any,
    fixtures_path: Path,
    total_fixture_count: int,
    classifiable_fixture_count: int,
    output_hashes: dict[str, str],
    seed: int,
    near_duplicate_threshold: float,
    curation_review_path: Path | None,
    human_review_path: Path | None,
    human_review_summary: dict[str, Any] | None,
) -> dict[str, Any]:
    manifest = {
        "schemaVersion": 1,
        "datasetVersion": DATASET_VERSION,
        "seed": seed,
        "ratios": {"train": 0.8, "validation": 0.1, "test": 0.1},
        "sources": [
            {
                "path": _relative(path),
                "sha256": file_sha256(path),
                "records": len(load_sources([path])),
            }
            for path in sources
        ],
        "curation": {
            "inputRecords": source_record_count,
            "retainedRecords": curated_record_count,
            "excludedDuplicates": duplicate_counts,
            "nearDuplicateMethod": "word-trigram Jaccard after SimHash LSH candidate search",
            "nearDuplicateThreshold": near_duplicate_threshold,
            "vagueEvalPolicy": "validation/test only",
        },
        "frozenEval": {
            "syntheticPath": _relative(DEFAULT_FROZEN_TEST),
            "syntheticSha256": output_hashes["test"],
            "shippedFixturesPath": _relative(fixtures_path),
            "shippedFixturesSha256": file_sha256(fixtures_path),
            "shippedFixtureRecords": total_fixture_count,
            "classifiableShippedFixtureRecords": classifiable_fixture_count,
            "excludedGeneralFixtureRecords": (
                total_fixture_count - classifiable_fixture_count
            ),
            "hardSliceDefinition": sorted(HARD_SLICES),
        },
        "splits": {
            "train": {
                "records": len(splits.train),
                "sha256": output_hashes["train"],
                "distribution": distribution(splits.train),
            },
            "validation": {
                "records": len(splits.validation),
                "sha256": output_hashes["validation"],
                "distribution": distribution(splits.validation),
            },
            "test": {
                "syntheticRecords": len(splits.test),
                "classifiableFixtureRecords": classifiable_fixture_count,
                "logicalRecords": splits.logical_test_count,
                "hardSyntheticRecords": sum(
                    record.value["slice"] in HARD_SLICES for record in splits.test
                ),
                "sha256": output_hashes["test"],
                "distribution": distribution(splits.test),
            },
        },
    }
    if curation_review_path is not None:
        manifest["curation"]["reviewPath"] = _relative(curation_review_path)
        manifest["curation"]["reviewSha256"] = file_sha256(curation_review_path)
    if human_review_path is not None:
        manifest["curation"]["humanReview"] = {
            "path": _relative(human_review_path),
            "sha256": file_sha256(human_review_path),
            "summary": human_review_summary,
        }
    return manifest


def prepare(
    *,
    sources: Sequence[Path],
    fixtures_path: Path,
    output_dir: Path,
    frozen_test_path: Path,
    manifest_path: Path,
    refresh_frozen_test: bool,
    seed: int,
    near_duplicate_threshold: float,
    curation_review_path: Path | None = None,
    human_review_path: Path | None = None,
) -> dict[str, Any]:
    records = load_sources(sources)
    fixtures = load_classifiable_fixtures(fixtures_path)
    total_fixture_count, classifiable_fixture_count = _fixture_counts(fixtures_path)
    if classifiable_fixture_count != len(fixtures):
        raise DataError("fixture accounting mismatch")

    lexical_curation = curate_records(
        records,
        fixtures,
        near_duplicate_threshold=near_duplicate_threshold,
    )
    reviewed_curation = (
        exclude_reviewed_duplicates(
            lexical_curation.records,
            load_reviewed_semantic_exclusions(curation_review_path),
        )
        if curation_review_path is not None
        else CurationResult(records=lexical_curation.records, duplicates=[])
    )
    human_review_summary = None
    reviewed_records = reviewed_curation.records
    if human_review_path is not None:
        human_review = apply_completed_human_review(
            human_review_path,
            reviewed_records,
            dataset_version=DATASET_VERSION,
        )
        reviewed_records = human_review.records
        human_review_summary = human_review.summary
    curated = CurationResult(
        records=reviewed_records,
        duplicates=lexical_curation.duplicates + reviewed_curation.duplicates,
    )
    splits = split_records(
        curated.records,
        fixture_count=classifiable_fixture_count,
        seed=seed,
    )

    train_values = _records(splits.train)
    validation_values = _records(splits.validation)
    test_values = _records(splits.test)
    candidate_frozen_test = jsonl_bytes(test_values)
    if refresh_frozen_test:
        frozen_test_path.parent.mkdir(parents=True, exist_ok=True)
        frozen_test_path.write_bytes(candidate_frozen_test)
    elif not frozen_test_path.exists():
        raise DataError(
            f"{frozen_test_path}: frozen test is missing; review the candidate then run "
            "--refresh-frozen-test"
        )
    elif frozen_test_path.read_bytes() != candidate_frozen_test:
        raise DataError(
            f"{frozen_test_path}: deterministic test split changed; inspect source/seed "
            "changes and use --refresh-frozen-test only when intentionally versioning it"
        )

    output_dir.mkdir(parents=True, exist_ok=True)
    write_jsonl(output_dir / "train.jsonl", train_values)
    write_jsonl(output_dir / "validation.jsonl", validation_values)
    write_jsonl(output_dir / "test.jsonl", test_values)
    write_jsonl(
        output_dir / "exclusions.jsonl",
        (
            {
                "promptHash": prompt_hash(duplicate.dropped.value["prompt"]),
                "matchedPromptHash": duplicate.matched_prompt_hash,
                "source": _relative(duplicate.dropped.source),
                "line": duplicate.dropped.line,
                "kind": duplicate.kind,
                "similarity": round(duplicate.similarity, 6),
            }
            for duplicate in curated.duplicates
        ),
    )

    duplicate_counts: dict[str, int] = {}
    for duplicate in curated.duplicates:
        duplicate_counts[duplicate.kind] = duplicate_counts.get(duplicate.kind, 0) + 1
    output_hashes = {
        "train": _sha256_bytes(jsonl_bytes(train_values)),
        "validation": _sha256_bytes(jsonl_bytes(validation_values)),
        "test": _sha256_bytes(candidate_frozen_test),
    }
    manifest = _build_manifest(
        sources=sources,
        source_record_count=len(records),
        curated_record_count=len(curated.records),
        duplicate_counts=dict(sorted(duplicate_counts.items())),
        splits=splits,
        fixtures_path=fixtures_path,
        total_fixture_count=total_fixture_count,
        classifiable_fixture_count=classifiable_fixture_count,
        output_hashes=output_hashes,
        seed=seed,
        near_duplicate_threshold=near_duplicate_threshold,
        curation_review_path=curation_review_path,
        human_review_path=human_review_path,
        human_review_summary=human_review_summary,
    )
    # The frozen path can be overridden in tests or experiments.
    manifest["frozenEval"]["syntheticPath"] = _relative(frozen_test_path)
    write_json(output_dir / "manifest.json", manifest)
    if refresh_frozen_test:
        write_json(manifest_path, manifest)
    elif manifest_path.exists():
        existing = json.loads(manifest_path.read_text(encoding="utf-8"))
        if existing != manifest:
            raise DataError(
                f"{manifest_path}: split manifest changed; inspect and refresh the frozen "
                "test intentionally"
            )
    else:
        raise DataError(
            f"{manifest_path}: frozen split manifest is missing; use --refresh-frozen-test"
        )
    return manifest


def build_parser() -> argparse.ArgumentParser:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument(
        "--source",
        action="append",
        type=Path,
        help="canonical source JSONL; repeat for multiple files (default: generation manifest)",
    )
    parser.add_argument("--fixtures", type=Path, default=DEFAULT_FIXTURES)
    parser.add_argument("--output-dir", type=Path, default=DEFAULT_OUTPUT_DIR)
    parser.add_argument("--frozen-test", type=Path, default=DEFAULT_FROZEN_TEST)
    parser.add_argument("--manifest", type=Path, default=DEFAULT_SPLIT_MANIFEST)
    parser.add_argument(
        "--refresh-frozen-test",
        action="store_true",
        help="replace the versioned test split and manifest after intentional review",
    )
    parser.add_argument("--seed", type=int, default=DEFAULT_SEED)
    parser.add_argument("--near-duplicate-threshold", type=float, default=0.92)
    parser.add_argument(
        "--curation-review",
        type=Path,
        default=DEFAULT_CURATION_REVIEW,
        help="completed semantic duplicate review applied after lexical deduplication",
    )
    parser.add_argument(
        "--human-review",
        type=Path,
        help=(
            "completed label-and-difficulty review ledger from review_data.py; "
            "omitted until human review is complete"
        ),
    )
    return parser


def main(argv: Sequence[str] | None = None) -> int:
    args = build_parser().parse_args(argv)
    try:
        sources = (
            [path.resolve() for path in args.source]
            if args.source
            else default_source_paths()
        )
        manifest = prepare(
            sources=sources,
            fixtures_path=args.fixtures.resolve(),
            output_dir=args.output_dir.resolve(),
            frozen_test_path=args.frozen_test.resolve(),
            manifest_path=args.manifest.resolve(),
            refresh_frozen_test=args.refresh_frozen_test,
            seed=args.seed,
            near_duplicate_threshold=args.near_duplicate_threshold,
            curation_review_path=args.curation_review.resolve(),
            human_review_path=(
                args.human_review.resolve() if args.human_review is not None else None
            ),
        )
    except (DataError, OSError, UnicodeError, json.JSONDecodeError) as exc:
        print(f"error: {exc}", file=sys.stderr)
        return 1

    splits = manifest["splits"]
    print(
        f"Prepared {manifest['curation']['retainedRecords']} curated records: "
        f"{splits['train']['records']} train, "
        f"{splits['validation']['records']} validation, "
        f"{splits['test']['logicalRecords']} frozen test "
        f"({splits['test']['classifiableFixtureRecords']} shipped fixtures)."
    )
    return 0


if __name__ == "__main__":
    raise SystemExit(main())
