#!/usr/bin/env python3
"""Loop check of the release files under candidate chat settings (28 Oct, for the card and a PocketPal Pal).

72 prompts shaped like the 19 Oct router sweep's (its loop-prone file, fp8check/aya_loopy4.jsonl, is gone, so the
40 ur/tr/bn/ne prompts are drawn from bench/aya_marco29_prompts.jsonl: same languages, not the same prompts), plus 16
in other languages and the 16 JSON prompts; each file
through its own llama-server on a spare port (compare_quality.serve: refuses an occupied port and checks /props names
the file; never the router), thinking off, the same seed per prompt. A loop is an output that hit max_tokens with a
16-gram repeat share >= 0.5 over its last 600 tokens (the sweep's rule).

    python scripts/sampling_check.py --model NAME=PATH.gguf ... [--ngl 0] [--port 18196]
"""
from __future__ import annotations

import argparse
import random
import json
import sys
from concurrent.futures import ThreadPoolExecutor
from pathlib import Path

sys.path.insert(0, str(Path(__file__).resolve().parent))
import compare_quality as cq  # noqa: E402
from fp8_check import loop_score  # noqa: E402
from router_sampling_sweep import CONFIGS, DATA  # noqa: E402

REPO = Path(__file__).resolve().parents[1]
CHECK = ["spec_t0.7_p0.9", "t0.7_rep1.1", "t0.3_p0.9"]  # a default-like setting, the candidate, a low temperature


def prompts() -> list[dict]:
    pool = [json.loads(line) for line in open(DATA / "bench/aya_marco29_prompts.jsonl")]
    rng = random.Random(0)
    out = []
    for lang in ("ur", "tr", "bn", "ne"):
        out += rng.sample([r for r in pool if r["language"] == lang], 10)
    for lang in ("en", "zh", "de", "ar", "ja", "ru", "es", "ko"):
        out += rng.sample([r for r in pool if r["language"] == lang], 2)
    out += [json.loads(line) for line in open(DATA / "bench/json_prompts.jsonl")]
    return out


def group(i: int) -> str:
    return "loop-prone" if i < 40 else "json" if i >= 56 else "other"


def main() -> None:
    ap = argparse.ArgumentParser()
    ap.add_argument("--server", default="/path/to/llama.cpp/build-ptq/bin/llama-server")
    ap.add_argument("--model", action="append", required=True, help="NAME=PATH.gguf")
    ap.add_argument("--out-dir", type=Path, default=REPO / "reports" / "sampling_check")
    ap.add_argument("--port", type=int, default=18196)
    ap.add_argument("--ngl", default="0")
    ap.add_argument("--max-tokens", type=int, default=1536)
    ap.add_argument("--workers", type=int, default=8)
    a = ap.parse_args()
    from tokenizers import Tokenizer
    tok = Tokenizer.from_file(str(DATA / "models/Marco-Nano-Instruct/tokenizer.json"))
    ps = prompts()
    a.out_dir.mkdir(parents=True, exist_ok=True)
    for spec in a.model:
        name, gguf = spec.split("=", 1)
        out = a.out_dir / f"{name}.jsonl"
        done = {(r["config"], r["i"]) for r in map(json.loads, out.open())} if out.exists() else set()
        todo = [(c, i, r) for c in CHECK for i, r in enumerate(ps) if (c, i) not in done]
        if todo:
            p = cq.serve(a.server, gguf, a.port, a.ngl, a.out_dir / f"{name}.server.log")
            try:
                def run(t):
                    c, i, r = t
                    body = {"messages": [{"role": "user", "content": r["text"]}], "max_tokens": a.max_tokens,
                            "seed": 1000 + i, "chat_template_kwargs": {"enable_thinking": False}, **CONFIGS[c]}
                    try:
                        return t, cq.post(a.port, "/v1/chat/completions", body)
                    except Exception as e:  # noqa: BLE001 — a server error (e.g. a reply cut inside a UTF-8 character)
                        return t, {"choices": [{"finish_reason": "error", "message": {"content": f"[{e}]"}}]}
                with ThreadPoolExecutor(a.workers) as ex, out.open("a") as f:
                    for (c, i, r), res in ex.map(run, todo):
                        ch = res["choices"][0]
                        text = ch["message"].get("content") or ""
                        ids = tok.encode(text, add_special_tokens=False).ids
                        f.write(json.dumps({"config": c, "i": i, "lang": r["language"], "finish": ch["finish_reason"],
                                            "n_tokens": len(ids), "loop_score": round(loop_score(ids), 3), "text": text},
                                           ensure_ascii=False) + "\n")
                        f.flush()
            finally:
                p.terminate()
                p.wait(60)
        rows = [json.loads(line) for line in out.open()]
        summ = {}
        for c in CHECK:
            for g in ("loop-prone", "other", "json"):
                rs = [r for r in rows if r["config"] == c and group(r["i"]) == g]
                cut = sum(r["finish"] == "length" for r in rs)
                loops = sum(r["finish"] == "length" and r["loop_score"] >= 0.5 for r in rs)
                summ[f"{c} / {g}"] = {"prompts": len(rs), "cut": cut, "loops": loops}
        (a.out_dir / f"{name}.summary.json").write_text(json.dumps(summ, indent=1))
        print(name, json.dumps({k: f"{v['cut']} cut / {v['loops']} loops of {v['prompts']}" for k, v in summ.items()}),
              flush=True)


if __name__ == "__main__":
    main()