# app.py
import os
import time
import asyncio
from typing import List, Optional
from fastapi import FastAPI, HTTPException, Depends, Header
from pydantic import BaseModel, Field
import openai
from openai import OpenAI, AsyncOpenAI
import logging
from functools import lru_cache

# 配置日志
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

app = FastAPI(title="Simple AI Gateway")

# --- 数据模型 ---
class ChatMessage(BaseModel):
    role: str = Field(..., description="system, user, or assistant")
    content: str

class ChatCompletionRequest(BaseModel):
    model: str = Field("gpt-3.5-turbo", description="逻辑模型名，由网关映射")
    messages: List[ChatMessage]
    temperature: Optional[float] = 0.7
    max_tokens: Optional[int] = 500

class ChatCompletionResponse(BaseModel):
    id: str
    model: str  # 返回实际调用的物理模型
    choices: List[dict]
    usage: dict

# --- 配置与客户端管理 ---
class ModelEndpoint:
    """代表一个物理模型端点"""
    def __init__(self, name: str, base_url: str, api_key: str, priority: int = 1, enabled: bool = True):
        self.name = name  # 物理标识，如 "azure-gpt4-us"
        self.base_url = base_url
        self.api_key = api_key
        self.priority = priority  # 优先级，用于路由
        self.enabled = enabled
        self.client = AsyncOpenAI(base_url=base_url, api_key=api_key)
        self.failure_count = 0  # 连续失败计数
        self.circuit_breaker_open = False  # 熔断器状态
        self.last_failure_time = 0

    async def is_healthy(self) -> bool:
        """简易健康检查：发送一个轻量级请求"""
        if self.circuit_breaker_open:
            # 检查是否应进入半开状态
            if time.time() - self.last_failure_time > 30:  # 30秒冷却
                self.circuit_breaker_open = False
                logger.info(f"Circuit breaker for {self.name} moved to HALF-OPEN")
            else:
                return False
        try:
            # 一个快速的列表模型请求作为健康检查
            await self.client.models.list(timeout=5.0)
            self.failure_count = 0  # 成功则重置失败计数
            return True
        except Exception as e:
            logger.warning(f"Health check failed for {self.name}: {e}")
            self.failure_count += 1
            self.last_failure_time = time.time()
            if self.failure_count >= 3:  # 连续失败3次，触发熔断
                self.circuit_breaker_open = True
                logger.error(f"Circuit breaker OPEN for {self.name}")
            return False

# 模拟配置：从数据库或配置中心加载
MODEL_ROUTING_TABLE = {
    "gpt-4": ["azure-gpt4-us", "openai-gpt4"],  # 逻辑模型映射到物理端点列表（按优先级）
    "gpt-3.5-turbo": ["openai-gpt35-turbo"],
    "claude-3": ["anthropic-claude-3-sonnet"]
}

ENDPOINTS = {
    "azure-gpt4-us": ModelEndpoint(
        name="azure-gpt4-us",
        base_url=os.getenv("AZURE_OPENAI_ENDPOINT", "https://your-resource.openai.azure.com/openai/deployments/gpt-4"),
        api_key=os.getenv("AZURE_OPENAI_KEY", "your-azure-key")
    ),
    "openai-gpt4": ModelEndpoint(
        name="openai-gpt4",
        base_url="https://api.openai.com/v1",
        api_key=os.getenv("OPENAI_API_KEY", "your-openai-key")
    ),
    "openai-gpt35-turbo": ModelEndpoint(
        name="openai-gpt35-turbo",
        base_url="https://api.openai.com/v1",
        api_key=os.getenv("OPENAI_API_KEY", "your-openai-key")
    ),
    # 可以继续添加 Anthropic, Cohere 等端点
}

# --- 依赖注入：简易鉴权 ---
async def verify_api_key(x_api_key: str = Header(...)):
    """验证网关自身的API密钥"""
    valid_keys = {os.getenv("GATEWAY_API_KEY", "gateway-secure-key")}
    if x_api_key not in valid_keys:
        raise HTTPException(status_code=401, detail="Invalid API Key")
    return x_api_key

# --- 核心路由逻辑 ---
async def intelligent_router(request: ChatCompletionRequest) -> ModelEndpoint:
    """
    智能路由策略：
    1. 根据逻辑模型名获取候选端点列表
    2. 过滤掉禁用或熔断的端点
    3. 进行健康检查（可缓存结果）
    4. 按优先级选择第一个健康的端点
    5. 若无健康端点，抛出异常
    """
    logical_model = request.model
    if logical_model not in MODEL_ROUTING_TABLE:
        raise HTTPException(status_code=400, detail=f"Unsupported logical model: {logical_model}")
    
    candidate_names = MODEL_ROUTING_TABLE[logical_model]
    
    for endpoint_name in candidate_names:
        if endpoint_name not in ENDPOINTS:
            continue
        endpoint = ENDPOINTS[endpoint_name]
        
        if not endpoint.enabled:
            continue
            
        # 这里可以加入更复杂的策略：基于成本、实时延迟等
        if await endpoint.is_healthy():
            logger.info(f"Routing request for model '{logical_model}' to endpoint '{endpoint.name}'")
            return endpoint
    
    # 所有候选都不可用，尝试降级策略（示例：gpt-4 降级到 gpt-3.5-turbo）
    if logical_model == "gpt-4":
        logger.warning("All GPT-4 endpoints unhealthy, attempting fallback to gpt-3.5-turbo")
        fallback_request = request.copy(update={"model": "gpt-3.5-turbo"})
        return await intelligent_router(fallback_request)
    
    raise HTTPException(status_code=503, detail=f"No healthy endpoint available for model: {logical_model}")

# --- API 端点 ---
@app.post("/v1/chat/completions", response_model=ChatCompletionResponse)
async def create_chat_completion(
    request: ChatCompletionRequest,
    api_key: str = Depends(verify_api_key)  # 网关鉴权
):
    """统一聊天补全接口"""
    start_time = time.time()
    
    # 1. 智能路由
    target_endpoint = await intelligent_router(request)
    
    # 2. 适配并转发请求
    try:
        # 注意：这里简化了参数映射，实际中需要处理不同供应商的参数差异
        response = await target_endpoint.client.chat.completions.create(
            model=target_endpoint.name,  # 或使用端点特定的模型名
            messages=[msg.dict() for msg in request.messages],
            temperature=request.temperature,
            max_tokens=request.max_tokens,
            timeout=30.0  # 设置超时
        )
        
        # 3. 标准化响应
        completion_response = ChatCompletionResponse(
            id=f"gateway-{response.id}",
            model=request.model,  # 返回逻辑模型名
            choices=[choice.dict() for choice in response.choices],
            usage=response.usage.dict() if response.usage else {}
        )
        
        # 4. 记录指标（可输出到 Prometheus）
        latency = (time.time() - start_time) * 1000  # 毫秒
        logger.info(f"Request completed. Endpoint: {target_endpoint.name}, Latency: {latency:.2f}ms")
        # 此处可添加 Token 计数和成本计算
        
        return completion_response
        
    except Exception as e:
        logger.error(f"Request failed for endpoint {target_endpoint.name}: {e}")
        # 此处可以触发熔断器逻辑（已在健康检查中实现）
        raise HTTPException(status_code=500, detail=f"Model service error: {str(e)}")

if __name__ == "__main__":
    import uvicorn
    uvicorn.run(app, host="0.0.0.0", port=8000)