import logging
import threading

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

logging.basicConfig(level=logging.INFO)

MODEL_ID = "Qwen/Qwen2.5-7B-Instruct"

SYSTEM_PROMPT = (
    "You are Tony-Flash, a helpful AI assistant. "
    "Answer in the user's language. Be accurate and clear."
)

# Load at startup; ZeroGPU attaches the real GPU only inside @spaces.GPU calls.
tok = AutoTokenizer.from_pretrained(MODEL_ID)
model = AutoModelForCausalLM.from_pretrained(MODEL_ID, torch_dtype=torch.bfloat16)
model.to("cuda")


@spaces.GPU(duration=90)
def chat(message, history):
    if not message or not message.strip():
        yield "Please enter a message."
        return

    messages = [{"role": "system", "content": SYSTEM_PROMPT}]
    for item in history or []:
        if isinstance(item, dict) and item.get("role") in ("user", "assistant"):
            content = item.get("content")
            if isinstance(content, str):
                messages.append({"role": item["role"], "content": content})
    messages.append({"role": "user", "content": message.strip()})

    inputs = tok.apply_chat_template(
        messages,
        add_generation_prompt=True,
        return_tensors="pt",
        return_dict=True,
    ).to("cuda")

    streamer = TextIteratorStreamer(tok, skip_prompt=True, skip_special_tokens=True)
    gen_kwargs = dict(
        **inputs,
        streamer=streamer,
        max_new_tokens=1024,
        do_sample=True,
        temperature=0.7,
    )
    threading.Thread(target=model.generate, kwargs=gen_kwargs).start()

    answer = ""
    for piece in streamer:
        answer += piece
        yield answer


demo = gr.ChatInterface(
    fn=chat,
    title="🧠 Tony-Flash",
    description=f"Chat with {MODEL_ID} on ZeroGPU",
    textbox=gr.Textbox(
        placeholder="اكتب رسالتك هنا... / Type your message",
        container=False,
    ),
)

if __name__ == "__main__":
    demo.queue().launch(show_error=True, ssr_mode=False)