import os

import spaces
import gradio as gr
import torch
from peft import AutoPeftModelForCausalLM
from transformers import AutoTokenizer


MODEL_ID = "Interact-ai/interact-ai-12b"
HF_TOKEN = os.getenv("HF_TOKEN")

if not HF_TOKEN:
    raise RuntimeError(
        "HF_TOKEN is required. Add a Hugging Face token with access to "
        "google/gemma-3-12b-it as a Space secret named HF_TOKEN."
    )

tokenizer = AutoTokenizer.from_pretrained(MODEL_ID, token=HF_TOKEN)
model = AutoPeftModelForCausalLM.from_pretrained(
    MODEL_ID,
    device_map={"": "cpu"},
    torch_dtype=torch.bfloat16,
    token=HF_TOKEN,
)
model = model.to("cuda")
model.eval()


@spaces.GPU(duration=120)
def respond(
    message: str,
    history: list[dict[str, str]],
    temperature: float,
    max_new_tokens: int,
) -> str:
    """Answer a coding question using the Interact AI 12B Coder model."""
    messages = [
        {
            "role": "system",
            "content": (
                "You are Interact AI 12B Coder, a precise and helpful coding "
                "assistant. Explain your reasoning briefly and provide safe, "
                "runnable code when appropriate."
            ),
        }
    ]
    messages.extend(
        {"role": item["role"], "content": item["content"]}
        for item in history
        if item["role"] in {"user", "assistant"}
    )
    messages.append({"role": "user", "content": message})

    inputs = tokenizer.apply_chat_template(
        messages,
        add_generation_prompt=True,
        return_tensors="pt",
        return_dict=True,
    )
    inputs = {key: value.to(model.device) for key, value in inputs.items()}

    with torch.inference_mode():
        output = model.generate(
            **inputs,
            do_sample=temperature > 0,
            temperature=max(temperature, 0.01),
            max_new_tokens=max_new_tokens,
            pad_token_id=tokenizer.eos_token_id,
        )

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


examples = [
    ["Write a Python function that retries an HTTP request with exponential backoff."],
    ["Explain the difference between an abstract class and a protocol in Python."],
    ["Review this SQL query for performance issues: SELECT * FROM users WHERE email LIKE '%@example.com';"],
]

with gr.Blocks(title="Interact AI 12B Coder") as demo:
    gr.Markdown(
        "# Interact AI 12B Coder\n"
        "Ask coding questions and get implementation-focused answers."
    )
    chatbot = gr.Chatbot(
        type="messages",
        height=560,
        show_copy_button=True,
    )
    gr.ChatInterface(
        fn=respond,
        chatbot=chatbot,
        additional_inputs=[
            gr.Slider(
                minimum=0.0,
                maximum=1.0,
                value=0.2,
                step=0.05,
                label="Temperature",
            ),
            gr.Slider(
                minimum=128,
                maximum=2048,
                value=512,
                step=128,
                label="Max new tokens",
            ),
        ],
        examples=examples,
        title=None,
        description="Powered by Interact-ai/interact-ai-12b",
    )

demo.launch(mcp_server=True)