import os
import time
import base64
import io
import threading
from PIL import Image
from fastapi import FastAPI, HTTPException
from fastapi.middleware.cors import CORSMiddleware
from pydantic import BaseModel
from typing import Optional, List
from google import genai

app = FastAPI()

app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],
    allow_methods=["*"],
    allow_headers=["*"],
)

raw_keys = os.getenv("GEMINI_KEYS", "")
KEY_POOL = [k.strip() for k in raw_keys.split(",") if k.strip()]

current_key_index = 0
index_lock = threading.Lock()

def get_next_key():
    global current_key_index
    with index_lock:
        if not KEY_POOL:
            return None, -1
        idx = current_key_index
        current_key_index = (current_key_index + 1) % len(KEY_POOL)
        return KEY_POOL[idx], idx

class AIRequest(BaseModel):
    prompt: str
    model: str = "gemini-2.5-flash"
    image_b64: Optional[str] = None
    images_b64: List[str] = [] # <--- THÊM BIẾN NÀY ĐỂ NHẬN NHIỀU ẢNH TỪ OCR

@app.get("/")
def home():
    return {"status": "M2W AI Hub is Online (Round-Robin + Multi-Image)", "total_keys_loaded": len(KEY_POOL)}

@app.post("/ask-ai")
def ask_gemini(req: AIRequest):
    if not KEY_POOL:
        raise HTTPException(status_code=500, detail="Server chưa được nạp GEMINI_KEYS!")

    contents = [req.prompt]
    
    # 1. Nhận ảnh đơn (Từ Tạo đề / Vẽ TikZ)
    if req.image_b64:
        try:
            contents.append(Image.open(io.BytesIO(base64.b64decode(req.image_b64))))
        except Exception as e:
            return {"success": False, "error": f"Lỗi giải mã ảnh: {str(e)}"}
            
    # 2. Nhận nhiều ảnh cùng lúc (Từ AI OCR)
    if req.images_b64:
        try:
            for b64_str in req.images_b64:
                contents.append(Image.open(io.BytesIO(base64.b64decode(b64_str))))
        except Exception as e:
            return {"success": False, "error": f"Lỗi giải mã mảng ảnh: {str(e)}"}

    max_retries = len(KEY_POOL) + 2 
    last_error = ""

    for attempt in range(max_retries):
        selected_key, key_idx = get_next_key()
        try:
            client = genai.Client(api_key=selected_key)
            response = client.models.generate_content(model=req.model, contents=contents)
            return {"success": True, "text": response.text, "key_used": f"Key_Index_{key_idx}"}
        except Exception as e:
            last_error = str(e)
            err_str = last_error.lower()
            if "400" in err_str or "invalid" in err_str or "quota" in err_str or "exhausted" in err_str:
                continue 
            elif "429" in err_str or "overload" in err_str:
                time.sleep(1)
                continue
            else:
                continue

    return {"success": False, "error": f"Toàn bộ Key trên Trạm trộn đều lỗi. Lỗi cuối: {last_error}"}