"""Test del modello distillato con prompt personalizzati da riga di comando."""
from __future__ import annotations

import os
import sys

# Add project root to path
sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
sys.stdout.reconfigure(encoding="utf-8", errors="replace")

import torch
from tokenizers import Tokenizer

from src.model.config import BoccaccioConfig
from src.model.transformer import BoccaccioForCausalLM


def load_model(path: str, config: BoccaccioConfig) -> BoccaccioForCausalLM:
    model = BoccaccioForCausalLM(config).to("cuda")
    sd = torch.load(path, map_location="cuda", weights_only=True)
    if isinstance(sd, dict) and "state_dict" in sd:
        sd = sd["state_dict"]
    cleaned = {}
    for k, v in sd.items():
        nk = k
        # Strip torch.compile prefix
        if nk.startswith("_orig_mod."):
            nk = nk[len("_orig_mod."):]
        # For Lightning .ckpt: strip outer "model." (LightningModule.model -> BoccaccioForCausalLM)
        # but only if double "model.model." (not model.pt which has single "model.")
        elif nk.startswith("model.model."):
            nk = nk[len("model."):]
        cleaned[nk] = v
    missing, unexpected = model.load_state_dict(cleaned, strict=False)
    if missing:
        print(f"WARNING: {len(missing)} missing keys")
    if unexpected:
        print(f"WARNING: {len(unexpected)} unexpected keys")
    model.eval()
    return model


def generate(model, tokenizer, prompt, max_new=300, temp=0.3, rep=1.2):
    ids = tokenizer.encode(prompt).ids
    x = torch.tensor([ids], device="cuda")
    with torch.no_grad():
        for _ in range(max_new):
            out = model(x)
            logits = out["logits"][0, -1, :]
            for tid in set(x[0].tolist()):
                logits[tid] /= rep
            if temp > 0:
                logits = logits / temp
                v, _ = torch.topk(logits, 40)
                logits[logits < v[-1]] = float("-inf")
                probs = torch.softmax(logits, dim=-1)
                tid = torch.multinomial(probs, num_samples=1).item()
            else:
                tid = torch.argmax(logits).item()
            if tid == 2:
                break
            x = torch.cat([x, torch.tensor([[tid]], device="cuda")], dim=1)
    new = x[0].cpu().tolist()[len(ids):]
    return tokenizer.decode(new).strip()


def main():
    config = BoccaccioConfig.from_yaml("configs/model.yaml", variant="model")
    tok = Tokenizer.from_file("tokenizer/boccaccio-32k.json")

    import os
    ckpt_dir = "checkpoints/finetune_distill"
    # Prefer model.pt (pure weights), fall back to .ckpt
    model_pt = os.path.join(ckpt_dir, "model.pt")
    if os.path.exists(model_pt):
        ckpt_path = model_pt
    else:
        ckpts = sorted([f for f in os.listdir(ckpt_dir) if f.endswith(".ckpt")])
        if not ckpts:
            print("Nessun checkpoint trovato in checkpoints/finetune_distill/")
            return
        ckpt_path = os.path.join(ckpt_dir, ckpts[-1])
    print(f"Using: {ckpt_path}\n", flush=True)
    model = load_model(ckpt_path, config)

    print("Modalita' chat. Digita 'quit' per uscire, 'ctx' per aggiungere contesto.\n")

    context = ""
    while True:
        try:
            user_input = input("Tu> ").strip()
        except (EOFError, KeyboardInterrupt):
            break
        if not user_input:
            continue
        if user_input.lower() == "quit":
            break
        if user_input.lower().startswith("ctx "):
            context = user_input[4:].strip()
            print(f"[Contesto impostato: {context[:80]}...]\n")
            continue
        if user_input.lower() == "ctx":
            context = ""
            print("[Contesto cancellato]\n")
            continue

        if context:
            prompt = f"### Contesto: {context}\n\n### Utente: {user_input}\n### Assistente:"
        else:
            prompt = f"### Utente: {user_input}\n### Assistente:"
        resp = generate(model, tok, prompt)
        print(f"\nAssistente: {resp}\n")


if __name__ == "__main__":
    main()