#!/usr/bin/env python3
"""Mine and freeze fresh hidden-test-validated tasks for Study 5 E16."""

from __future__ import annotations

import argparse
from dataclasses import asdict
import json
from pathlib import Path

from agent_harness.specs import load_tasks
from build_study2_tasks import (
    candidate_commits,
    changed_symbols,
    load_repository,
    manifest_text,
    precheck,
    run_git,
    validate_candidate,
)


TARGETS = {"R002": 10, "R003": 7}


def select(root: Path, repository_id: str, write: bool) -> dict:
    rule = load_repository(root, repository_id)
    target = TARGETS[repository_id]
    catalog = load_tasks(root)
    excluded = {task.gold_commit for task in catalog.values()}
    selected: list[str] = []
    audits = []
    manifest_dir = root / "tasks" / "manifests"
    patch_dir = root / "tasks" / "patches"
    validation_dir = root / "tasks" / "validation" / "study5"
    selection_dir = root / "tasks" / "selection" / "study5"
    if write:
        validation_dir.mkdir(parents=True, exist_ok=True)
        selection_dir.mkdir(parents=True, exist_ok=True)

    for commit in candidate_commits(rule):
        if len(selected) >= target:
            break
        audit = precheck(rule, commit, excluded)
        audits.append(audit)
        if audit.status != "eligible":
            continue
        assert audit.parent is not None
        source_patch = str(
            run_git(rule.path, "diff", "--binary", audit.parent, commit, "--", *audit.source_paths)
        )
        test_patch = str(
            run_git(rule.path, "diff", "--binary", audit.parent, commit, "--", *audit.test_paths)
        )
        validation = validate_candidate(rule, audit, source_patch, test_patch)
        audit.validation = validation
        audit.elapsed_seconds += float(validation["elapsed_seconds"])
        if not validation["valid"]:
            audit.status = "rejected"
            audit.reason = str(validation["reason"])
            print(f"REJECT {repository_id} {commit[:12]} {audit.reason}", flush=True)
            continue

        task_id = f"TASK_S5_{repository_id}_{len(selected) + 1:03d}"
        source_name = f"{task_id}_source.patch"
        test_name = f"{task_id}_tests.patch"
        audit.task_id = task_id
        audit.status = "selected"
        audit.reason = None
        selected.append(task_id)
        excluded.add(commit)
        if write:
            text = manifest_text(
                rule,
                audit,
                task_id,
                source_name,
                test_name,
                changed_symbols(rule, audit),
            ).replace(
                "deterministic Study 2 held-out validation",
                "deterministic Study 5 fresh validation",
            )
            (manifest_dir / f"{task_id}.toml").write_text(text, encoding="utf-8")
            (patch_dir / source_name).write_text(source_patch, encoding="utf-8")
            (patch_dir / test_name).write_text(test_patch, encoding="utf-8")
            (validation_dir / f"{task_id}.json").write_text(
                json.dumps(
                    {
                        "schema_version": 1,
                        "study": "Study 5 fresh validation",
                        "task_id": task_id,
                        "repository_id": repository_id,
                        "commit": commit,
                        "parent": audit.parent,
                        "valid_end_to_end": True,
                        "validation": validation,
                    },
                    indent=2,
                    sort_keys=True,
                )
                + "\n",
                encoding="utf-8",
            )
        print(
            f"SELECT {task_id} {commit[:12]} source={len(audit.source_paths)} "
            f"tests={len(audit.test_paths)} lines={audit.source_changed_lines}",
            flush=True,
        )

    report = {
        "schema_version": 1,
        "study": "Study 5 fresh validation",
        "repository": asdict(rule) | {"path": str(rule.path)},
        "selection_policy": {
            "target": target,
            "excluded_prior_task_commits": len({task.gold_commit for task in catalog.values()}),
            "order": "reverse chronological first-parent history since 2023-01-01",
            "outcome_blind": True,
        },
        "selected_task_ids": selected,
        "selected_count": len(selected),
        "complete": len(selected) == target,
        "audits": [asdict(item) for item in audits],
    }
    if len(selected) != target:
        raise RuntimeError(f"{repository_id} yielded {len(selected)}/{target} valid tasks")
    if write:
        (selection_dir / f"{repository_id}_selection.json").write_text(
            json.dumps(report, indent=2, sort_keys=True) + "\n", encoding="utf-8"
        )
    return report


def freeze_split(root: Path) -> tuple[str, ...]:
    catalog = load_tasks(root)
    expected = tuple(
        f"TASK_S5_{repository_id}_{number:03d}"
        for repository_id, target in TARGETS.items()
        for number in range(1, target + 1)
    )
    missing = [task_id for task_id in expected if task_id not in catalog]
    if missing:
        raise RuntimeError(f"cannot freeze Study 5 split; missing {missing}")
    (root / "tasks" / "splits" / "study5_fresh.txt").write_text(
        "# Fresh Study 5 validation split; frozen before E13--E15 screening outcomes.\n"
        + "\n".join(expected)
        + "\n",
        encoding="utf-8",
    )
    return expected


def main() -> None:
    parser = argparse.ArgumentParser()
    parser.add_argument("--root", type=Path, default=Path(__file__).resolve().parents[1])
    parser.add_argument("--repository-id", choices=tuple(TARGETS))
    parser.add_argument("--write", action="store_true")
    parser.add_argument("--freeze-split", action="store_true")
    args = parser.parse_args()
    root = args.root.resolve()
    if args.freeze_split:
        split = freeze_split(root)
        print(json.dumps({"split_count": len(split), "task_ids": split}, indent=2))
        return
    if not args.repository_id:
        parser.error("--repository-id is required unless --freeze-split is used")
    report = select(root, args.repository_id, args.write)
    print(json.dumps({"repository_id": args.repository_id, "selected": report["selected_count"]}))


if __name__ == "__main__":
    main()