"""Voxel model-eval harness — one CLI for the box-layout generation test.

Processes ONE model at a time (download -> test -> optionally delete its
download -> next), so disk peak is a single model. Results merge into
report_data.js after each model, so partial progress is never lost. Grammar
enforcement (response_format json_object + schema) is always on — without it you
measure prompt-following noise, not capability (see wiki/model-selection-spike.md).

Requires llama-cpp-python, which lives in the sibling venv, NOT this repo:
  PYTHON_BIN=/Users/chenzhikai/Documents/Project/build-a-buddy-hf/.venv/bin/python
  $PYTHON_BIN tools/voxel-model-eval/eval.py --delete-after

`llama_cpp` and `huggingface_hub` are imported lazily inside functions so the
repo's pytest suite (which only touches parse_models_arg / core) does not need
the sibling venv.
"""
from __future__ import annotations

import argparse
import json
import shutil
import time
from pathlib import Path

import core
from models import MODELS

OUT = Path(__file__).resolve().parent
HUB = Path.home() / ".cache/huggingface/hub"
DATA_FILE = OUT / "report_data.js"
FEWSHOT_FILE = OUT / "report_data_fewshot.js"


def parse_models_arg(s: str):
    """'label=repo,label=repo' -> [(label, repo), ...]. Empty -> []. Whitespace-tolerant."""
    out = []
    for chunk in s.split(","):
        chunk = chunk.strip()
        if not chunk:
            continue
        label, _, repo = chunk.partition("=")
        out.append((label.strip(), repo.strip()))
    return out


def pick_q4(repo: str):
    """Resolve the Q4_K_M GGUF filename in a repo (falls back to the first gguf)."""
    from huggingface_hub import list_repo_files
    fs = [f for f in list_repo_files(repo) if f.lower().endswith(".gguf")]
    for f in fs:
        if "q4_k_m" in f.lower():
            return f
    return fs[0] if fs else None


def resolve_or_download(repo: str):
    """Return a local path to the repo's Q4_K_M GGUF, downloading if absent."""
    from huggingface_hub import hf_hub_download
    fname = pick_q4(repo)
    if not fname:
        return None
    return hf_hub_download(repo, fname)


def load_existing(path: Path):
    """Parse the window.SPIKE_DATA array out of an existing report data file."""
    if not path.exists():
        return []
    txt = path.read_text(encoding="utf-8").strip()
    s, e = txt.find("["), txt.rfind("]")
    return json.loads(txt[s:e + 1]) if s >= 0 else []


def write_data(path: Path, rows):
    path.write_text(
        "window.SPIKE_DATA = " + json.dumps(rows, ensure_ascii=False) + ";\n",
        encoding="utf-8",
    )


def merge_rows(existing, new):
    """Merge `new` result rows into `existing`, with same-model OVERWRITE:
    any existing rows whose `model` label is being re-tested are dropped, so a
    re-run updates a model in place instead of duplicating it. Rows for models
    not in `new` are kept; brand-new models append. Order: kept, then new.
    """
    new_models = {r["model"] for r in new}
    kept = [r for r in existing if r.get("model") not in new_models]
    return kept + new


def run_model(llm, label: str, themes, sys_prompt: str):
    """Run one loaded model across the themes; return result rows."""
    rows = []
    for tid, theme in themes:
        t1 = time.time()
        try:
            out = llm.create_chat_completion(
                messages=[{"role": "system", "content": sys_prompt},
                          {"role": "user", "content": theme}],
                max_tokens=1500, temperature=0.8, top_p=0.95,
                response_format={"type": "json_object", "schema": core.BOX_SCHEMA},
            )
            text = out["choices"][0]["message"]["content"]
            ctoks = out.get("usage", {}).get("completion_tokens", 0)
        except Exception as e:
            text, ctoks = f"__ERR__ {e}", 0
        gen_s = time.time() - t1
        obj = core.extract_json(text)
        boxes = obj.get("boxes", []) if isinstance(obj, dict) else []
        boxes = [b for b in boxes if isinstance(b, dict) and all(k in b for k in ("x", "y", "z", "w", "h", "d"))]
        rows.append({"model": label, "prompt": tid,
                     "name": (obj.get("name") if isinstance(obj, dict) else None) or "(unnamed)",
                     "boxes": boxes, "json_ok": bool(boxes), "gen_s": round(gen_s, 1)})
        print(f"    [{label}/{tid}] {gen_s:.1f}s tok={ctoks} json={bool(boxes)} boxes={len(boxes)}", flush=True)
    return rows


def main():
    ap = argparse.ArgumentParser(description="Voxel box-layout model eval (grammar-enforced).")
    ap.add_argument("--models", default="", help="'label=repo,...' override for models.py")
    ap.add_argument("--delete-after", action="store_true", help="remove each model's HF cache dir after testing")
    ap.add_argument("--fewshot", action="store_true", help="BASE-vs-REF exemplar A/B instead of the 8-theme sweep")
    ap.add_argument("--fresh", action="store_true", help="overwrite the data file instead of merging")
    args = ap.parse_args()

    from llama_cpp import Llama  # lazy: only needed for an actual run

    model_list = parse_models_arg(args.models) if args.models else list(MODELS)
    if not model_list:
        print("No models to test. Edit models.py or pass --models 'label=repo'.", flush=True)
        return

    if args.fewshot:
        import exemplars
        out_file = FEWSHOT_FILE
        conditions = [("BASE", core.SYS), ("REF", core.SYS + exemplars.REF_NOTE)]
        themes = exemplars.FEWSHOT_THEMES
    else:
        out_file = DATA_FILE
        conditions = [("", core.SYS)]
        themes = core.THEMES

    if args.fresh and out_file.exists():
        out_file.unlink()

    for label, repo in model_list:
        print(f"\n=== {label} ({repo}) ===", flush=True)
        try:
            path = resolve_or_download(repo)
        except Exception as e:
            print(f"    DOWNLOAD FAIL: {e}", flush=True); continue
        if not path:
            print(f"    SKIP: no gguf in {repo}", flush=True); continue
        rows = []
        try:
            llm = Llama(model_path=path, n_ctx=4096, n_gpu_layers=-1, verbose=False)
            for cond, sys_prompt in conditions:
                run_label = f"{label}-{cond}" if cond else label
                rows += run_model(llm, run_label, themes, sys_prompt)
            del llm
        except Exception as e:
            print(f"    LOAD/RUN FAIL: {e}", flush=True)
        if rows:
            existing = load_existing(out_file)
            merged = merge_rows(existing, rows)
            replaced = len(existing) + len(rows) - len(merged)
            write_data(out_file, merged)
            note = f" (replaced {replaced} stale)" if replaced else ""
            print(f"    merged +{len(rows)} -> {len(merged)} total in {out_file.name}{note}", flush=True)
        if args.delete_after:
            cache_dir = HUB / f"models--{repo.replace('/', '--')}"
            if cache_dir.exists():
                shutil.rmtree(cache_dir, ignore_errors=True)
                print(f"    reclaimed disk: removed {cache_dir.name}", flush=True)
    print("\n=== EVAL RUN DONE ===", flush=True)


if __name__ == "__main__":
    main()