import os

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


MODEL_ID = "Qwen/Qwen2.5-Coder-3B-Instruct"
MAX_INPUT_TOKENS = 8192
MAX_NEW_TOKENS = 512
DEFAULT_SYSTEM_PROMPT = (
    "You are a helpful coding assistant. "
    "Give accurate, clear answers and use code examples when useful."
)

torch.set_num_threads(min(4, os.cpu_count() or 1))

tokenizer = AutoTokenizer.from_pretrained(MODEL_ID)
model = AutoModelForCausalLM.from_pretrained(
    MODEL_ID,
    torch_dtype=torch.float16,
    attn_implementation="sdpa",
)
model.to("cuda")
model.eval()


def normalize_history(history):
    messages = []

    for item in history or []:
        if isinstance(item, dict):
            role = item.get("role")
            content = item.get("content")
            if role in ("user", "assistant") and isinstance(content, str):
                messages.append({"role": role, "content": content})
        elif isinstance(item, (list, tuple)) and len(item) == 2:
            user_message, assistant_message = item
            if isinstance(user_message, str):
                messages.append({"role": "user", "content": user_message})
            if isinstance(assistant_message, str):
                messages.append({"role": "assistant", "content": assistant_message})

    return messages


def prepare_inputs(system_prompt, messages):
    chat_messages = [{"role": "system", "content": system_prompt}, *messages]
    encoded = tokenizer.apply_chat_template(
        chat_messages,
        tokenize=True,
        add_generation_prompt=True,
        return_tensors="pt",
        return_dict=True,
    )
    return encoded


@spaces.GPU(size="xlarge", duration=60)
def generate(message, history, system_prompt):
    system_prompt = (
        DEFAULT_SYSTEM_PROMPT if system_prompt is None else system_prompt
    )
    messages = normalize_history(history)
    messages.append({"role": "user", "content": message})

    inputs = prepare_inputs(system_prompt, messages)
    while inputs["input_ids"].shape[-1] > MAX_INPUT_TOKENS:
        if len(messages) <= 1:
            raise gr.Error(
                "This message or system prompt is too long. Shorten it and try again."
            )
        messages = messages[2:] if len(messages) > 2 else messages[1:]
        inputs = prepare_inputs(system_prompt, messages)

    inputs = {key: value.to("cuda") for key, value in inputs.items()}
    with torch.inference_mode():
        output = model.generate(
            **inputs,
            max_new_tokens=MAX_NEW_TOKENS,
            do_sample=True,
            temperature=0.2,
            top_p=0.95,
            repetition_penalty=1.05,
            pad_token_id=tokenizer.eos_token_id,
            eos_token_id=tokenizer.eos_token_id,
        )

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


system_prompt_input = gr.Textbox(
    label="System prompt",
    value=DEFAULT_SYSTEM_PROMPT,
    lines=4,
)

demo = gr.ChatInterface(
    fn=generate,
    chatbot=gr.Chatbot(type="messages", height=500),
    additional_inputs=[system_prompt_input],
    title="Qwen2.5 Coder 3B",
    description=(
        "A coding assistant powered by Qwen2.5-Coder-3B on Hugging Face ZeroGPU."
    ),
)

demo.queue()

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