import os
from fastapi import FastAPI, Request, HTTPException
from llama_cpp import Llama
import logging

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

app = FastAPI()

try:
    logger.info("⏳ جاري تحميل النماذج...")
    allam_model = Llama(model_path="./allam.gguf", n_ctx=4096, n_threads=4)
    math_model = Llama(model_path="./qwen_math.gguf", n_ctx=4096, n_threads=4)
    logger.info("✅ النماذج جاهزة!")
except Exception as e:
    logger.error(f"❌ النماذج غير جاهزة بعد: {e}")

@app.post("/ask")
async def ask_assistant(request: Request):
    try:
        data = await request.json()
        prompt = data.get("prompt")
        is_math = data.get("is_math", False)

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

        model = math_model if is_math else allam_model
        
        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": "الدماغ اليمني ينبض بالحياة!"}