"""Quantize an f16 GGUF to a K-quant using llama-cpp-python's bundled libllama
(llama_model_quantize) — no separate llama-quantize binary / C++ build needed.

Usage:
    uv run --group train --with llama-cpp-python python scripts/quantize_gguf.py \
        data/gguf/yui-brain1-sft1.f16.gguf \
        data/gguf/yui-brain1-sft1.Q4_K_M.gguf Q4_K_M
"""

from __future__ import annotations

import sys

import llama_cpp


# ftype values from llama.cpp (llama_ftype enum); the common K-quants.
_FTYPE = {
    "Q4_K_M": llama_cpp.LLAMA_FTYPE_MOSTLY_Q4_K_M,
    "Q5_K_M": llama_cpp.LLAMA_FTYPE_MOSTLY_Q5_K_M,
    "Q8_0": llama_cpp.LLAMA_FTYPE_MOSTLY_Q8_0,
}


def main() -> None:
    src, dst, quant = sys.argv[1], sys.argv[2], sys.argv[3]
    params = llama_cpp.llama_model_quantize_default_params()
    params.ftype = _FTYPE[quant]
    rc = llama_cpp.llama_model_quantize(
        src.encode("utf-8"), dst.encode("utf-8"), params
    )
    if rc != 0:
        raise SystemExit(f"llama_model_quantize failed rc={rc}")
    print(f"[quantize] {src} -> {dst} ({quant})")


if __name__ == "__main__":
    main()