# --- 1. استيراد المكتبات ---
from fastapi import FastAPI
from pydantic import BaseModel
import torch
from transformers import AutoModelForCausalLM, AutoTokenizer, pipeline
import os

# --- 2. إعدادات النموذج ---
# هذا هو اسم نموذجك الذي رفعته
MODEL_ID = "fady-50/F-Chat-Model-GPTQ"
# المنصة بتتوقع إن الـ app يشتغل على بورت 7860
APP_PORT = int(os.environ.get("PORT", 7860))


# --- 3. تحميل النموذج والـ Tokenizer (يحدث مرة واحدة عند بدء التشغيل) ---
print("--- بدء تحميل النموذج ---")
try:
    tokenizer = AutoTokenizer.from_pretrained(MODEL_ID, use_fast=True)
    
    # نحتاج لإعدادات خاصة لتحميل نموذج GPTQ
    model = AutoModelForCausalLM.from_pretrained(
        MODEL_ID,
        device_map="auto", # توزيع النموذج على الـ CPU/GPU المتاح
        torch_dtype=torch.float16, # استخدام دقة أقل لسرعة أفضل
        trust_remote_code=True # ضروري لنماذج GPTQ
    )
    print("--- تم تحميل النموذج بنجاح ---")

    # إنشاء Pipeline جاهز للاستدلال
    # هذا يغلف كل التعقيدات الخاصة بالـ prompt
    pipe = pipeline(
        "text-generation",
        model=model,
        tokenizer=tokenizer,
        max_new_tokens=512, # الحد الأقصى لعدد الكلمات الجديدة التي سينتجها
        do_sample=True,
        temperature=0.7,
        top_p=0.95,
        top_k=40,
        repetition_penalty=1.1
    )

except Exception as e:
    print(f"!!! خطأ فادح أثناء تحميل النموذج: {e}")
    # في حالة الفشل، نجعل الـ pipeline لا شيء لنعرف ذلك
    pipe = None

# --- 4. إعداد FastAPI ---
app = FastAPI()

# نموذج الإدخال (ما سيرسله الباك إند)
class PromptRequest(BaseModel):
    prompt: str

# --- 5. تعريف نقطة النهاية (Endpoint) للـ API ---

@app.get("/")
def read_root():
    """ رسالة ترحيبية للتأكد أن الـ API يعمل """
    return {"status": "success", "message": "F-Chat API is running"}

@app.post("/generate")
async def generate_text(request: PromptRequest):
    """ هذه هي نقطة النهاية الرئيسية التي سيستخدمها الباك إند """
    
    if pipe is None:
        return {"status": "error", "message": "Model failed to load. Check Space logs."}

    try:
        # الحصول على الـ prompt من الطلب
        prompt = request.prompt
        
        # تنسيق الـ prompt كما يتوقعه نموذجك (قد تحتاج لتعديل هذا)
        # هذا مجرد مثال شائع لنماذج الدردشة
        formatted_prompt = f"Human: {prompt}\nAssistant:"
        
        print(f"--- استلام طلب لـ: {formatted_prompt} ---")

        # تشغيل الـ Pipeline
        outputs = pipe(formatted_prompt)
        
        # استخراج النتيجة
        generated_text = outputs[0]['generated_text']
        
        # تنظيف الرد (اختياري: إزالة الـ prompt الأصلي من الرد)
        # final_response = generated_text.split("Assistant:")[-1].strip()
        
        print(f"--- تم إنشاء الرد ---")
        return {"status": "success", "response": generated_text}

    except Exception as e:
        print(f"!!! خطأ أثناء الاستدلال: {e}")
        return {"status": "error", "message": str(e)}

# (اختياري) كود لتشغيل الخادم محلياً للاختبار
if __name__ == "__main__":
    import uvicorn
    # لازم نحدد البورت ده عشان Hugging Face يشتغل صح
    uvicorn.run(app, host="0.0.0.0", port=APP_PORT)