import os
from fastapi import FastAPI, Request, HTTPException
from pydantic import BaseModel  # <-- تمت إضافة الاستيراد
from huggingface_hub import hf_hub_download
from llama_cpp import Llama
import logging

logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

# --- تعريف نموذج البيانات (Pydantic Schema) ---
class QueryRequest(BaseModel):
    prompt: str
    is_math: bool = False

app = FastAPI()

# --- تحميل ALLaM ---
try:
    allam_path = hf_hub_download(
        repo_id="eltay89/ALLaM-7B-Instruct-GGUF",
        filename="ALLaM-7B-Instruct-Q4_K_M.gguf"
    )
    logger.info(f"✅ ALLaM: {allam_path}")
    allam_model = Llama(model_path=allam_path, n_ctx=4096, n_threads=4)
except Exception as e:
    logger.error(f"❌ فشل ALLaM: {e}")
    allam_model = None

# --- تحميل Qwen (اختياري) ---
math_model = None
try:
    qwen_path = hf_hub_download(
        repo_id="Qwen/Qwen2.5-1.5B-Instruct-GGUF",
        filename="qwen2.5-1.5b-instruct-q4_k_m.gguf"
    )
    logger.info(f"✅ Qwen: {qwen_path}")
    math_model = Llama(model_path=qwen_path, n_ctx=4096, n_threads=4)
except Exception as e:
    logger.warning(f"⚠️ Qwen غير متاح، سيتم استخدام ALLaM فقط: {e}")

@app.post("/ask")
async def ask_assistant(query: QueryRequest):  # <-- تم استخدام النموذج هنا
    try:
        prompt = query.prompt
        is_math = query.is_math

        if not prompt:
            raise HTTPException(status_code=400, detail="البرومبت فارغ!")

        model = math_model if (is_math and math_model) else allam_model
        if not model:
            return {"error": "النماذج غير جاهزة بعد. حاول مرة أخرى لاحقاً."}

        formatted_prompt = f"<s>[INST] {prompt} [/INST]" if not is_math else f"<|im_start|>user\n{prompt}<|im_end|>\n<|im_start|>assistant\n"

        output = model(
            formatted_prompt,
            max_tokens=1024,
            stop=["</s>", "<|im_end|>"],
            echo=False
        )
        
        return {"response": output["choices"][0]["text"].strip()}

    except Exception as e:
        logger.error(f"⚠️ خطأ: {e}")
        return {"error": "حدث خطأ. النماذج قد لا تكون جاهزة بعد."}

@app.get("/")
def home():
    return {"status": "online", "message": "الدماغ اليمني ينبض بالحياة!"}