# llm_manager/api_client.py
"""API 客户端模块 - 处理 OpenAI 兼容 API 调用"""

import httpx
import random
from typing import Optional, Generator
from .config import ConfigManager, APIEndpoint


class APIClient:
    """OpenAI 兼容 API 客户端"""
    
    def __init__(self, config: ConfigManager):
        self.config = config
        self.timeout = httpx.Timeout(60.0, connect=10.0)
    
    def _get_headers(self, token: str) -> dict:
        return {
            "Authorization": f"Bearer {token}",
            "Content-Type": "application/json"
        }
    
    def _normalize_url(self, url: str) -> str:
        """标准化 URL，确保格式正确"""
        url = url.rstrip("/")
        if not url.startswith(("http://", "https://")):
            url = "https://" + url
        return url
    
    def fetch_models(self, endpoint: APIEndpoint) -> list[str]:
        """从 API 获取可用模型列表"""
        url = f"{self._normalize_url(endpoint.url)}/v1/models"
        
        try:
            with httpx.Client(timeout=self.timeout) as client:
                response = client.get(url, headers=self._get_headers(endpoint.token))
                response.raise_for_status()
                data = response.json()
                
                models = []
                for model in data.get("data", []):
                    model_id = model.get("id")
                    if model_id:
                        models.append(model_id)
                
                return sorted(models)
        except httpx.HTTPStatusError as e:
            raise RuntimeError(f"API 请求失败 (HTTP {e.response.status_code}): {e.response.text}")
        except httpx.RequestError as e:
            raise RuntimeError(f"网络请求错误: {e}")
    
    def chat_completion(
        self,
        endpoint: APIEndpoint,
        model: str,
        messages: list[dict],
        stream: bool = True,
        **kwargs
    ) -> Generator[str, None, None] | dict:
        """发送聊天请求"""
        url = f"{self._normalize_url(endpoint.url)}/v1/chat/completions"
        
        payload = {
            "model": model,
            "messages": messages,
            "stream": stream,
            **kwargs
        }
        
        try:
            with httpx.Client(timeout=self.timeout) as client:
                if stream:
                    with client.stream(
                        "POST",
                        url,
                        headers=self._get_headers(endpoint.token),
                        json=payload
                    ) as response:
                        response.raise_for_status()
                        for line in response.iter_lines():
                            if line.startswith("data: "):
                                data = line[6:]
                                if data == "[DONE]":
                                    break
                                try:
                                    import json
                                    chunk = json.loads(data)
                                    content = chunk.get("choices", [{}])[0].get("delta", {}).get("content", "")
                                    if content:
                                        yield content
                                except json.JSONDecodeError:
                                    continue
                else:
                    response = client.post(
                        url,
                        headers=self._get_headers(endpoint.token),
                        json=payload
                    )
                    response.raise_for_status()
                    return response.json()
        except httpx.HTTPStatusError as e:
            raise RuntimeError(f"API 请求失败 (HTTP {e.response.status_code}): {e.response.text}")
        except httpx.RequestError as e:
            raise RuntimeError(f"网络请求错误: {e}")
    
    def get_random_endpoint_and_model(
        self,
        group_name: Optional[str] = None
    ) -> tuple[APIEndpoint, str]:
        """随机获取一个端点和模型"""
        endpoints = self.config.get_all_endpoints()
        
        if not endpoints:
            raise RuntimeError("没有配置任何 API 端点")
        
        if group_name:
            if group_name not in endpoints:
                raise RuntimeError(f"未找到名为 '{group_name}' 的端点")
            endpoint = endpoints[group_name]
            if not endpoint.models:
                raise RuntimeError(f"端点 '{group_name}' 没有可用的模型")
            model = random.choice(endpoint.models)
        else:
            # 随机选择端点和模型
            available = [(ep, m) for ep in endpoints.values() for m in ep.models]
            if not available:
                raise RuntimeError("没有可用的模型")
            endpoint, model = random.choice(available)
        
        return endpoint, model
