#!/usr/bin/env python3
"""Build a training-only dataset augmented with labeled first prompts."""

from __future__ import annotations

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

from purpose_data import (
    DataError,
    SourceRecord,
    curate_records,
    distribution,
    file_sha256,
    jsonl_bytes,
    load_classifiable_fixtures,
    load_jsonl,
    normalized_key,
    prompt_hash,
    validate_source_record,
    write_json,
    write_jsonl,
)


SCRIPT_DIR = Path(__file__).resolve().parent
REPOSITORY_ROOT = SCRIPT_DIR.parent.parent
DEFAULT_BASE_DATASET = SCRIPT_DIR / ".artifacts" / "dataset-v1"
DEFAULT_HISTORY = (
    SCRIPT_DIR / ".artifacts" / "nucleic-history-first-prompts.labeled.jsonl"
)
DEFAULT_FIXTURES = (
    REPOSITORY_ROOT
    / "Tests"
    / "NucleicCoreTests"
    / "Fixtures"
    / "purpose-prompts.json"
)
DEFAULT_OUTPUT = SCRIPT_DIR / ".artifacts" / "dataset-v1-history-first-prompts"
EXPERIMENT_VERSION = "history-first-prompts-training-augmentation-v1"


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


def _validate_records(records: Sequence[dict[str, Any]], path: Path) -> None:
    if not records:
        raise DataError(f"{path}: split is empty")
    for line, record in enumerate(records, 1):
        validate_source_record(record, f"{path}:{line}")


def _indexed_records(
    records: Sequence[dict[str, Any]], path: Path
) -> list[SourceRecord]:
    return [
        SourceRecord(value=record, source=path, line=line)
        for line, record in enumerate(records, 1)
    ]


def _unique_split_keys(
    splits: Sequence[tuple[str, Sequence[dict[str, Any]]]]
) -> dict[str, tuple[str, str]]:
    seen: dict[str, tuple[str, str]] = {}
    for split_name, records in splits:
        for line, record in enumerate(records, 1):
            key = normalized_key(record["prompt"])
            previous = seen.get(key)
            if previous is not None:
                previous_split, previous_label = previous
                detail = (
                    "conflicting labels"
                    if previous_label != record["purpose"]
                    else "duplicate prompt"
                )
                raise DataError(
                    f"{split_name}:{line}: {detail} also present in {previous_split}"
                )
            seen[key] = (split_name, record["purpose"])
    return seen


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


def _exclusion(
    record: SourceRecord,
    *,
    reason: str,
    matched_prompt_hash: str | None = None,
    similarity: float | None = None,
) -> dict[str, Any]:
    value: dict[str, Any] = {
        "promptHash": prompt_hash(record.value["prompt"]),
        "sourceLine": record.line,
        "reason": reason,
    }
    if matched_prompt_hash is not None:
        value["matchedPromptHash"] = matched_prompt_hash
    if similarity is not None:
        value["similarity"] = round(similarity, 6)
    return value


