"""
app.py — Gradio web UI for the Mastermind 46M-parameter GPT model.

Fixes applied vs. the broken version:
  1. Load weights from checkpoint["model_state"] (not checkpoint["model"]).
     train.py saves checkpoints as a dict with keys:
         {"model_state": <OrderedDict>, "cfg": <dict>, "iter": int, "val_loss": float}
     So we have to pull the actual weights out of the "model_state" key before
     calling load_state_dict().

  2. Set bias = False to match train.py's default. The trained model has NO
     bias parameters — if you set bias = True here, you'd get a "missing key"
     error for every *.bias parameter.

  3. Robust config: pull values from the checkpoint's saved "cfg" dict when
     possible, so vocab_size / n_layer / etc. always match what was trained.
     Falls back to manual values only if the cfg dict isn't present.
"""

import gradio as gr
import torch

# Only import the GPT class — train.py doesn't expose GPTConfig
from train import GPT


# ----------------------------------------------------------------------------
# 1. Build the config — prefer the values saved inside the checkpoint itself,
#    so we never have a mismatch with what was actually trained.
# ----------------------------------------------------------------------------

ckpt_path = "model.pt"
print(f"[load] reading checkpoint: {ckpt_path}")
checkpoint = torch.load(ckpt_path, map_location=torch.device("cpu"), weights_only=False)

# train.py saves cfg as a plain dict inside the checkpoint
saved_cfg = checkpoint.get("cfg", {})

class SimpleConfig:
    # Pull from saved_cfg when available, otherwise fall back to defaults
    # that match train.py's Config dataclass.
    block_size = saved_cfg.get("block_size", 256)
    vocab_size = saved_cfg.get("vocab_size", 8000)
    n_layer    = saved_cfg.get("n_layer", 12)
    n_head     = saved_cfg.get("n_head", 8)
    n_embd     = saved_cfg.get("n_embd", 512)
    dropout    = 0.0                       # inference mode — dropout disabled
    bias       = saved_cfg.get("bias", False)   # train.py default is False
    ffn_mult   = saved_cfg.get("ffn_mult", 4)

config = SimpleConfig()

print(f"[load] config: vocab={config.vocab_size} layers={config.n_layer} "
      f"embd={config.n_embd} heads={config.n_head} bias={config.bias}")


# ----------------------------------------------------------------------------
# 2. Build the model and load the trained weights
# ----------------------------------------------------------------------------

model = GPT(config)

# --- THE FIX: extract the actual weight dict from checkpoint["model_state"] ---
if "model_state" in checkpoint:
    state_dict = checkpoint["model_state"]
    print(f"[load] found 'model_state' key with {len(state_dict)} tensors")
elif "model" in checkpoint:
    # Backwards-compat: if anyone saved with the older "model" key
    state_dict = checkpoint["model"]
    print(f"[load] found 'model' key with {len(state_dict)} tensors")
else:
    # Last resort: assume the checkpoint IS the state dict directly
    state_dict = checkpoint
    print(f"[load] using checkpoint directly as state_dict ({len(state_dict)} tensors)")

# strict=True ensures every parameter matches; if this still fails, the
# trained model's config differs from what we built above.
model.load_state_dict(state_dict, strict=True)
model.eval()
print("[load] weights loaded successfully — model is ready.")


# ----------------------------------------------------------------------------
# 3. Generation function
# ----------------------------------------------------------------------------

# Default sampling params — feel free to tune
MAX_NEW_TOKENS = 200
TEMPERATURE    = 0.8
TOP_K          = 40

@torch.no_grad()
def generate_text(prompt):
    """Generate a completion from the user's prompt."""
    if not prompt.strip():
        return "Please enter a prompt!"

    try:
        # Tokenize the prompt. We assume you've also uploaded tokenizer.json
        # alongside model.pt. If not, you'll need to load it some other way.
        try:
            from tokenizers import Tokenizer
            tokenizer = Tokenizer.from_file("tokenizer.json")
        except Exception as e:
            return (f"Tokenizer load failed: {e}\n"
                    f"Make sure tokenizer.json is in the same directory as app.py.")

        enc = tokenizer.encode(prompt)
        input_ids = torch.tensor([enc.ids], dtype=torch.long)

        # Autoregressive generation
        for _ in range(MAX_NEW_TOKENS):
            # Crop to block_size if the context has grown too long
            if input_ids.size(1) > config.block_size:
                input_ids = input_ids[:, -config.block_size:]

            logits, _ = model(input_ids)
            next_logits = logits[0, -1] / TEMPERATURE

            # Optional top-k filtering
            if TOP_K is not None and TOP_K > 0:
                v, _ = torch.topk(next_logits, min(TOP_K, next_logits.size(-1)))
                next_logits = next_logits.masked_fill(next_logits < v[-1], float("-inf"))

            probs = torch.softmax(next_logits, dim=-1)
            next_id = torch.multinomial(probs, num_samples=1)
            input_ids = torch.cat([input_ids, next_id.unsqueeze(0)], dim=1)

            # Stop on EOS if your tokenizer has one
            eos_id = tokenizer.token_to_id("<eos>")
            if eos_id is not None and next_id.item() == eos_id:
                break

        # Decode and return
        output_ids = input_ids[0].tolist()
        return tokenizer.decode(output_ids, skip_special_tokens=True)

    except Exception as e:
        import traceback
        return f"Error during generation:\n{traceback.format_exc()}"


# ----------------------------------------------------------------------------
# 4. Gradio UI
# ----------------------------------------------------------------------------

demo = gr.Interface(
    fn=generate_text,
    inputs=gr.Textbox(
        lines=3,
        placeholder="Type your prompt here...",
        label="Prompt",
    ),
    outputs=gr.Textbox(label="Response", lines=10),
    title="Mastermind",
    description="Self-trained 46M parameter GPT-style language model.",
    examples=[
        ["User: Hello! How are you today?"],
        ["User: Can you explain what a neural network is?"],
        ["User: Write a short poem about the ocean."],
    ],
)

if __name__ == "__main__":
    demo.launch()