"""
Gradio demo for the SLM-264M data transformer fine-tune.

Pattern: code lives here (in the Space), weights live in the model repo and are
pulled at startup with huggingface_hub. You bundle model.py alongside this file.

Files this Space expects:
  - app.py            (this file)
  - model.py          (YOUR architecture file — copy it from your training repo)
  - requirements.txt
  - README.md         (Space card, has the sdk: gradio frontmatter)

Weights + tokenizer are downloaded from MODEL_REPO below. If that repo is private,
add an HF_TOKEN secret to the Space (Settings -> Variables and secrets).
"""

import os
import torch
import gradio as gr
from tokenizers import Tokenizer
from huggingface_hub import hf_hub_download
import spaces
# ---- point these at your fine-tuned transform model repo -------------------
# MODEL_REPO = os.environ.get("MODEL_REPO", "uday210/slm-264m-transform")  # <-- EDIT
# CKPT_FILENAME = os.environ.get("CKPT_FILENAME", "slm-264m-transform.pt")  # <-- EDIT

MODEL_REPO = os.environ.get("MODEL_REPO", "uday210/slm-transform-264m")
CKPT_FILENAME = os.environ.get("CKPT_FILENAME", "slm-transform-final.pt")
TOKENIZER_REPO = os.environ.get("TOKENIZER_REPO", "uday210/slm-264m-base")   # <-- ADD
TOKENIZER_FILENAME = os.environ.get("TOKENIZER_FILENAME", "tokenizer.json")
# ---- prompt template: MUST match how the fine-tune was trained -------------
# Your notes show the inference format as:
#   "Convert this CSV to JSON:\n<data>\nOutput:\n"
# If your training pairs used a different wording, change it here or the model
# will be prompted off-distribution.
PROMPT_TEMPLATE = "Convert this {src} to {tgt}:\n{data}\nOutput:\n"

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

device = "cuda" if torch.cuda.is_available() else "cpu"
if device == "cpu":
    torch.set_num_threads(os.cpu_count() or 4)

# ---------------------------------------------------------------------------
# Load architecture + weights
# ---------------------------------------------------------------------------
from model import GPT  # noqa: E402  (your import-safe model.py)

_token = os.environ.get("HF_TOKEN")  # only needed if the model repo is private

ckpt_path = hf_hub_download(MODEL_REPO, CKPT_FILENAME, token=_token)
# tokenizer is bundled in the Space, next to app.py
tok_path = os.path.join(os.path.dirname(os.path.abspath(__file__)), TOKENIZER_FILENAME)
ckpt = torch.load(ckpt_path, map_location=device, weights_only=False)
model = GPT(ckpt["config"])
model.load_state_dict(ckpt["model"])
model.eval().to(device)

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


# ---------------------------------------------------------------------------
# Generation (greedy by default — this task is deterministic transformation)
# ---------------------------------------------------------------------------
@spaces.GPU
@torch.no_grad()
def generate(prompt: str, max_new_tokens: int = 512, temperature: float = 0.0):
    ids = tok.encode(prompt).ids
    x = torch.tensor(ids, dtype=torch.long, device=device)[None]
    start = len(ids)

    for _ in range(max_new_tokens):
        logits, _ = model(x[:, -BLOCK_SIZE:])
        logits = logits[:, -1, :]
        if temperature and temperature > 0:
            probs = torch.softmax(logits / temperature, 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
        # stream the decoded suffix so far
        yield tok.decode(x[0, start:].tolist())

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


def transform(data, src_fmt, tgt_fmt, max_new_tokens, temperature):
    data = (data or "").strip()
    if not data:
        yield "Paste some input data first."
        return
    if src_fmt == tgt_fmt:
        yield "Source and target format are the same — pick different formats."
        return
    prompt = PROMPT_TEMPLATE.format(src=src_fmt, tgt=tgt_fmt, data=data)
    for partial in generate(prompt, int(max_new_tokens), float(temperature)):
        yield partial


# ---------------------------------------------------------------------------
# UI
# ---------------------------------------------------------------------------
FORMATS = ["CSV", "JSON", "XML"]

EXAMPLE_CSV = "name,city,age\nalice,london,31\nbob,paris,47\ncarol,tokyo,29"

with gr.Blocks(title="SLM-264M Data Transformer", theme=gr.themes.Soft()) as demo:
    gr.Markdown(
        "# 🔄 SLM-264M Data Transformer\n"
        "A 264M-parameter small language model, trained from scratch and fine-tuned to "
        "convert between **CSV, JSON, and XML**. Runs on CPU — expect a few tokens/sec.\n\n"
        "*Note: numbers can occasionally be miscopied; text and structure are reliable.*"
    )
    with gr.Row():
        with gr.Column():
            data_in = gr.Textbox(
                label="Input data", value=EXAMPLE_CSV, lines=8,
                placeholder="Paste CSV, JSON, or XML here…",
            )
            with gr.Row():
                src = gr.Dropdown(FORMATS, value="CSV", label="From")
                tgt = gr.Dropdown(FORMATS, value="JSON", label="To")
            with gr.Accordion("Advanced", open=False):
                max_tok = gr.Slider(64, 1024, value=512, step=64, label="Max new tokens")
                temp = gr.Slider(
                    0.0, 1.0, value=0.0, step=0.1,
                    label="Temperature (0 = greedy, recommended for exact transforms)",
                )
            go = gr.Button("Transform", variant="primary")
        with gr.Column():
            out = gr.Textbox(label="Output", lines=14, show_copy_button=True)

    go.click(transform, [data_in, src, tgt, max_tok, temp], out)
    data_in.submit(transform, [data_in, src, tgt, max_tok, temp], out)

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