import os

os.environ.setdefault("PYTORCH_CUDA_ALLOC_CONF", "expandable_segments:True")

import json

import spaces
import torch
import gradio as gr
from threading import Thread
from transformers import AutoTokenizer, AutoModelForCausalLM, TextIteratorStreamer

BUILD = "tools-v1"

MODEL_ID = "Qwen/Qwen2.5-Coder-7B-Instruct-GGUF"
GGUF_FILE = "qwen2.5-coder-7b-instruct-q4_k_m.gguf"
SYSTEM_PROMPT = "You are Qwen2.5-Coder, a helpful and concise coding assistant."

tokenizer = AutoTokenizer.from_pretrained(MODEL_ID, gguf_file=GGUF_FILE)
model = (
    AutoModelForCausalLM.from_pretrained(
        MODEL_ID,
        gguf_file=GGUF_FILE,
        dtype=torch.float32,
        attn_implementation="sdpa",
    )
    .to("cuda")
    .eval()
)


def _build_prompt(message, history):
    messages = [{"role": "system", "content": SYSTEM_PROMPT}]
    for turn in history:
        role, content = turn.get("role"), turn.get("content")
        if role in ("user", "assistant") and isinstance(content, str) and content:
            messages.append({"role": role, "content": content})
    messages.append({"role": "user", "content": message})
    return tokenizer.apply_chat_template(
        messages, tokenize=False, add_generation_prompt=True
    )


def _stream(message, history, max_new_tokens, temperature, top_p):
    text = _build_prompt(message, history)
    inputs = tokenizer(text, return_tensors="pt").to("cuda")
    streamer = TextIteratorStreamer(
        tokenizer, skip_prompt=True, skip_special_tokens=True
    )
    kwargs = dict(
        **inputs,
        max_new_tokens=int(max_new_tokens),
        do_sample=temperature > 0,
        temperature=max(float(temperature), 1e-5),
        top_p=float(top_p),
        streamer=streamer,
        pad_token_id=tokenizer.eos_token_id,
    )
    Thread(target=model.generate, kwargs=kwargs).start()
    partial = ""
    for chunk in streamer:
        partial += chunk
        yield partial


@spaces.GPU(duration=90)
def respond(chat_history, temperature, top_p, max_new_tokens):
    message = chat_history[-1]["content"]
    history = chat_history[:-1]
    chat_history = chat_history + [{"role": "assistant", "content": ""}]
    for partial in _stream(message, history, max_new_tokens, temperature, top_p):
        chat_history[-1]["content"] = partial
        yield chat_history


def add_user(message, chat_history):
    if not message or not message.strip():
        return "", chat_history
    return "", chat_history + [{"role": "user", "content": message}]


def _to_qwen_messages(messages):
    out = []
    for m in messages or []:
        role = m.get("role")
        content = m.get("content") or ""
        if role == "assistant" and m.get("tool_calls"):
            parts = [content] if content else []
            for tc in m["tool_calls"]:
                fn = (tc or {}).get("function", {}) or {}
                args = fn.get("arguments")
                if isinstance(args, str):
                    try:
                        args = json.loads(args)
                    except Exception:
                        args = {}
                parts.append(
                    "<tool_call>\n"
                    + json.dumps(
                        {"name": fn.get("name"), "arguments": args},
                        ensure_ascii=False,
                    )
                    + "\n</tool_call>"
                )
            content = "\n".join(parts)
        if role in ("system", "user", "assistant", "tool"):
            out.append({"role": role, "content": content})
    return out


@spaces.GPU(duration=90)
def tool_chat(messages_json, tools_json, temperature, top_p, max_new_tokens):
    messages = json.loads(messages_json) if messages_json else []
    tools = json.loads(tools_json) if tools_json else None
    chat = _to_qwen_messages(messages)
    if not chat or chat[0]["role"] != "system":
        chat = [{"role": "system", "content": SYSTEM_PROMPT}] + chat
    kwargs = {}
    if tools:
        kwargs["tools"] = tools
    try:
        text = tokenizer.apply_chat_template(
            chat, tokenize=False, add_generation_prompt=True, **kwargs
        )
    except TypeError:
        text = tokenizer.apply_chat_template(
            chat, tokenize=False, add_generation_prompt=True
        )
    inputs = tokenizer(text, return_tensors="pt").to("cuda")
    out = model.generate(
        **inputs,
        max_new_tokens=int(max_new_tokens),
        do_sample=float(temperature) > 0,
        temperature=max(float(temperature), 1e-5),
        top_p=float(top_p),
        pad_token_id=tokenizer.eos_token_id,
    )
    return tokenizer.decode(
        out[0][inputs["input_ids"].shape[1]:], skip_special_tokens=True
    )


with gr.Blocks(title="Qwen2.5 Coder 7B (GGUF)") as demo:
    gr.Markdown(
        "# Qwen2.5-Coder-7B-Instruct (GGUF)\n"
        "Streaming chat demo running on **ZeroGPU**. " + BUILD
    )
    chatbot = gr.Chatbot(
        type="messages", height=440, label="Qwen2.5-Coder", allow_tags=False
    )
    with gr.Row():
        msg = gr.Textbox(
            placeholder="Ask me to write or explain code...",
            scale=8,
            container=False,
        )
        send = gr.Button("Send", variant="primary", scale=1)
    with gr.Row():
        temperature = gr.Slider(0.0, 1.5, value=0.7, step=0.05, label="Temperature")
        top_p = gr.Slider(0.1, 1.0, value=0.9, step=0.05, label="Top-p")
        max_tokens = gr.Slider(64, 2048, value=512, step=64, label="Max new tokens")
    clear = gr.Button("Clear")
    gr.Examples(
        [
            ["Write a Python function to merge two sorted lists."],
            ["Debug this SQL: SELECT * FORM users;"],
        ],
        inputs=msg,
    )

    with gr.Row(visible=False):
        tc_msgs = gr.Textbox()
        tc_tools = gr.Textbox()
        tc_temp = gr.Number(value=0.2)
        tc_top_p = gr.Number(value=0.9)
        tc_max = gr.Number(value=512)
        tc_btn = gr.Button("tool")
        tc_out = gr.Textbox()
    tc_btn.click(
        tool_chat,
        [tc_msgs, tc_tools, tc_temp, tc_top_p, tc_max],
        tc_out,
        api_name="tool_chat",
    )

    msg.submit(add_user, [msg, chatbot], [msg, chatbot], queue=False).then(
        respond, [chatbot, temperature, top_p, max_tokens], chatbot
    )
    send.click(add_user, [msg, chatbot], [msg, chatbot], queue=False).then(
        respond, [chatbot, temperature, top_p, max_tokens], chatbot
    )
    clear.click(lambda: None, None, chatbot, queue=False)

demo.queue()
demo.launch(mcp_server=True)