import torch
from huggingface_hub import hf_hub_download
from bot_model import GPT, decode, greedy_until_eos, extract_numbers, patch_final_answer, render_add_steps

# Step 1: Download file (returns local path)
ckpt_path = hf_hub_download(
    repo_id="codeninja3d/minigpt-model",
    filename="model_latest.pt",
    repo_type="model"
)

# Step 2: Load checkpoint
checkpoint = torch.load(ckpt_path, map_location="cpu")

# 🔎 Sanity check
if not isinstance(checkpoint, dict):
    raise RuntimeError(
        f"❌ Checkpoint format invalid! Expected dict, got {type(checkpoint)}. "
        f"Value preview: {str(checkpoint)[:200]}"
    )
missing_keys = [k for k in ["model", "vocab", "block_size"] if k not in checkpoint]
if missing_keys:
    raise RuntimeError(
        f"❌ Checkpoint missing keys: {missing_keys}. Found keys: {list(checkpoint.keys())}"
    )

# Step 3: Restore vocab + block size
stoi, itos = checkpoint["vocab"]
block_size = checkpoint["block_size"]

# Step 4: Build model + load weights
model = GPT(len(stoi), max_len=block_size)
model.load_state_dict(checkpoint["model"])
model.eval()

def chatbot_reply(user_text: str):
    nums = extract_numbers(user_text)
    if len(nums) < 2:
        return "Tell me two or more numbers to add."
    if len(nums) >= 3:
        return render_add_steps(nums)

    # Exactly 2 numbers → GPT inference
    expr = " + ".join(str(n) for n in nums)
    prompt = f"q: {expr} a:"
    space_id = stoi.get(" ", next(iter(stoi.values())))  # safe fallback
    idx = torch.tensor([[stoi.get(c, space_id) for c in prompt]], dtype=torch.long)

    out = greedy_until_eos(model, idx, itos)
    resp = decode(out[0].tolist(), itos)

    if "<EOS>" in resp:
        resp = resp.split("<EOS>")[0]
    resp = patch_final_answer(resp, nums)
    return resp.split("a:", 1)[-1].strip()