#!/usr/bin/env python3
"""T2.7 · GGUF export check: packed ternary students → llama.cpp Q2_0 (D-51: N-bit students → Q2_0 / Q4_0 /
Q3_K), round trip and greedy agreement.

    # template: llama.cpp's own conversion of the base model (metadata, tensor names and order)
    python scripts/t27_gguf.py build --template T.gguf --out S.gguf --packed <packed_dir>        # a student
    python scripts/t27_gguf.py build --template T.gguf --out S.gguf --rtn-dense <hf_dir>          # RTN Qwen3-0.6B
    python scripts/t27_gguf.py compare --gguf S.gguf (--packed <dir> | --rtn-dense <hf_dir>) --out R.json

`compare` runs both models greedy on the same prompt token ids (200 held-out prompts, chat template, ≤512 prompt
tokens, 64 new tokens, end-of-sequence suppressed in both) and reports position-wise token agreement, the first
divergence per prompt, and llama.cpp's own timings. PyTorch is the reference: the packed student dequantised to
bf16 (`load_packed_student`), on CPU. llama.cpp runs CPU-only (`llama-server`, GPU hidden: gen's GPU is busy).
"""
from __future__ import annotations

import argparse
import json
import os
import random
import subprocess
import sys
import time
import urllib.request
from pathlib import Path

import torch

REPO = Path(__file__).resolve().parent.parent
sys.path.insert(0, str(REPO / "src"))
from lowbit.export import gguf as gx  # noqa: E402
from lowbit.quant.ternary import QuantSpec, dequant_packed, export_packed  # noqa: E402

LLAMA = Path(os.environ.get("MT_LLAMA", "/path/to/llama.cpp"))  # gguf-py, the converter (spark: third_party)
# the pegfix build: mainline's chat-output parser returns HTTP 500 on some outputs (invalid UTF-8 from a degraded
# model), which the fork's UTF-8 fix removes; same source base, so the same Q2_0 kernels
SERVER_LIBS = "/path/to/llama.cpp-fork/build-pegfix/bin"
SERVER = Path(SERVER_LIBS) / "llama-server"
HELDOUT = Path("/path/to/data/prompts/v2/heldout.jsonl")


def rtn_dense(hf_dir: str):
    """A dense model with every MLP projection ternarised by round-to-nearest (Q-A, g128): (dense, packed, model)."""
    from transformers import AutoModelForCausalLM
    model = AutoModelForCausalLM.from_pretrained(hf_dir, dtype=torch.bfloat16)
    dense, packed = {}, {}
    for n, p in model.named_parameters():
        if any(n.endswith(f"mlp.{k}_proj.weight") for k in ("gate", "up", "down")):
            pk = export_packed(p.detach().float(), QuantSpec(method="absmean", group_size=128))
            packed[n[: -len(".weight")]] = pk
            with torch.no_grad():
                p.copy_(dequant_packed(pk))  # the PyTorch reference holds exactly the exported weights
        else:
            dense[n] = p.detach()
    return dense, packed, model


def packed_student(packed_dir: str):
    """(plain tensors, packed experts); D-51's N-bit students carry their method per tensor (`read_packed`)."""
    from lowbit.model.student import read_packed
    sd, packed, _ = read_packed(packed_dir)
    return sd, packed


def build(a) -> None:
    g = gx.gguf_module(LLAMA)
    if a.packed:
        dense, packed = packed_student(a.packed)
    else:
        dense, packed, _ = rtn_dense(a.rtn_dense)
    names = {t.name for t in g.GGUFReader(a.template).tensors}
    if "output.weight" in names and "lm_head.weight" not in dense:  # tied: the converter still writes output
        dense["lm_head.weight"] = dense["model.embed_tokens.weight"]
    dense = {k: v for k, v in dense.items() if gx.gguf_name(k) in names}  # a tied lm_head absent from the GGUF
    t0 = time.time()
    meta = {}
    if a.name:  # a release: its own name, and the teacher recorded as its base model
        meta = {"general.name": a.name, "general.finetune": "QAT", "general.base_model.count": 1,
                "general.base_model.0.name": "Marco Nano Instruct", "general.base_model.0.organization": "AIDC-AI",
                "general.base_model.0.repo_url": "https://huggingface.co/AIDC-AI/Marco-Nano-Instruct"}
    info = gx.write_gguf(a.template, a.out, gx.student_tensors(dense, packed, g, a.ternary_as, a.dense_as), g, meta)
    info["write_s"] = round(time.time() - t0, 1)
    if a.ternary_as == "q2_0":
        info.update(gx.check_round_trip(a.out, packed, g))
    print(json.dumps(info))


def prompts(tok, n: int, max_len: int = 512) -> list[list[int]]:
    rows = [json.loads(l) for l in open(HELDOUT)]
    random.Random(0).shuffle(rows)
    out = []
    for r in rows:
        text = tok.apply_chat_template([{"role": "user", "content": r["text"]}], add_generation_prompt=True,
                                       tokenize=False)
        ids = tok(text, add_special_tokens=False)["input_ids"]
        if len(ids) <= max_len:
            out.append(list(ids))
        if len(out) == n:
            return out
    return out


