#!/usr/bin/env -S uv run --script
# /// script
# requires-python = ">=3.11"
# dependencies = [
#   "gguf>=0.10",
# ]
# ///
"""
quantize-all.py — drive transcribe-quantize across the shipped preset matrix.

Given a reference-dtype GGUF (F32, F16, or BF16 — the converter's output),
produces one quantized GGUF per preset in the matrix, skipping any preset
that equals the source dtype. Output filenames follow the llama.cpp
convention: <variant>-<PRESET>.gguf in the same directory as the input.

The default matrix is DERIVED_PRESETS, but an architecture may register a
narrower override in scripts/lib/quant_policy.py::FAMILY_PRESETS (used
when k-quant tiers would degenerate into Q8_0 for that model's shapes).
The architecture is read from the input GGUF's `general.architecture`
KV; pass --presets to override explicitly.

Usage:
    uv run scripts/quantize-all.py \\
        models/parakeet-tdt-0.6b-v2/parakeet-tdt-0.6b-v2-F32.gguf

    # Pick a subset:
    uv run scripts/quantize-all.py --presets Q8_0,Q4_K_M \\
        models/parakeet-tdt-0.6b-v2/parakeet-tdt-0.6b-v2-F32.gguf

Assumes build/bin/transcribe-quantize has been built (cmake --build build
--target transcribe-quantize). See docs/tools/quantization.md.
"""

from __future__ import annotations

import argparse
import subprocess
import sys
from pathlib import Path

# Import via path so the script stays `uv run`-compatible (no project context).
HERE = Path(__file__).resolve().parent
sys.path.insert(0, str(HERE))
from lib.quant_policy import (  # noqa: E402
    DERIVED_PRESETS,
    derived_presets_for_arch,
    validate_preset,
)


def detect_source_preset(path: Path) -> str | None:
    """Return the F32/F16/BF16 suffix embedded in the filename, or None."""
    stem = path.stem
    for tier in ("F32", "F16", "BF16"):
        if stem.endswith(f"-{tier}"):
            return tier
    return None


def read_architecture(gguf_path: Path) -> str | None:
    """Return the general.architecture string from a GGUF, or None on failure."""
    try:
        from gguf.gguf_reader import GGUFReader
    except ImportError:
        return None
    try:
        reader = GGUFReader(str(gguf_path))
        fld = reader.fields.get("general.architecture")
        if fld is None or not fld.parts:
            return None
        return bytes(fld.parts[-1]).decode("utf-8")
    except Exception:
        return None


def main(argv: list[str]) -> int:
    p = argparse.ArgumentParser(description=__doc__.splitlines()[0])
    p.add_argument("input", type=Path, help="Reference-dtype GGUF produced by convert-<family>.py")
    p.add_argument("--presets", type=str, default=None,
                   help="Comma-separated preset subset (default: per-architecture matrix from "
                        "quant_policy.FAMILY_PRESETS, falling back to DERIVED_PRESETS)")
    p.add_argument("--quantize-bin", type=Path, default=Path("build/bin/transcribe-quantize"),
                   help="Path to transcribe-quantize (default: build/bin/transcribe-quantize)")
    args = p.parse_args(argv)

    src = args.input.resolve()
    if not src.is_file():
        print(f"error: input not found: {src}", file=sys.stderr)
        return 2
    if not args.quantize_bin.is_file():
        print(f"error: transcribe-quantize not built at {args.quantize_bin}. "
              f"Run: cmake --build build --target transcribe-quantize", file=sys.stderr)
        return 2

    if args.presets is None:
        arch = read_architecture(src)
        matrix = derived_presets_for_arch(arch)
        if matrix is not DERIVED_PRESETS:
            print(f"[quantize-all] arch={arch!r}: per-family preset matrix {matrix}", flush=True)
        presets_str = ",".join(matrix)
    else:
        presets_str = args.presets
    presets = [validate_preset(x) for x in presets_str.split(",") if x.strip()]
    src_tier = detect_source_preset(src)
    stem_base = src.stem.removesuffix(f"-{src_tier}") if src_tier else src.stem

    results: list[tuple[str, bool, str]] = []
    for preset in presets:
        if preset == src_tier:
            results.append((preset, True, "skipped (equals source dtype)"))
            continue
        out = src.parent / f"{stem_base}-{preset}.gguf"
        if out.exists():
            results.append((preset, True, f"already present at {out}"))
            continue
        cmd = [str(args.quantize_bin), str(src), str(out), "--quant", preset]
        print(f"[quantize-all] {preset}: {' '.join(cmd)}", flush=True)
        rc = subprocess.call(cmd)
        ok = rc == 0
        results.append((preset, ok, str(out) if ok else f"transcribe-quantize exited {rc}"))

    print()
    print(f"{'preset':<10} {'status':<5}  detail")
    print(f"{'-'*10} {'-'*5}  {'-'*40}")
    for preset, ok, detail in results:
        print(f"{preset:<10} {('OK' if ok else 'FAIL'):<5}  {detail}")

    return 0 if all(ok for _, ok, _ in results) else 1


if __name__ == "__main__":
    raise SystemExit(main(sys.argv[1:]))