"""Load our GGUF container back into TinyLiquid and chat.

The GGUF file is the standard container (Q8_0 / F16); this reader dequantizes
and maps names back to the native architecture, so the shipped GGUF is fully
usable on-device without llama.cpp.

Usage:
  .venv/bin/python -m model.gguf_runtime --gguf hf_repo/tiny-liquid-q8.gguf \
      --prompt "Verify: the bridge was painted in 2019."
"""

import argparse
import sys
from pathlib import Path

import numpy as np
import torch

from gguf import GGUFReader
from gguf.constants import GGMLQuantizationType as Q
from gguf.quants import dequantize

from model.config import TinyLiquidConfig
from model.tiny_liquid import TinyLiquid
from data.tokenizer import load_tokenizer

sys.path.insert(0, str(Path(__file__).resolve().parents[1]))
from hf.export_gguf import back_name


def load_gguf(gguf_path: str, cfg: TinyLiquidConfig, model: torch.nn.Module):
    reader = GGUFReader(gguf_path)
    state = {}
    for t in reader.tensors:
        arr = t.data
        if t.tensor_type == Q.Q8_0:
            arr = dequantize(arr, Q.Q8_0).astype(np.float32)
        elif t.tensor_type == Q.F16:
            arr = arr.astype(np.float32)
        state[back_name(t.name)] = torch.from_numpy(np.ascontiguousarray(arr))
    missing = [k for k in model.state_dict() if k not in state]
    assert not missing, f"gguf missing tensors: {missing[:5]}"
    model.load_state_dict(state)
    return model


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--gguf", default="hf_repo/tiny-liquid-q8.gguf")
    ap.add_argument("--tok", default="data/tokenizer.json")
    ap.add_argument("--prompt", default="<|analyst|><|user|>Hey, how's it going?<|assistant|>")
    ap.add_argument("--max-new", type=int, default=120)
    ap.add_argument("--threads", type=int, default=8)
    args = ap.parse_args()

    torch.set_num_threads(args.threads)
    tok = load_tokenizer(args.tok)
    cfg = TinyLiquidConfig(vocab_size=tok.get_vocab_size())
    model = TinyLiquid(cfg)
    load_gguf(args.gguf, cfg, model)
    model.eval()
    print(f"loaded {args.gguf} into TinyLiquid ({sum(p.numel() for p in model.parameters())} params)")
    ids = tok.encode(args.prompt).ids
    out = model.generate(tok, ids, persona_id=1, max_new=args.max_new,
                         temperature=0.6, top_k=40, repetition_penalty=1.4,
                         no_repeat_ngram_size=4)
    print(tok.decode(out[len(ids):]).strip())


if __name__ == "__main__":
    main()