@torch.no_grad()
def torch_greedy(model, tok, ps: list[list[int]], new: int, batch: int) -> list[list[int]]:
    tok.padding_side = "left"
    pad = tok.pad_token_id if tok.pad_token_id is not None else tok.eos_token_id
    out = []
    for i in range(0, len(ps), batch):
        chunk = ps[i:i + batch]
        L = max(len(p) for p in chunk)
        ids = torch.tensor([[pad] * (L - len(p)) + p for p in chunk])
        att = torch.tensor([[0] * (L - len(p)) + [1] * len(p) for p in chunk])
        g = model.generate(input_ids=ids, attention_mask=att, max_new_tokens=new, min_new_tokens=new,
                           do_sample=False, pad_token_id=pad)
        out += [row[L:].tolist() for row in g]
        print(json.dumps({"torch_greedy": i + len(chunk), "of": len(ps)}), flush=True)
    return out


def start_server(gguf: str, threads: int, port: int) -> subprocess.Popen:
    env = {**os.environ, "CUDA_VISIBLE_DEVICES": "", "LD_LIBRARY_PATH": SERVER_LIBS}
    # --reasoning-format none: plain completions, no think-block parsing (a parser hang stalled one request)
    proc = subprocess.Popen([str(SERVER), "-m", gguf, "--port", str(port), "--host", "127.0.0.1", "-ngl", "0",
                             "-t", str(threads), "-c", "8192", "--parallel", "1", "--no-webui",
                             "--reasoning-format", "none"],
                            env=env, stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL)
    for _ in range(600):
        try:
            urllib.request.urlopen(f"http://127.0.0.1:{port}/health", timeout=2)
            return proc
        except Exception:
            if proc.poll() is not None:
                raise RuntimeError(f"llama-server exited with {proc.returncode}")
            time.sleep(1)
    raise RuntimeError("llama-server did not come up")


def stop_server(proc: subprocess.Popen) -> None:
    proc.terminate()  # it can sit on SIGTERM mid-request; a slow shutdown must not discard results (it did once)
    try:
        proc.wait(30)
    except subprocess.TimeoutExpired:
        proc.kill()
        proc.wait()


