import os
import requests
from fastapi import FastAPI, HTTPException
from pydantic import BaseModel
from typing import List, Dict, Optional
import uuid

app = FastAPI(title="Gemini Multi-Model Proxy")

# Хранилище сессий (в памяти)
chat_sessions: Dict[str, List[dict]] = {}

class UserRequest(BaseModel):
    message: str
    session_id: Optional[str] = None

# Берем ключ из Secret Variables на Hugging Face
API_KEY = os.getenv("GEMINI_API_KEY")

# СПИСОК МОДЕЛЕЙ ДЛЯ ПЕРЕБОРА (от лучшей к худшей)
MODELS_TO_TRY = [
    "gemini-2.0-flash-exp",    # New 2.0 (Free tier available)
    "gemini-1.5-flash",        # Current stable fast
    "gemini-1.5-flash-latest", # Alias for latest flash
    "gemini-1.5-flash-001",    # Fixed version
    "gemini-1.5-pro",          # Pro version
    "gemini-pro"               # Legacy 1.0
]

@app.get("/")
def home():
    return {"status": "running", "strategy": "failover_models"}

@app.post("/newchat")
def new_chat():
    new_id = str(uuid.uuid4())[:8]
    chat_sessions[new_id] = []
    return {"session_id": new_id, "message": "Новый чат создан."}

@app.post("/chat")
def chat(req: UserRequest):
    if not API_KEY:
        raise HTTPException(status_code=500, detail="GEMINI_API_KEY не установлен")

    sid = req.session_id
    if not sid or sid not in chat_sessions:
        sid = str(uuid.uuid4())[:8]
        chat_sessions[sid] = []

    # 1. Сохраняем сообщение юзера
    user_msg_obj = {"role": "user", "parts": [{"text": req.message}]}
    chat_sessions[sid].append(user_msg_obj)

    # 2. Пытаемся отправить запрос, перебирая модели
    last_error = ""
    success = False
    ai_text = ""
    used_model = ""

    payload = {"contents": chat_sessions[sid]}
    headers = {"Content-Type": "application/json"}

    for model_name in MODELS_TO_TRY:
        url = f"https://generativelanguage.googleapis.com/v1beta/models/{model_name}:generateContent?key={API_KEY}"
        
        try:
            print(f"Trying model: {model_name}...") # Лог в консоль HF
            response = requests.post(url, json=payload, headers=headers, timeout=30)
            
            # Если 404 (модель не найдена) или 5xx (ошибка сервера) -> идем к следующей
            if response.status_code != 200:
                error_detail = response.text
                print(f"Failed {model_name}: {response.status_code} - {error_detail}")
                last_error = f"{response.status_code} on {model_name}"
                continue # Пробуем следующую модель
            
            # Если 200 OK
            data = response.json()
            try:
                ai_text = data["candidates"][0]["content"]["parts"][0]["text"]
                used_model = model_name
                success = True
                break # Выходим из цикла, успех!
            except (KeyError, IndexError):
                # Бывает, если сработал фильтр безопасности
                ai_text = "[Gemini заблокировал ответ (Safety Filter)]"
                success = True
                break

        except requests.exceptions.RequestException as e:
            print(f"Connection error on {model_name}: {e}")
            last_error = str(e)
            continue

    # 3. Обработка результата
    if not success:
        # Если перепробовали все модели и ничего не вышло
        chat_sessions[sid].pop() # Удаляем вопрос юзера, чтобы не портить историю
        raise HTTPException(status_code=502, detail=f"Все модели недоступны. Last error: {last_error}")

    # 4. Сохраняем ответ (и дописываем, какая модель ответила, для дебага можно убрать)
    # ai_text += f"\n\n_(model: {used_model})_" 
    
    ai_msg_obj = {"role": "model", "parts": [{"text": ai_text}]}
    chat_sessions[sid].append(ai_msg_obj)

    return {
        "response": ai_text,
        "session_id": sid,
        "model_used": used_model
    }