#!/usr/bin/env python3
"""Compare a deterministic llama-completion output with the HF BF16 reference."""

from __future__ import annotations

import argparse
import json
from pathlib import Path


def main() -> None:
    parser = argparse.ArgumentParser()
    parser.add_argument("--artifact", required=True)
    parser.add_argument("--hf-reference", type=Path, required=True)
    parser.add_argument("--gguf-output", type=Path, required=True)
    parser.add_argument("--case", default="greedy_capitals_12")
    parser.add_argument("--prompt-text", required=True)
    parser.add_argument(
        "--mismatch-policy",
        choices=("fail", "report"),
        default="fail",
        help="Whether a valid but different quantized generation is fatal.",
    )
    parser.add_argument("--output", type=Path, required=True)
    args = parser.parse_args()

    reference = json.loads(args.hf_reference.read_text(encoding="utf-8"))
    reference_case = reference["generations"][args.case]
    expected = reference_case["new_text"]
    expected_prompt = reference_case["prompt"]
    raw = args.gguf_output.read_text(encoding="utf-8")
    transport_suffix_present = raw.endswith("\n\n")
    actual = raw[:-2] if transport_suffix_present else raw
    execution_failures = []
    if args.prompt_text != expected_prompt:
        execution_failures.append("llama.cpp prompt differs from HF reference prompt")
    if not transport_suffix_present:
        execution_failures.append("missing llama-completion terminal newline suffix")
    result = {
        "artifact": args.artifact,
        "case": args.case,
        "prompt": args.prompt_text,
        "prompt_exact_match": args.prompt_text == expected_prompt,
        "expected_hf_bf16_text": expected,
        "actual_gguf_text": actual,
        "exact_text_match": actual == expected,
        "transport_suffix_present": transport_suffix_present,
        "execution_failures": execution_failures,
        "mismatch_policy": args.mismatch_policy,
    }
    args.output.parent.mkdir(parents=True, exist_ok=True)
    args.output.write_text(
        json.dumps(result, ensure_ascii=False, indent=2, sort_keys=True) + "\n",
        encoding="utf-8",
    )
    print(json.dumps(result, ensure_ascii=False, sort_keys=True))
    if execution_failures:
        raise SystemExit("; ".join(execution_failures))
    if not result["exact_text_match"] and args.mismatch_policy == "fail":
        raise SystemExit("greedy generation differs from HF BF16 reference")


if __name__ == "__main__":
    main()