def server_greedy(gguf: str, ps: list[list[int]], new: int, threads: int, port: int,
                  timeout_s: int = 120) -> tuple[list, dict]:
    """Greedy tokens per prompt; a request over `timeout_s` is recorded as None and the server restarted."""
    proc = start_server(gguf, threads, port)
    outs, pp, tg, timeouts = [], [], [], 0
    try:
        for p in ps:
            body = {"prompt": p, "n_predict": new, "temperature": 0.0, "top_k": 1, "ignore_eos": True,
                    "cache_prompt": False, "return_tokens": True, "samplers": ["top_k"]}
            req = urllib.request.Request(f"http://127.0.0.1:{port}/completion", json.dumps(body).encode(),
                                         {"Content-Type": "application/json"})
            try:
                r = json.loads(urllib.request.urlopen(req, timeout=timeout_s).read())
                outs.append(r["tokens"])
                pp.append(r["timings"]["prompt_per_second"])
                tg.append(r["timings"]["predicted_per_second"])
            except (TimeoutError, OSError) as ex:
                outs.append(None)
                timeouts += 1
                print(json.dumps({"llama_timeout": len(outs) - 1, "error": repr(ex)[:120]}), flush=True)
                stop_server(proc)
                proc = start_server(gguf, threads, port)
            if len(outs) % 20 == 0:
                print(json.dumps({"llama_greedy": len(outs), "of": len(ps), "gen_tok_s": round(tg[-1], 1) if tg else None,
                                  "prompt_tok_s": round(pp[-1], 1) if pp else None}), flush=True)
        return outs, {"prompt_tok_s_median": sorted(pp)[len(pp) // 2], "gen_tok_s_median": sorted(tg)[len(tg) // 2],
                      "threads": threads, "timeouts": timeouts}
    finally:
        stop_server(proc)


def compare(a) -> None:
    from transformers import AutoTokenizer
    from lowbit.model import student as st
    torch.set_num_threads(a.threads)
    if a.packed:
        model = st.load_packed_student(a.packed, attn_impl="sdpa", experts_impl="grouped_mm")
        tok = AutoTokenizer.from_pretrained(a.tokenizer)
    else:
        _, _, model = rtn_dense(a.rtn_dense)
        tok = AutoTokenizer.from_pretrained(a.rtn_dense)
    model.eval()
    ps = prompts(tok, a.n)
    t0 = time.time()
    ref = torch_greedy(model, tok, ps, a.new, a.batch)
    torch_s = time.time() - t0
    del model
    got, timing = server_greedy(a.gguf, ps, a.new, a.threads, a.port)
    pairs = [(r, g) for r, g in zip(ref, got) if g is not None]  # timed-out requests are reported, not scored
    agree = sum(x == y for r, g in pairs for x, y in zip(r, g))
    first = [next((i for i, (x, y) in enumerate(zip(r, g)) if x != y), a.new) for r, g in pairs]
    res = {"prompts": len(pairs), "new_tokens": a.new, "token_agreement": agree / (len(pairs) * a.new),
           "identical_sequences": sum(f == a.new for f in first), "first_divergence_median": sorted(first)[len(first) // 2],
           "first_divergence_mean": sum(first) / len(first), "torch_cpu_s": round(torch_s), **timing,
           "gguf_bytes": Path(a.gguf).stat().st_size}
    Path(a.out).write_text(json.dumps({**res, "torch": ref, "llama": got}))
    print(json.dumps(res))


def forced(a) -> None:
    """For every position of PyTorch's greedy continuation, does llama.cpp pick the same next token given the same
    context? No cascade after a first difference; llama.cpp's prompt cache makes each step one new token."""
    from transformers import AutoTokenizer
    ref = json.loads(Path(a.compare).read_text())
    ps = prompts(AutoTokenizer.from_pretrained(a.tokenizer), ref["prompts"])
    proc = start_server(a.gguf, a.threads, a.port)
    hits = n = 0
    per_pos = [0] * ref["new_tokens"]
    try:
        for k, (p, cont) in enumerate(zip(ps, ref["torch"])):
            for i, want in enumerate(cont):
                body = {"prompt": p + cont[:i], "n_predict": 1, "temperature": 0.0, "top_k": 1, "ignore_eos": True,
                        "cache_prompt": True, "return_tokens": True, "samplers": ["top_k"]}
                req = urllib.request.Request(f"http://127.0.0.1:{a.port}/completion", json.dumps(body).encode(),
                                             {"Content-Type": "application/json"})
                got = json.loads(urllib.request.urlopen(req, timeout=120).read())["tokens"][0]
                hits += got == want
                per_pos[i] += got == want
                n += 1
            if (k + 1) % 25 == 0:
                print(json.dumps({"forced": k + 1, "of": len(ps), "agreement": round(hits / n, 4)}), flush=True)
    finally:
        stop_server(proc)
    res = {"gguf": a.gguf, "positions": n, "forced_top1_agreement": hits / n,
           "agreement_first_16": sum(per_pos[:16]) / (16 * len(ps)), "agreement_last_16": sum(per_pos[-16:]) / (16 * len(ps))}
    Path(a.out).write_text(json.dumps(res, indent=1))
    print(json.dumps(res))


def main() -> None:
    p = argparse.ArgumentParser()
    sub = p.add_subparsers(dest="op", required=True)
    b = sub.add_parser("build")
    b.add_argument("--template", required=True)
    b.add_argument("--out", required=True)
    src = b.add_mutually_exclusive_group(required=True)
    src.add_argument("--packed")
    src.add_argument("--rtn-dense")
    b.add_argument("--ternary-as", choices=["q2_0", "bf16"], default="q2_0",
                   help="q2_0: each packed tensor in its native type (ternary/int2 Q2_0, int4 Q4_0, int3 Q3_K + Q4_0 "
                        "down; D-51); bf16: the lossless control")
    b.add_argument("--name", help="a release name for general.name (and base-model metadata)")
    b.add_argument("--dense-as", choices=["bf16", "q8_0"], default="bf16",
                   help="attention and embedding: bf16 as trained, or q8_0 (post hoc, as the Q8 teacher; faster on CPU)")
    tf = sub.add_parser("forced", help="teacher-forced next-token agreement against a compare run's PyTorch tokens")
    tf.add_argument("--gguf", required=True)
    tf.add_argument("--compare", required=True, help="the compare run's JSON (its PyTorch greedy tokens)")
    tf.add_argument("--tokenizer", required=True)
    tf.add_argument("--out", required=True)
    tf.add_argument("--threads", type=int, default=16)
    tf.add_argument("--port", type=int, default=18190)
    c = sub.add_parser("compare")
    c.add_argument("--gguf", required=True)
    c.add_argument("--out", required=True)
    src = c.add_mutually_exclusive_group(required=True)
    src.add_argument("--packed")
    src.add_argument("--rtn-dense")
    c.add_argument("--tokenizer", default="/path/to/data/models/Marco-Nano-Instruct")
    c.add_argument("--n", type=int, default=200)
    c.add_argument("--new", type=int, default=64)
    c.add_argument("--batch", type=int, default=25)
    c.add_argument("--threads", type=int, default=16)
    c.add_argument("--port", type=int, default=18190)
    a = p.parse_args()
    {"build": build, "compare": compare, "forced": forced}[a.op](a)


if __name__ == "__main__":
    main()