"""`crit eval` command line. Also runs as `python eval/crit_eval` from a crit checkout."""
import argparse
import json
from pathlib import Path
import sys


def main(argv=None):
    if hasattr(sys.stdout, "reconfigure"):
        sys.stdout.reconfigure(encoding="utf-8", errors="replace")
    parser = argparse.ArgumentParser(prog="crit eval", description="Run the common crit speed, quality and speculative-exactness suite")
    subs = parser.add_subparsers(dest="command", required=True)

    run = subs.add_parser("run", help="Evaluate one engine (prepares data and reference first when missing)")
    run.add_argument("--adapter", required=True, help="Python file exporting create_eval_engine(context)")
    run.add_argument("--model", required=True, help="Model directory passed to the adapter")
    run.add_argument("--data", required=True, help="Prepared suite and reference directory for this model")
    run.add_argument("--out", required=True, help="Result JSON path")
    run.add_argument("--config", default="{}", help="JSON configuration overrides passed to the adapter")
    run.add_argument("--only", choices=["speed", "quality", "exactness"], action="append",
                     help="Run only these parts (repeatable); default is all")
    run.add_argument("--repetitions", type=int, help="Override the protocol's measured repetitions")
    run.add_argument("--warmups", type=int, help="Override the protocol's warmup runs")
    run.add_argument("--label", help="Name shown in comparisons")
    run.add_argument("--bandwidth", type=float, metavar="GBS",
                     help="Achievable device memory bandwidth in GB/s (from a copy test) for the pass floor")

    prepare = subs.add_parser("prepare", help="Freeze the suite as token ids for one model")
    prepare.add_argument("--model", required=True)
    prepare.add_argument("--data", required=True)

    reference = subs.add_parser("reference", help="Build reference log-probabilities (torch + transformers, GPU)")
    reference.add_argument("--model", required=True)
    reference.add_argument("--data", required=True)

    compare = subs.add_parser("compare", help="Print a side-by-side table of result files")
    compare.add_argument("results", nargs="+")

    args = parser.parse_args(argv)
    if args.command == "prepare":
        from .prepare import prepare as do_prepare
        prepared = do_prepare(args.model, args.data)
        print(f"prepared {len(prepared['speed'])} speed cases and {len(prepared['quality'])} quality conversations in {args.data}")
        return 0
    if args.command == "reference":
        return _reference(args.model, args.data)
    if args.command == "compare":
        from .report import compare as do_compare
        print(do_compare(args.results))
        return 0
    from .prepare import PREPARED
    from .reference import REFERENCE
    from .runner import run as do_run
    parts = tuple(args.only or ("speed", "quality", "exactness"))
    data = Path(args.data)
    if not (data / PREPARED).is_file():
        from .prepare import prepare as do_prepare
        do_prepare(args.model, data)
    if "quality" in parts and not (data / REFERENCE).is_file():
        _reference(args.model, data)
    result = do_run(args.adapter, args.model, data, args.out, json.loads(args.config), parts,
                    args.warmups, args.repetitions, args.label, bandwidth_gbs=args.bandwidth)
    print(f"wrote {args.out}")
    speed, quality = result.get("speed"), result.get("quality")
    if speed and speed["decode_score"]:
        print(f"decode score {speed['decode_score']:.1f} tok/s (geometric mean over {len(speed['cases'])} cases)")
    if quality:
        print("quality " + ("passed" if quality["passed"] else "FAILED"))
    room = result.get("headroom")
    if room and room["pass_efficiency"] is not None:
        print(f"model pass {room['step_ms']:.2f} ms vs floor {room['pass_floor_ms']:.2f} ms "
              f"({room['pass_efficiency']:.0%}{', polished' if room['pass_polished'] else ''})")
    projection = (room or {}).get("projection") or {}
    if projection.get("speed_of_light_score") and speed and speed["decode_score"]:
        print(f"speed of light {projection['speed_of_light_score']:.1f} tok/s at draft width "
              f"{projection['best_width_at_floor']}: this engine is at {projection['fraction_of_speed_of_light']:.0%}")
    return 0


def _reference(model, data):
    from . import load_json
    from .reference import build_reference
    protocol = load_json(Path(__file__).resolve().parents[1] / "protocol.json")
    out = build_reference(model, data, protocol["quality"]["top_k"])
    print(f"reference perplexity {out['perplexity']:.4f} over {sum(len(c['targets']) for c in out['conversations'])} targets")
    return 0
