import spaces  # لازم يتعمل import قبل torch
import json
import torch
import gradio as gr
from transformers import AutoModelForCausalLM, AutoTokenizer

MODEL_NAME = "Qwen/Qwen2.5-3B-Instruct"
SYSTEM_PROMPT = "You are a helpful assistant. Answer in the same language as the user."

print(f"Loading Model: {MODEL_NAME} ...")
tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
model = AutoModelForCausalLM.from_pretrained(MODEL_NAME, dtype=torch.bfloat16)
model.to("cuda")   # على ZeroGPU مسموح تعمليها على مستوى الملف
model.eval()
print("Model loaded!")


def _text(content):
    if isinstance(content, str):
        return content
    if isinstance(content, list):
        return "".join(p.get("text", "") if isinstance(p, dict) else str(p) for p in content)
    return str(content)


def build_messages(message, history):
    messages = [{"role": "system", "content": SYSTEM_PROMPT}]
    for turn in history or []:
        if isinstance(turn, dict):                       # Gradio 5/6
            messages.append({"role": turn["role"], "content": _text(turn["content"])})
        else:                                            # tuples قديمة
            human, assistant = turn
            messages.append({"role": "user", "content": _text(human)})
            if assistant:
                messages.append({"role": "assistant", "content": _text(assistant)})
    messages.append({"role": "user", "content": _text(message)})
    return messages


@spaces.GPU(duration=60)
def generate_response(message, history):
    messages = build_messages(message, history)
    text = tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
    inputs = tokenizer([text], return_tensors="pt").to("cuda")

    with torch.no_grad():
        outputs = model.generate(
            **inputs,
            max_new_tokens=384,
            temperature=0.7,
            top_p=0.9,
            do_sample=True,
        )
    new_tokens = outputs[0][inputs.input_ids.shape[1]:]
    return tokenizer.decode(new_tokens, skip_special_tokens=True)


def api_generate(message: str, history_json: str = "[]") -> str:
    """Endpoint لأي تطبيق خارجي (Lovable). history_json = JSON list of {role, content}"""
    try:
        history = json.loads(history_json)
    except Exception:
        history = []
    return generate_response(message, history)


with gr.Blocks() as demo:
    gr.ChatInterface(
        fn=generate_response,
        title="Qwen 2.5 1.5B API & Chat",
        description="تحدث مع الموديل هنا أو اربطه بأي تطبيق عبر الـ API.",
    )
    gr.api(api_generate, api_name="generate")

if __name__ == "__main__":
    demo.queue(max_size=10).launch(server_name="0.0.0.0", server_port=7860)