#!/usr/bin/env python3
"""Compare GGUF token IDs with the pinned Hugging Face tokenizer outputs."""

from __future__ import annotations

import argparse
import ast
import json
import subprocess
from pathlib import Path


def main() -> None:
    parser = argparse.ArgumentParser()
    parser.add_argument("--llama-tokenize", type=Path, required=True)
    parser.add_argument("--model", type=Path, required=True)
    parser.add_argument("--hf-reference", type=Path, required=True)
    parser.add_argument("--output", type=Path, required=True)
    args = parser.parse_args()

    reference = json.loads(args.hf_reference.read_text(encoding="utf-8"))
    cases = reference.get("tokenizer_cases")
    if not isinstance(cases, list) or len(cases) != 6:
        raise SystemExit("HF reference must contain exactly six tokenizer cases")
    if any(
        not isinstance(case, dict)
        or not isinstance(case.get("text"), str)
        or not case["text"]
        or not isinstance(case.get("token_ids"), list)
        or not case["token_ids"]
        for case in cases
    ):
        raise SystemExit("HF reference contains an empty or malformed tokenizer case")

    comparisons = []
    for case in cases:
        completed = subprocess.run(
            [
                str(args.llama_tokenize),
                "--model",
                str(args.model),
                "--prompt",
                case["text"],
                "--ids",
                "--no-bos",
                "--offline",
                "--log-disable",
            ],
            check=True,
            text=True,
            capture_output=True,
        )
        gguf_ids = ast.literal_eval(completed.stdout.strip())
        match = gguf_ids == case["token_ids"]
        comparisons.append(
            {
                "text": case["text"],
                "hf_token_ids": case["token_ids"],
                "gguf_token_ids": gguf_ids,
                "exact_match": match,
            }
        )

    result = {
        "cases": comparisons,
        "case_count": len(comparisons),
        "all_exact": all(case["exact_match"] for case in comparisons),
    }
    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 not result["all_exact"]:
        raise SystemExit("tokenizer mismatch")


if __name__ == "__main__":
    main()