"""Command-line entry point shared by the Kai and Lex Core ML toolkits."""

from __future__ import annotations

import argparse
import json
from pathlib import Path

from huggingface_hub import snapshot_download
from typed_coreml import (
    KINDS,
    compare_native,
    export_coreml,
    load_native,
    verify_coreml,
)


def main() -> None:
    parser = argparse.ArgumentParser()
    parser.add_argument("--toolkit", required=True, type=Path)
    parser.add_argument("--native-dir", type=Path)
    parser.add_argument("--output-dir", type=Path)
    parser.add_argument("--kind", choices=KINDS)
    parser.add_argument("--tokens", type=int, default=128)
    parser.add_argument("--candidates", type=int)
    parser.add_argument("--fixture-index", type=int, default=0)
    parser.add_argument("--flexible", action="store_true")
    parser.add_argument("--marker-map", action="store_true", help="use an FP16 one-hot marker map instead of indices")
    parser.add_argument("--verify-only", action="store_true")
    parser.add_argument("--verify-limit", type=int, help="maximum pinned real rows to verify per type")
    parser.add_argument("--verify-package", type=Path)
    parser.add_argument("--verify-packages-dir", type=Path)
    args = parser.parse_args()
    if args.marker_map and args.flexible:
        parser.error("--marker-map currently requires a fixed token and candidate shape")
    if args.verify_limit is not None and args.verify_limit < 1:
        parser.error("--verify-limit must be positive")
    if args.verify_package is not None and args.kind is None:
        parser.error("--verify-package requires --kind")
    if args.verify_package is not None and args.verify_packages_dir is not None:
        parser.error("Choose one package verification mode")
    toolkit = args.toolkit.resolve()
    lock = json.loads((toolkit / "assets.lock.json").read_text())
    native_dir = args.native_dir
    if native_dir is None:
        snapshot = snapshot_download(
            repo_id=lock["source_repo"],
            revision=lock["source_revision"],
            allow_patterns=["native/**"],
        )
        native_dir = Path(snapshot) / "native"
    model, collator, _ = load_native(native_dir)
    records = []
    for fixture in ("upstream-decisions.jsonl", "upstream-system-one-rows.jsonl"):
        records.extend(json.loads(line) for line in (toolkit / fixture).read_text().splitlines() if line)
    kinds = [args.kind] if args.kind else list(KINDS)
    failed_packages = []
    for kind in kinds:
        selected = [record for record in records if record["question"]["type"].lower() == kind]
        if not selected:
            raise ValueError(f"Missing upstream {kind} fixture")
        if args.verify_only:
            print(json.dumps(compare_native(
                model, collator, selected, kind, tokens=args.tokens, candidates=args.candidates,
                marker_map=args.marker_map,
            )))
            continue
        package = args.verify_package
        if args.verify_packages_dir is not None:
            package = args.verify_packages_dir / (kind + ".mlpackage")
        if package is not None:
            report = verify_coreml(
                model, collator, selected[:args.verify_limit], kind, package
            )
            output = package.parent / (kind + "-validation.json")
            output.write_text(json.dumps(report, indent=2) + "\n")
            metadata_path = package.with_suffix(".json")
            metadata = json.loads(metadata_path.read_text())
            passed = report["choice_agreement"] and report["max_abs_probability_error"] <= 0.02
            metadata["validated"] = passed
            metadata["validation_report"] = output.name
            metadata_path.write_text(json.dumps(metadata, indent=2) + "\n")
            print(json.dumps(report))
            if not passed:
                failed_packages.append(kind)
            continue
        if args.output_dir is None:
            raise ValueError("--output-dir is required for conversion")
        if args.fixture_index < 0 or args.fixture_index >= len(selected):
            raise ValueError("--fixture-index is outside the upstream fixture set")
        suffix = "-flex" if args.flexible else "-marker-map" if args.marker_map else ""
        report = export_coreml(
            model,
            collator,
            [selected[args.fixture_index]],
            kind,
            args.output_dir / (kind + suffix + ".mlpackage"),
            source_repo=lock["source_repo"],
            source_revision=lock["source_revision"],
            tokens=args.tokens,
            candidates=args.candidates,
            flexible=args.flexible,
            marker_map=args.marker_map,
        )
        print(json.dumps(report))
    if failed_packages:
        raise SystemExit(f"Core ML native-parity gate failed for: {', '.join(failed_packages)}")


if __name__ == "__main__":
    main()