"""
Gradio demo for the SLM-264M recipe generator fine-tune (ZeroGPU).
"""

import os
import spaces
import torch
import gradio as gr
from tokenizers import Tokenizer
from huggingface_hub import hf_hub_download

MODEL_REPO = os.environ.get("MODEL_REPO", "uday210/slm-recipe-264m")   # <-- EDIT to your recipe repo
CKPT_FILENAME = os.environ.get("CKPT_FILENAME", "slm-recipe-final.pt")  # <-- EDIT to real filename
TOKENIZER_FILENAME = os.environ.get("TOKENIZER_FILENAME", "tokenizer.json")

PROMPT_TEMPLATE = "Make a recipe using: {ingredients}\n"

BLOCK_SIZE = 2048
EOT_TOKEN = "<|endoftext|>"

from model import GPT  # noqa: E402  (your import-safe model.py)

_token = os.environ.get("HF_TOKEN")

ckpt_path = hf_hub_download(MODEL_REPO, CKPT_FILENAME, token=_token)
tok_path = os.path.join(os.path.dirname(os.path.abspath(__file__)), TOKENIZER_FILENAME)

ckpt = torch.load(ckpt_path, map_location="cpu", weights_only=False)
model = GPT(ckpt["config"])
model.load_state_dict(ckpt["model"])
model.eval()

tok = Tokenizer.from_file(tok_path)
EOT_ID = tok.token_to_id(EOT_TOKEN)

_moved = False


@spaces.GPU(duration=120)
@torch.no_grad()
def generate(prompt: str, max_new_tokens: int = 400, temperature: float = 0.8, top_k: int = 50):
    global _moved
    if not _moved:
        model.to("cuda")
        _moved = True

    ids = tok.encode(prompt).ids
    x = torch.tensor(ids, dtype=torch.long, device="cuda")[None]
    start = len(ids)

    for _ in range(max_new_tokens):
        logits, _ = model(x[:, -BLOCK_SIZE:])
        logits = logits[:, -1, :]
        if temperature and temperature > 0:
            logits = logits / temperature
            if top_k:
                v, _ = torch.topk(logits, min(top_k, logits.size(-1)))
                logits[logits < v[:, [-1]]] = -float("inf")
            probs = torch.softmax(logits, dim=-1)
            nxt = torch.multinomial(probs, num_samples=1)
        else:
            nxt = logits.argmax(dim=-1, keepdim=True)
        x = torch.cat([x, nxt], dim=1)
        if EOT_ID is not None and nxt.item() == EOT_ID:
            break
        yield tok.decode(x[0, start:].tolist())

    yield tok.decode(x[0, start:].tolist())


def make_recipe(ingredients, max_new_tokens, temperature):
    ingredients = (ingredients or "").strip()
    if not ingredients:
        yield "Enter a few ingredients first (e.g. chicken, rice, lemon)."
        return
    prompt = PROMPT_TEMPLATE.format(ingredients=ingredients)
    for partial in generate(prompt, int(max_new_tokens), float(temperature)):
        yield partial


with gr.Blocks(title="SLM-264M Recipe Generator", theme=gr.themes.Soft()) as demo:
    gr.Markdown(
        "# 🍳 SLM-264M Recipe Generator\n"
        "A 264M small language model, trained from scratch and fine-tuned to turn a list "
        "of ingredients into a recipe.\n\n"
        "*Note: format is reliable; it can occasionally repeat an ingredient or drift on "
        "long recipes — it's a 264M model.*"
    )
    with gr.Row():
        with gr.Column():
            ingredients = gr.Textbox(
                label="Ingredients", value="chicken, rice, lemon", lines=3,
                placeholder="e.g. chicken, rice, lemon",
            )
            with gr.Accordion("Advanced", open=False):
                max_tok = gr.Slider(100, 800, value=400, step=50, label="Max new tokens")
                temp = gr.Slider(
                    0.0, 1.2, value=0.8, step=0.1,
                    label="Temperature (higher = more varied; ~0.8 is a good default)",
                )
            go = gr.Button("Generate recipe", variant="primary")
            gr.Examples(
                examples=[
                    ["chicken, rice, lemon"],
                    ["flour, sugar, butter, eggs, chocolate"],
                    ["tomato, basil, garlic, pasta"],
                ],
                inputs=[ingredients],
            )
        with gr.Column():
            out = gr.Textbox(label="Recipe", lines=18, show_copy_button=True)

    go.click(make_recipe, [ingredients, max_tok, temp], out)
    ingredients.submit(make_recipe, [ingredients, max_tok, temp], out)

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