import torch
import gradio as gr
from transformers import AutoModelForCausalLM, AutoTokenizer

try:
    import spaces  # HF ZeroGPU: required so the Space detects a GPU-capable function
    HAS_SPACES = True
except ImportError:
    # Allows running locally (outside HF Spaces) without the `spaces` package.
    HAS_SPACES = False

    class _NoOpSpaces:
        @staticmethod
        def GPU(func):
            return func

    spaces = _NoOpSpaces()

# --------------------------------------------------------------------------
# Config
# --------------------------------------------------------------------------
MODEL_ID = "SupraLabs/Supra2-100M-Instruct"
ZERO_GPU = HAS_SPACES  # HF ZeroGPU: GPU only exists inside @spaces.GPU-decorated calls
DEVICE = "cpu" if ZERO_GPU else ("cuda" if torch.cuda.is_available() else "cpu")
DTYPE = torch.float32  # 100M params — CPU inference is fast enough, no need for bf16/GPU
MAX_CONTEXT_TOKENS = 1024  # model was trained at 1024; 2048 config but untested beyond 1024

# --------------------------------------------------------------------------
# Load model + tokenizer once at startup
# --------------------------------------------------------------------------
print(f"[*] Loading {MODEL_ID} on {DEVICE}...")

tokenizer = AutoTokenizer.from_pretrained(MODEL_ID, trust_remote_code=True)
model = AutoModelForCausalLM.from_pretrained(
    MODEL_ID,
    dtype=DTYPE,
    trust_remote_code=True,
)
model.to(DEVICE)
model.eval()

print("[*] Model loaded.")

# --------------------------------------------------------------------------
# Generation logic
# --------------------------------------------------------------------------
def build_messages(history, user_message):
    """Convert Gradio chat history (list of dicts) + new message into
    the messages format expected by the model's chat template."""
    messages = []
    for turn in history:
        messages.append({"role": turn["role"], "content": turn["content"]})
    messages.append({"role": "user", "content": user_message})
    return messages


def truncate_messages_to_fit(messages, max_tokens):
    """Drop oldest turns (keeping the latest user message) until the
    tokenized prompt fits within max_tokens. Small model = tiny context,
    so this matters in multi-turn chats."""
    while len(messages) > 1:
        prompt = tokenizer.apply_chat_template(
            messages, tokenize=False, add_generation_prompt=True
        )
        n_tokens = len(tokenizer(prompt)["input_ids"])
        if n_tokens <= max_tokens:
            return messages
        messages.pop(0)  # drop oldest turn
    return messages


@spaces.GPU(duration=30)  # ZeroGPU: allocates a GPU for the duration of this call only
def respond(user_message, history, max_new_tokens, temperature, top_p, top_k):
    if not user_message or not user_message.strip():
        return history, ""

    history = history or []
    messages = build_messages(history, user_message)
    messages = truncate_messages_to_fit(messages, MAX_CONTEXT_TOKENS - max_new_tokens)

    prompt_text = tokenizer.apply_chat_template(
        messages, tokenize=False, add_generation_prompt=True
    )
    inputs = tokenizer(prompt_text, return_tensors="pt").to(DEVICE)

    with torch.no_grad():
        output_ids = model.generate(
            **inputs,
            max_new_tokens=int(max_new_tokens),
            do_sample=True,
            temperature=float(temperature),
            top_p=float(top_p),
            top_k=int(top_k),
            no_repeat_ngram_size=3,
            pad_token_id=tokenizer.pad_token_id,
            eos_token_id=tokenizer.eos_token_id,
        )

    generated_tokens = output_ids[0][inputs["input_ids"].shape[-1]:]
    response = tokenizer.decode(generated_tokens, skip_special_tokens=True).strip()

    history = history + [
        {"role": "user", "content": user_message},
        {"role": "assistant", "content": response},
    ]
    return history, ""


def clear_chat():
    return [], ""


# --------------------------------------------------------------------------
# UI
# --------------------------------------------------------------------------
DESCRIPTION = """
# 🧠 Supra2-100M-Instruct — Chat Demo

A **100M-parameter** decoder-only model trained from scratch by **SupraLabs**
on ~30B tokens of English web text (Qwen3 architecture, custom 32K tokenizer).

[Model card](https://huggingface.co/SupraLabs/Supra2-100M-Instruct) ·
[Base model](https://huggingface.co/SupraLabs/Supra2-100M-Base) ·
[SupraLabs on HF](https://huggingface.co/SupraLabs)
"""

with gr.Blocks(title="Supra2-100M-Instruct Chat", theme=gr.themes.Soft()) as demo:
    gr.Markdown(DESCRIPTION)

    chatbot = gr.Chatbot(
        label="Supra2-100M-Instruct",
        type="messages",
        height=500,
        avatar_images=(None, "https://cdn-avatars.huggingface.co/v1/production/uploads/697f2832c2c5e4daa93cece7/IQMtz5gg-vLFP7Gn75POT.png"),
    )

    with gr.Row():
        msg = gr.Textbox(
            placeholder="What is AI?",
            show_label=False,
            scale=8,
            container=False,
        )
        submit_btn = gr.Button("Enter", variant="primary", scale=1)

    with gr.Row():
        clear_btn = gr.Button("🗑️ Delete chat")

    with gr.Accordion("⚙️ Hyperparams", open=False):
        max_new_tokens = gr.Slider(16, 512, value=200, step=8, label="Max new tokens")
        temperature = gr.Slider(0.1, 1.5, value=0.7, step=0.05, label="Temperature")
        top_p = gr.Slider(0.1, 1.0, value=0.85, step=0.05, label="Top-p")
        top_k = gr.Slider(1, 100, value=25, step=1, label="Top-k")

    gr.Examples(
        examples=[
            "What is AI?",
            "Write a short poem about the sea.",
            "Give me pros and cons for eating fast food.",
            "Who was Albert Einstein?",
        ],
        inputs=msg,
    )

    gen_inputs = [msg, chatbot, max_new_tokens, temperature, top_p, top_k]
    gen_outputs = [chatbot, msg]

    msg.submit(respond, gen_inputs, gen_outputs)
    submit_btn.click(respond, gen_inputs, gen_outputs)
    clear_btn.click(clear_chat, None, [chatbot, msg])

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