def prepare(
    *,
    base_dataset: Path,
    history_path: Path,
    fixtures_path: Path,
    output_dir: Path,
    near_duplicate_threshold: float,
    overwrite_output: bool,
) -> dict[str, Any]:
    base_paths = {
        split: base_dataset / f"{split}.jsonl"
        for split in ("train", "validation", "test")
    }
    base = {split: load_jsonl(path) for split, path in base_paths.items()}
    for split, path in base_paths.items():
        _validate_records(base[split], path)

    base_keys = _unique_split_keys(
        [(split, base[split]) for split in ("train", "validation", "test")]
    )
    history = load_jsonl(history_path)
    _validate_records(history, history_path)
    fixtures = load_classifiable_fixtures(fixtures_path)

    eval_fixtures = [
        {"prompt": record["prompt"], "purpose": record["purpose"]}
        for split in ("validation", "test")
        for record in base[split]
    ] + fixtures
    eval_hashes = {prompt_hash(record["prompt"]) for record in eval_fixtures}

    eligible: list[SourceRecord] = []
    exclusions: list[dict[str, Any]] = []
    for record in _indexed_records(history, history_path):
        if record.value["slice"] == "vague-eval":
            exclusions.append(_exclusion(record, reason="vague-eval"))
            continue

        key = normalized_key(record.value["prompt"])
        previous = base_keys.get(key)
        if previous is not None:
            split_name, previous_label = previous
            if previous_label != record.value["purpose"]:
                raise DataError(
                    f"{history_path}:{record.line}: label conflicts with exact "
                    f"{split_name} prompt"
                )
            exclusions.append(
                _exclusion(
                    record,
                    reason=(
                        "exact-base-train-overlap"
                        if split_name == "train"
                        else "exact-eval-overlap"
                    ),
                    matched_prompt_hash=prompt_hash(record.value["prompt"]),
                    similarity=1.0,
                )
            )
            continue
        eligible.append(record)

    curated = curate_records(
        eligible,
        eval_fixtures,
        near_duplicate_threshold=near_duplicate_threshold,
    )
    for duplicate in curated.duplicates:
        exclusions.append(
            _exclusion(
                duplicate.dropped,
                reason=(
                    f"{duplicate.kind}-eval-overlap"
                    if duplicate.matched_prompt_hash in eval_hashes
                    else f"{duplicate.kind}-history-duplicate"
                ),
                matched_prompt_hash=duplicate.matched_prompt_hash,
                similarity=duplicate.similarity,
            )
        )

    accepted_history = [record.value for record in curated.records]
    output_train = base["train"] + accepted_history
    train_bytes = jsonl_bytes(output_train)
    validation_bytes = base_paths["validation"].read_bytes()
    test_bytes = base_paths["test"].read_bytes()

    if output_dir.exists() and any(output_dir.iterdir()):
        if not overwrite_output:
            raise DataError(
                f"{output_dir}: output is not empty; pass --overwrite-output "
                "intentionally"
            )
        shutil.rmtree(output_dir)
    output_dir.mkdir(parents=True, exist_ok=True)
    (output_dir / "train.jsonl").write_bytes(train_bytes)
    (output_dir / "validation.jsonl").write_bytes(validation_bytes)
    (output_dir / "test.jsonl").write_bytes(test_bytes)
    write_jsonl(
        output_dir / "history-exclusions.jsonl",
        sorted(exclusions, key=lambda value: value["sourceLine"]),
    )

    exclusion_counts = dict(
        sorted(Counter(value["reason"] for value in exclusions).items())
    )
    accepted_sources = _indexed_records(accepted_history, history_path)
    manifest = {
        "schemaVersion": 1,
        "experimentVersion": EXPERIMENT_VERSION,
        "policy": {
            "historyUsage": "training-only",
            "vagueEval": "excluded from optimization",
            "exactBaseTrainOverlap": "excluded",
            "exactEvaluationOverlap": "excluded",
            "nearEvaluationOverlap": "excluded",
            "nearHistoryDuplicates": "excluded",
            "nearDuplicateThreshold": near_duplicate_threshold,
            "validationAndTestBytes": "identical to base dataset",
        },
        "sources": {
            "baseDataset": {
                "path": _relative(base_dataset),
                "splits": {
                    split: {
                        "path": _relative(path),
                        "records": len(base[split]),
                        "sha256": file_sha256(path),
                    }
                    for split, path in base_paths.items()
                },
            },
            "history": {
                "path": _relative(history_path),
                "records": len(history),
                "sha256": file_sha256(history_path),
            },
            "fixtures": {
                "path": _relative(fixtures_path),
                "classifiableRecords": len(fixtures),
                "sha256": file_sha256(fixtures_path),
            },
        },
        "augmentation": {
            "inputHistoryRecords": len(history),
            "acceptedHistoryRecords": len(accepted_history),
            "excludedHistoryRecords": len(exclusions),
            "exclusions": exclusion_counts,
            "acceptedDistribution": distribution(accepted_sources),
        },
        "outputs": {
            "train": {
                "records": len(output_train),
                "sha256": _sha256_bytes(train_bytes),
            },
            "validation": {
                "records": len(base["validation"]),
                "sha256": _sha256_bytes(validation_bytes),
            },
            "test": {
                "records": len(base["test"]),
                "sha256": _sha256_bytes(test_bytes),
            },
        },
    }
    write_json(output_dir / "manifest.json", manifest)
    return manifest


def build_parser() -> argparse.ArgumentParser:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--base-dataset", type=Path, default=DEFAULT_BASE_DATASET)
    parser.add_argument("--history", type=Path, default=DEFAULT_HISTORY)
    parser.add_argument("--fixtures", type=Path, default=DEFAULT_FIXTURES)
    parser.add_argument("--output-dir", type=Path, default=DEFAULT_OUTPUT)
    parser.add_argument("--near-duplicate-threshold", type=float, default=0.92)
    parser.add_argument("--overwrite-output", action="store_true")
    return parser


def main(argv: Sequence[str] | None = None) -> int:
    args = build_parser().parse_args(argv)
    try:
        manifest = prepare(
            base_dataset=args.base_dataset.resolve(),
            history_path=args.history.resolve(),
            fixtures_path=args.fixtures.resolve(),
            output_dir=args.output_dir.resolve(),
            near_duplicate_threshold=args.near_duplicate_threshold,
            overwrite_output=args.overwrite_output,
        )
    except (DataError, OSError, UnicodeError, json.JSONDecodeError) as exc:
        print(f"error: {exc}", file=sys.stderr)
        return 1

    augmentation = manifest["augmentation"]
    outputs = manifest["outputs"]
    print(
        f"Prepared {outputs['train']['records']} training records "
        f"({augmentation['acceptedHistoryRecords']} from history); validation and "
        f"test remain byte-identical to the base dataset."
    )
    return 0


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