#!/usr/bin/env python3
"""Convert this repository's Hugging Face checkpoint to Q4_K_M GGUF.

The tokenizer is a small custom ByteLevel BPE.  Its behavior is GPT-2-style,
but its tokenizer probe hash is not yet present in llama.cpp's generated
pre-tokenizer lookup table.  This wrapper validates the tokenizer structure
before supplying the corresponding ``gpt-2`` pre-tokenizer identifier.
"""

from __future__ import annotations

import argparse
import json
import runpy
import shutil
import subprocess
import sys
from pathlib import Path


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument(
        "--llama-cpp",
        type=Path,
        required=True,
        help="Path to a llama.cpp checkout containing convert_hf_to_gguf.py",
    )
    quantizer = shutil.which("llama-quantize")
    parser.add_argument(
        "--quantizer",
        type=Path,
        default=Path(quantizer) if quantizer else None,
        help="Path to llama-quantize (default: resolve it from PATH)",
    )
    parser.add_argument(
        "--outfile",
        type=Path,
        default=Path("gguf/tinygemma3-2m-Q4_K_M.gguf"),
    )
    return parser.parse_args()


def is_expected_bytelevel_bpe(model_dir: Path) -> bool:
    tokenizer = json.loads((model_dir / "tokenizer.json").read_text())
    return (
        tokenizer.get("normalizer") is None
        and tokenizer.get("model", {}).get("type") == "BPE"
        and tokenizer.get("model", {}).get("byte_fallback") is False
        and tokenizer.get("pre_tokenizer")
        == {
            "type": "ByteLevel",
            "add_prefix_space": False,
            "trim_offsets": False,
            "use_regex": True,
        }
    )


def main() -> None:
    args = parse_args()
    repo_dir = Path(__file__).resolve().parent
    model_dir = repo_dir / "hf"
    converter = args.llama_cpp.resolve() / "convert_hf_to_gguf.py"
    quantizer = args.quantizer.resolve() if args.quantizer else None
    outfile = args.outfile if args.outfile.is_absolute() else repo_dir / args.outfile
    intermediate = outfile.parent / ".tinygemma3-2m-f16.intermediate.gguf"

    if not converter.is_file():
        raise SystemExit(f"llama.cpp converter not found: {converter}")
    if quantizer is None or not quantizer.is_file():
        raise SystemExit("llama-quantize not found; pass it with --quantizer")
    if not is_expected_bytelevel_bpe(model_dir):
        raise SystemExit("Unexpected tokenizer structure; refusing to guess GGUF metadata")

    sys.path.insert(0, str(args.llama_cpp.resolve()))
    from conversion import TextModel  # noqa: PLC0415
    from conversion.gemma import Gemma3Model  # noqa: PLC0415

    original = TextModel.get_vocab_base_pre

    def get_vocab_base_pre(self: TextModel, tokenizer: object) -> str:
        if Path(self.dir_model).resolve() == model_dir.resolve():
            return "gpt-2"
        return original(self, tokenizer)

    def modify_tensors(
        self: Gemma3Model, data_torch: object, name: str, bid: int | None
    ) -> object:
        # This checkpoint intentionally pads the embedding matrix and logits
        # from 1,003 tokenizer entries to config.vocab_size=1,024.  Keep those
        # rows so the GGUF architecture and HF reference logits stay aligned.
        f_shift = self.norm_shift(name)
        if f_shift != 0.0:
            data_torch = data_torch + f_shift
        yield from super(Gemma3Model, self).modify_tensors(data_torch, name, bid)

    TextModel.get_vocab_base_pre = get_vocab_base_pre
    Gemma3Model.modify_tensors = modify_tensors
    outfile.parent.mkdir(parents=True, exist_ok=True)
    sys.argv = [
        str(converter),
        str(model_dir),
        "--outfile",
        str(intermediate),
        "--outtype",
        "f16",
        "--model-name",
        "tinygemma3-2m",
    ]
    try:
        runpy.run_path(str(converter), run_name="__main__")
        subprocess.run(
            [str(quantizer), str(intermediate), str(outfile), "Q4_K_M"],
            check=True,
        )
    finally:
        intermediate.unlink(missing_ok=True)


if __name__ == "__main__":
    main()