import os
import base64
from fastapi import FastAPI
from fastapi.responses import StreamingResponse
from fastapi.staticfiles import StaticFiles
from pydantic import BaseModel
from google import genai
from google.genai import types

app = FastAPI()

# جلب المفاتيح من متغيرات البيئة
API_KEYS_RAW = os.environ.get("API_KEYS", os.environ.get("GEMINI_API_KEY", ""))
API_KEYS =[k.strip() for k in API_KEYS_RAW.split(",") if k.strip()]
CURRENT_KEY_INDEX = 0

clients =[]
for key in API_KEYS:
    clients.append(genai.Client(api_key=key))

# تحديث هيكل الطلب لدعم الملفات
class FileData(BaseModel):
    mime_type: str
    data: str # Base64
    name: str

class ChatRequest(BaseModel):
    message: str
    history: list
    files: list[FileData] =[]

def get_client():
    global CURRENT_KEY_INDEX
    if not clients:
        raise ValueError("لم يتم العثور على مفاتيح API (API_KEYS) صالحة.")
    return clients[CURRENT_KEY_INDEX]

@app.post("/chat")
async def chat_endpoint(request: ChatRequest):
    global CURRENT_KEY_INDEX
    
    attempts = 0
    while attempts < len(clients):
        try:
            client = get_client()
            
            # بناء تاريخ المحادثة
            contents =[]
            for entry in request.history:
                if entry.get('user'):
                    contents.append(types.Content(role="user", parts=[types.Part.from_text(text=entry['user'])]))
                if entry.get('bot'):
                    contents.append(types.Content(role="model", parts=[types.Part.from_text(text=entry['bot'])]))
            
            # تجهيز رسالة المستخدم الحالية مع الملفات
            user_parts =[]
            
            # إضافة الملفات المرفقة (الصور وغيرها) إن وجدت
            if request.files:
                for f in request.files:
                    try:
                        file_bytes = base64.b64decode(f.data)
                        user_parts.append(types.Part.from_bytes(data=file_bytes, mime_type=f.mime_type))
                    except Exception as e:
                        print(f"Error decoding file {f.name}: {e}")
            
            # إضافة النص
            user_parts.append(types.Part.from_text(text=request.message))
            
            contents.append(types.Content(role="user", parts=user_parts))
            
            # إعدادات الأداة والنظام
            tools =[types.Tool(googleSearch=types.GoogleSearch())]
            system_instruction = """انت Genisi نموذج ذكاء اصطناعي متطور من قبل AnesNT
AnesNT مبادرة جزائرية من ولاية باتنة صاحبها انس كامش باللاتينية Anes Kameche وهو حاليا العضو الوحيد في هذي المبادرة وهي مبادرة تكنولوجية
كن مفيدا وصديقا للمستخدم تحدث اللفة التي تحدث بها المستخدم"""

            config = types.GenerateContentConfig(
                temperature=2,
                thinking_config=types.ThinkingConfig(thinking_level="MINIMAL"),
                media_resolution="MEDIA_RESOLUTION_HIGH",
                tools=tools,
                system_instruction=[types.Part.from_text(text=system_instruction)]
            )

            stream = client.models.generate_content_stream(
                model="gemma-4-31b-it", # يمكن تغييره إلى gemini-1.5-pro-latest لدعم أقوى للصور
                contents=contents,
                config=config
            )

            def stream_generator():
                try:
                    for chunk in stream:
                        if chunk.text:
                            yield chunk.text
                except Exception as stream_err:
                    yield f"\n\n⚠️ **خطأ أثناء توليد الرد:**\n`{str(stream_err)}`"

            return StreamingResponse(stream_generator(), media_type="text/plain")

        except Exception as e:
            error_msg = str(e).lower()
            if "429" in error_msg or "quota" in error_msg:
                CURRENT_KEY_INDEX = (CURRENT_KEY_INDEX + 1) % len(clients)
                attempts += 1
            else:
                def err_gen(): yield f"⚠️ **خطأ في السيرفر:** {str(e)}"
                return StreamingResponse(err_gen(), media_type="text/plain")

    def limit_gen(): yield "⚠️ تم الوصول للحد الأقصى لجميع المفاتيح المتاحة."
    return StreamingResponse(limit_gen(), media_type="text/plain")

app.mount("/", StaticFiles(directory=".", html=True), name="static")