from __future__ import annotations

import argparse
import json
from pathlib import Path

from .data import DecisionRow, load_rows, validate_corpus
from .oracles import fixtures, oracle
from .teacher import dry_run_plan


def main() -> None:
    parser = argparse.ArgumentParser(description="Rook offline curation and guarded teacher utilities")
    commands = parser.add_subparsers(dest="command", required=True)
    from .native.commands import add_commands
    add_commands(commands)
    validate = commands.add_parser("validate-data")
    validate.add_argument("paths", nargs="+", type=Path, help="All partitions together, to check leakage")
    commands.add_parser("teacher-plan")
    generate = commands.add_parser("write-fixtures")
    generate.add_argument("path", type=Path)
    schema = commands.add_parser("write-schema")
    schema.add_argument("path", type=Path)
    for name in ("curate-data", "teacher-dry-run", "export-training", "apply-review", "teacher-run", "teacher-ledger-status", "finalize-pilot", "review-ui"):
        command = commands.add_parser(name)
        command.add_argument("--root", type=Path, default=Path.cwd())
        if name not in {"curate-data", "teacher-ledger-status"}:
            command.add_argument("--dataset", type=Path, default=Path("data/rook-phase2a-v1"))
        if name == "export-training":
            command.add_argument("--output", type=Path, required=True)
        if name == "apply-review":
            command.add_argument("--responses", type=Path, required=True)
            command.add_argument("--version", required=True)
            command.add_argument("--family-audits", type=Path)
        if name in {"teacher-dry-run", "teacher-run"}:
            command.add_argument("--pilot", type=Path)
        if name == "finalize-pilot":
            command.add_argument("--config", type=Path, default=Path("configs/teacher.phase2b-small.json"))
            command.add_argument("--count", type=int, default=5)
            command.add_argument("--version", required=True)
        if name == "review-ui":
            command.add_argument("--store", type=Path, default=Path("runs/review.sqlite3"))
            command.add_argument("--port", type=int, default=8765)
        if name in {"teacher-run", "teacher-ledger-status"}:
            command.add_argument("--ledger", type=Path, default=Path("runs/teacher.sqlite3"))
        if name == "teacher-run":
            command.add_argument("--authorization", type=Path, required=True)
    args = parser.parse_args()
    if args.command.startswith("native-"):
        from .native.commands import run
        try:
            result = run(args)
        except (ValueError, RuntimeError, OSError, ImportError, TimeoutError) as exc:
            parser.exit(2, f"Native command stopped: {exc}\n")
        print(json.dumps({"status": result.get("status", "recorded"), "output": str(args.output),
                          "check": result.get("check"), "training_rows": result.get("training_rows"),
                          "cuda_checks": result.get("cuda_checks")}, ensure_ascii=False))
        return
    if args.command == "validate-data":
        try:
            rows = load_rows(args.paths)
            for row in rows:
                if row.provenance.label_method == "program_oracle" and row.source_id == "rook-original-fixtures":
                    if row.options[row.label].key != oracle(row.provenance.oracle_id, row.state):
                        raise ValueError(f"{row.id}: label disagrees with program oracle")
        except (ValueError, TypeError) as exc:
            parser.exit(2, f"Validation failed: {exc}\n")
        print(json.dumps({"status": "valid", "rows": len(rows), "splits": sorted({r.split for r in rows})}))
    elif args.command == "teacher-plan":
        print(json.dumps(dry_run_plan(), indent=2))
    elif args.command == "write-fixtures":
        rows = fixtures()
        validate_corpus(rows)
        args.path.parent.mkdir(parents=True, exist_ok=True)
        args.path.write_text("".join(r.model_dump_json(exclude_none=True) + "\n" for r in rows), encoding="utf-8")
    elif args.command == "write-schema":
        value = DecisionRow.model_json_schema()
        value["$schema"] = "https://json-schema.org/draft/2020-12/schema"
        value["$comment"] = "Use rook validate-data for label bounds, mappings, provenance and whole-corpus isolation; schema alone cannot enforce these."
        args.path.write_text(json.dumps(value, ensure_ascii=False, indent=2) + "\n", encoding="utf-8")
    else:
        from .curation.build import build
        from .curation.commands import export_training, load_pilot, run_pilot
        from .curation.review import apply_batch
        from .teacher import BudgetExceeded
        from .teacher_ledger import DurableLedger, OwnershipError
        root = args.root.resolve()
        dataset = (root / args.dataset).resolve() if hasattr(args, "dataset") else None
        try:
            if args.command == "curate-data":
                build(root)
                return
            if args.command == "teacher-dry-run":
                _, _, _, result = load_pilot(root, dataset, root / args.pilot if args.pilot else None)
            elif args.command == "export-training":
                result = export_training(root, dataset, args.output)
            elif args.command == "apply-review":
                result = apply_batch(root, dataset, args.responses, args.version, args.family_audits)
            elif args.command == "finalize-pilot":
                from .curation.finalize import finalize_pilot
                result = finalize_pilot(root, dataset, root / args.config, args.version, args.count)
            elif args.command == "review-ui":
                import uvicorn
                from .review_app import create_review_app
                uvicorn.run(create_review_app(root, dataset, root / args.store), host="127.0.0.1", port=args.port, workers=1)
                return
            elif args.command == "teacher-run":
                result = run_pilot(root, dataset, args.authorization, root / args.ledger, root / args.pilot if args.pilot else None)
            else:
                with DurableLedger(root / args.ledger) as ledger:
                    result = ledger.snapshot()
        except (ValueError, PermissionError, OwnershipError, BudgetExceeded, OSError) as exc:
            parser.exit(2, f"Command stopped: {exc}\n")
        print(json.dumps(result, ensure_ascii=False, indent=2))