"""Quantize the fine-tuned model to GGUF for llama.cpp (🦙 Llama Champion · 🔌 Off the Grid).

Runs entirely on Modal (CPU) so we never download the 16 GB model locally:
  1. build llama.cpp (convert script + llama-quantize + llama-cli)
  2. pull the merged fine-tuned model from the Hub
  3. convert HF → GGUF f16
  4. quantize to Q4_K_M (laptop sweet spot) and Q8_0 (high quality)
  5. verify the GGUF headers (and a best-effort inference smoke test)
  6. push the GGUF files to the Hub

Usage:
    modal run modal_app/quantize_modal.py::main
"""

from __future__ import annotations

import modal

SRC_REPO = "eldinosaur/cuentas-claras-sat-classifier-minicpm"
BASE_REPO = "openbmb/MiniCPM4.1-8B"   # source of the SentencePiece tokenizer.model
OUT_REPO = "eldinosaur/cuentas-claras-sat-classifier-gguf"
QUANTS = ["Q4_K_M", "Q8_0"]

app = modal.App("cuentas-claras-quantize")
volume = modal.Volume.from_name("cuentas-claras-gguf-vol", create_if_missing=True)
hf_secret = modal.Secret.from_name("huggingface-secret")

image = (
    modal.Image.debian_slim(python_version="3.12")
    .apt_install("git", "build-essential", "cmake", "libcurl4-openssl-dev")
    .pip_install("torch", index_url="https://download.pytorch.org/whl/cpu")
    .pip_install("transformers>=4.44", "sentencepiece", "protobuf", "numpy",
                 "safetensors", "huggingface_hub>=0.23")
    .run_commands(
        "git clone --depth 1 https://github.com/ggml-org/llama.cpp /llama.cpp",
        "pip install -e /llama.cpp/gguf-py",
        "cmake -S /llama.cpp -B /llama.cpp/build -DLLAMA_CURL=OFF -DGGML_NATIVE=OFF",
        "cmake --build /llama.cpp/build --config Release -j --target llama-quantize llama-cli",
    )
)


GGUF_README = """\
---
license: apache-2.0
base_model: eldinosaur/cuentas-claras-sat-classifier-minicpm
tags: [gguf, llama.cpp, accounting, sat, mexico]
language: [es, en]
---

# Cuentas Claras — SAT Transaction Classifier (GGUF)

GGUF quantizations of the fine-tuned MiniCPM4.1-8B SAT transaction classifier, for
local inference with **llama.cpp** (🦙 / 🔌 — no cloud, runs on a laptop).

| File | Quant | ~Size | Use |
|---|---|---|---|
| `cuentas-claras-sat-Q4_K_M.gguf` | Q4_K_M | ~4.9 GB | laptop default |
| `cuentas-claras-sat-Q8_0.gguf` | Q8_0 | ~8.5 GB | higher quality |

## Run

```bash
llama-cli -m cuentas-claras-sat-Q4_K_M.gguf --jinja \\
  -sys "Eres un clasificador contable mexicano. Responde solo JSON." \\
  -p 'Clasifica: "Suscripción anual a Adobe Creative Cloud"'
```

Returns the SAT account, deductibility and IVA treatment as JSON. See the base model
card for training details (eval_loss 0.155, token accuracy 96.1%).

> Not tax advice. Account codes/rules are simplified and must be verified vs the SAT catalogue.
"""


@app.function(image=image, secrets=[hf_secret], volumes={"/work": volume},
              timeout=60 * 60, cpu=8.0, memory=32768)
def quantize():
    import shutil
    import subprocess
    from pathlib import Path
    from huggingface_hub import HfApi, hf_hub_download, snapshot_download

    work = Path("/work")
    model_dir = work / "model"
    out_dir = work / "gguf"
    out_dir.mkdir(exist_ok=True)

    print(f"Downloading {SRC_REPO} …")
    snapshot_download(SRC_REPO, local_dir=str(model_dir),
                      ignore_patterns=["*.pt", "*.bin"])  # safetensors only

    # Our fine-tuned repo ships tokenizer.json but not the SentencePiece
    # tokenizer.model that llama.cpp's MiniCPM converter requires. The vocab is
    # unchanged from the base, so fetch it from there.
    if not (model_dir / "tokenizer.model").exists():
        print("Fetching tokenizer.model from base model …")
        tm = hf_hub_download(BASE_REPO, "tokenizer.model")
        shutil.copy(tm, model_dir / "tokenizer.model")

    f16 = out_dir / "model-f16.gguf"
    print("Converting HF → GGUF f16 …")
    subprocess.run(
        ["python3", "/llama.cpp/convert_hf_to_gguf.py", str(model_dir),
         "--outfile", str(f16), "--outtype", "f16"],
        check=True,
    )

    produced = {}
    for qt in QUANTS:
        dst = out_dir / f"cuentas-claras-sat-{qt}.gguf"
        print(f"Quantizing → {qt} …")
        subprocess.run(["/llama.cpp/build/bin/llama-quantize", str(f16), str(dst), qt],
                       check=True)
        produced[qt] = dst

    # --- verify GGUF headers ---
    import gguf
    for qt, path in produced.items():
        reader = gguf.GGUFReader(str(path))
        arch = reader.fields.get("general.architecture")
        arch_val = bytes(arch.parts[arch.data[0]]).decode() if arch else "?"
        size_gb = path.stat().st_size / 1e9
        print(f"[verify] {path.name}: arch={arch_val}, tensors={len(reader.tensors)}, "
              f"size={size_gb:.2f} GB")

    # --- best-effort inference smoke test on the Q4 model ---
    try:
        q4 = produced.get("Q4_K_M")
        if q4:
            res = subprocess.run(
                ["/llama.cpp/build/bin/llama-cli", "-m", str(q4), "--jinja", "-no-cnv",
                 "-n", "80", "-sys",
                 "Eres un clasificador contable mexicano. Responde solo JSON.",
                 "-p", 'Clasifica: "Suscripción anual a Adobe Creative Cloud"'],
                capture_output=True, text=True, timeout=300)
            print("[smoke-test output]\n", (res.stdout or "")[-600:])
    except Exception as e:
        print("[smoke-test skipped]", e)

    # --- push ---
    api = HfApi()
    api.create_repo(OUT_REPO, exist_ok=True)
    (out_dir / "README.md").write_text(GGUF_README, encoding="utf-8")
    api.upload_file(path_or_fileobj=str(out_dir / "README.md"),
                    path_in_repo="README.md", repo_id=OUT_REPO)
    for qt, path in produced.items():
        print(f"Uploading {path.name} …")
        api.upload_file(path_or_fileobj=str(path), path_in_repo=path.name, repo_id=OUT_REPO)

    url = f"https://huggingface.co/{OUT_REPO}"
    print("Done →", url)
    return url


@app.local_entrypoint()
def main():
    print("GGUF repo:", quantize.remote())