from fastapi import FastAPI, Request
from transformers import pipeline
import torch
import uvicorn

app = FastAPI(title="Qwen 0.5B AI Chat API")

# Check for GPU (even though free Space uses CPU)
device = "cuda" if torch.cuda.is_available() else "cpu"
dtype = torch.bfloat16 if torch.cuda.is_available() else torch.float32

print(f"Loading Qwen 2.5 0.5B model on {device}...")

try:
    pipe = pipeline(
        "text-generation",
        model="Qwen/Qwen2.5-0.5B-Instruct",
        torch_dtype=dtype,
        device_map="auto",
    )
    print("Model loaded successfully!")
except Exception as e:
    print(f"Error loading model: {e}")
    pipe = None

@app.get("/")
def read_root():
    return {
        "status": "Chat API is active",
        "model": "Qwen 0.5B",
        "device": device
    }

@app.post("/analyze")
async def chat_endpoint(request: Request):
    if pipe is None:
        return {"error": "Model not loaded properly. Check logs."}
    
    try:
        data = await request.json()
    except Exception:
        return {"error": "لا يوجد نص في الرسالة."}

    messages_history = data.get("messages", [])
    prompt = data.get("prompt", "")

    if not messages_history and not prompt:
        return {"error": "لا يوجد نص في الرسالة."}

    system_prompt = (
        "أنت مساعد طبي ذكي متخصص في أمراض الجلد، تعمل ضمن تطبيق DermexAI.\n"
        "مهمتك: الإجابة على أسئلة المستخدمين المتعلقة بأمراض الجلد وتحليل الحالات الجلدية.\n"
        "قواعد يجب اتباعها:\n"
        "- أجب دائماً بنفس لغة المستخدم (عربي أو إنجليزي).\n"
        "- كن واضحاً وموجزاً ومحترفاً.\n"
        "- لا تقدم تشخيصاً نهائياً — وجّه المستخدم لزيارة طبيب متخصص في الحالات الخطيرة.\n"
        "- ركز على أمراض الجلد فقط، وإذا كان السؤال خارج نطاقك أخبر المستخدم بلطف."
    )

    messages = [{"role": "system", "content": system_prompt}]

    if messages_history:
        for msg in messages_history[-10:]:
            role = "user" if msg.get("role") == "user" else "assistant"
            messages.append({"role": role, "content": msg.get("text", "")})
    else:
        messages.append({"role": "user", "content": prompt})
    
    try:
        # Run inference using the chat template directly
        output = pipe(
            messages,
            max_new_tokens=400,
            do_sample=True,
            temperature=0.7,
            top_p=0.9
        )
        
        # The output includes the system, user, and assistant messages. We take the last one.
        result = output[0]["generated_text"][-1]["content"]
        
        return {
            "analysis": result,
            "success": True
        }
    except Exception as e:
        return {"error": f"Failed to generate response: {str(e)}"}

if __name__ == "__main__":
    uvicorn.run(app, host="0.0.0.0", port=7860)