"""
LLM 客户端抽象层

提供统一的 LLM API 调用接口：
- 支持 OpenAI 兼容 API
- 自动重试和错误处理
- 流式/非流式响应
- 代理支持
- Token 估算

使用示例:
    from utils.llm_client import LLMClient, LLMConfig

    config = LLMConfig(
        api_base="https://api.openai.com/v1",
        api_key="sk-xxx",
        model="gpt-4",
    )
    client = LLMClient(config)

    response = await client.chat_completion(
        messages=[{"role": "user", "content": "Hello"}],
        tools=[...],
    )
"""

from __future__ import annotations

import asyncio
import json
import time
from dataclasses import dataclass, field
from typing import Any, AsyncGenerator, Dict, List, Optional, Union

import aiohttp
from loguru import logger

# 可选代理支持
try:
    from aiohttp_socks import ProxyConnector
    PROXY_SUPPORT = True
except ImportError:
    PROXY_SUPPORT = False


@dataclass
class LLMConfig:
    """LLM 配置"""
    api_base: str = "https://api.openai.com/v1"
    api_key: str = ""
    model: str = "gpt-4"
    temperature: float = 0.7
    max_tokens: int = 4096
    timeout: int = 120
    max_retries: int = 3
    retry_delay: float = 1.0

    # 代理配置
    proxy_enabled: bool = False
    proxy_type: str = "socks5"
    proxy_host: str = "127.0.0.1"
    proxy_port: int = 7890

    # 额外参数
    extra_params: Dict[str, Any] = field(default_factory=dict)

    @classmethod
    def from_dict(cls, config: Dict[str, Any]) -> "LLMConfig":
        """从配置字典创建"""
        api_config = config.get("api", {})
        proxy_config = config.get("proxy", {})

        return cls(
            api_base=api_config.get("base_url", "https://api.openai.com/v1"),
            api_key=api_config.get("api_key", ""),
            model=api_config.get("model", "gpt-4"),
            temperature=api_config.get("temperature", 0.7),
            max_tokens=api_config.get("max_tokens", 4096),
            timeout=api_config.get("timeout", 120),
            max_retries=api_config.get("max_retries", 3),
            retry_delay=api_config.get("retry_delay", 1.0),
            proxy_enabled=proxy_config.get("enabled", False),
            proxy_type=proxy_config.get("type", "socks5"),
            proxy_host=proxy_config.get("host", "127.0.0.1"),
            proxy_port=proxy_config.get("port", 7890),
        )


@dataclass
class LLMResponse:
    """LLM 响应"""
    content: str = ""
    tool_calls: List[Dict[str, Any]] = field(default_factory=list)
    finish_reason: str = ""
    usage: Dict[str, int] = field(default_factory=dict)
    raw_response: Dict[str, Any] = field(default_factory=dict)
    error: Optional[str] = None

    @property
    def has_tool_calls(self) -> bool:
        return len(self.tool_calls) > 0

    @property
    def success(self) -> bool:
        return self.error is None


class LLMClient:
    """
    LLM 客户端

    提供统一的 API 调用接口，支持：
    - OpenAI 兼容 API
    - 自动重试
    - 代理
    - 流式响应
    """

    def __init__(self, config: LLMConfig):
        self.config = config
        self._session: Optional[aiohttp.ClientSession] = None

    async def _get_session(self) -> aiohttp.ClientSession:
        """获取或创建 HTTP 会话"""
        if self._session is None or self._session.closed:
            connector = None

            # 配置代理
            if self.config.proxy_enabled and PROXY_SUPPORT:
                proxy_url = (
                    f"{self.config.proxy_type}://"
                    f"{self.config.proxy_host}:{self.config.proxy_port}"
                )
                connector = ProxyConnector.from_url(proxy_url)
                logger.debug(f"[LLMClient] 使用代理: {proxy_url}")

            timeout = aiohttp.ClientTimeout(total=self.config.timeout)
            self._session = aiohttp.ClientSession(
                connector=connector,
                timeout=timeout,
            )

        return self._session

    async def close(self):
        """关闭会话"""
        if self._session and not self._session.closed:
            await self._session.close()
            self._session = None

    def _build_headers(self) -> Dict[str, str]:
        """构建请求头"""
        return {
            "Content-Type": "application/json",
            "Authorization": f"Bearer {self.config.api_key}",
        }

    def _build_payload(
        self,
        messages: List[Dict[str, Any]],
        tools: Optional[List[Dict[str, Any]]] = None,
        stream: bool = False,
        **kwargs,
    ) -> Dict[str, Any]:
        """构建请求体"""
        payload = {
            "model": self.config.model,
            "messages": messages,
            "temperature": self.config.temperature,
            "max_tokens": self.config.max_tokens,
            "stream": stream,
        }

        if tools:
            payload["tools"] = tools
            payload["tool_choice"] = kwargs.get("tool_choice", "auto")

        # 合并额外参数
        payload.update(self.config.extra_params)
        payload.update(kwargs)

        return payload

    async def chat_completion(
        self,
        messages: List[Dict[str, Any]],
        tools: Optional[List[Dict[str, Any]]] = None,
        **kwargs,
    ) -> LLMResponse:
        """
        非流式聊天补全

        Args:
            messages: 消息列表
            tools: 工具列表（可选）
            **kwargs: 额外参数

        Returns:
            LLMResponse 对象
        """
        session = await self._get_session()
        url = f"{self.config.api_base}/chat/completions"
        headers = self._build_headers()
        payload = self._build_payload(messages, tools, stream=False, **kwargs)

        last_error = None

        for attempt in range(self.config.max_retries):
            try:
                async with session.post(url, headers=headers, json=payload) as resp:
                    if resp.status == 200:
                        data = await resp.json()
                        return self._parse_response(data)

                    error_text = await resp.text()
                    last_error = f"HTTP {resp.status}: {error_text[:200]}"
                    logger.warning(f"[LLMClient] 请求失败 (尝试 {attempt + 1}): {last_error}")

                    # 不可重试的错误
                    if resp.status in [400, 401, 403]:
                        break

            except asyncio.TimeoutError:
                last_error = f"请求超时 ({self.config.timeout}s)"
                logger.warning(f"[LLMClient] {last_error} (尝试 {attempt + 1})")

            except Exception as e:
                last_error = str(e)
                logger.warning(f"[LLMClient] 请求异常 (尝试 {attempt + 1}): {e}")

            # 重试延迟
            if attempt < self.config.max_retries - 1:
                await asyncio.sleep(self.config.retry_delay * (attempt + 1))

        return LLMResponse(error=last_error)

    async def chat_completion_stream(
        self,
        messages: List[Dict[str, Any]],
        tools: Optional[List[Dict[str, Any]]] = None,
        **kwargs,
    ) -> AsyncGenerator[str, None]:
        """
        流式聊天补全

        Args:
            messages: 消息列表
            tools: 工具列表（可选）
            **kwargs: 额外参数

        Yields:
            文本片段
        """
        session = await self._get_session()
        url = f"{self.config.api_base}/chat/completions"
        headers = self._build_headers()
        payload = self._build_payload(messages, tools, stream=True, **kwargs)

        try:
            async with session.post(url, headers=headers, json=payload) as resp:
                if resp.status != 200:
                    error_text = await resp.text()
                    logger.error(f"[LLMClient] 流式请求失败: HTTP {resp.status}")
                    return

                async for line in resp.content:
                    line = line.decode("utf-8").strip()
                    if not line or not line.startswith("data: "):
                        continue

                    data_str = line[6:]
                    if data_str == "[DONE]":
                        break

                    try:
                        data = json.loads(data_str)
                        delta = data.get("choices", [{}])[0].get("delta", {})
                        content = delta.get("content", "")
                        if content:
                            yield content
                    except json.JSONDecodeError:
                        continue

        except Exception as e:
            logger.error(f"[LLMClient] 流式请求异常: {e}")

    def _parse_response(self, data: Dict[str, Any]) -> LLMResponse:
        """解析 API 响应"""
        try:
            choice = data.get("choices", [{}])[0]
            message = choice.get("message", {})

            content = message.get("content", "") or ""
            tool_calls = message.get("tool_calls", [])
            finish_reason = choice.get("finish_reason", "")
            usage = data.get("usage", {})

            # 标准化 tool_calls
            parsed_tool_calls = []
            for tc in tool_calls:
                parsed_tool_calls.append({
                    "id": tc.get("id", ""),
                    "type": tc.get("type", "function"),
                    "function": {
                        "name": tc.get("function", {}).get("name", ""),
                        "arguments": tc.get("function", {}).get("arguments", "{}"),
                    }
                })

            return LLMResponse(
                content=content,
                tool_calls=parsed_tool_calls,
                finish_reason=finish_reason,
                usage=usage,
                raw_response=data,
            )

        except Exception as e:
            logger.error(f"[LLMClient] 解析响应失败: {e}")
            return LLMResponse(error=f"解析响应失败: {e}")

    # ==================== Token 估算 ====================

    @staticmethod
    def estimate_tokens(text: str) -> int:
        """
        估算文本的 token 数量

        使用简化规则：
        - 英文约 4 字符 = 1 token
        - 中文约 1.5 字符 = 1 token
        """
        if not text:
            return 0

        chinese_chars = sum(1 for c in text if '\u4e00' <= c <= '\u9fff')
        other_chars = len(text) - chinese_chars

        chinese_tokens = chinese_chars / 1.5
        other_tokens = other_chars / 4

        return int(chinese_tokens + other_tokens)

    @staticmethod
    def estimate_message_tokens(message: Dict[str, Any]) -> int:
        """估算单条消息的 token 数量"""
        content = message.get("content", "")

        if isinstance(content, str):
            return LLMClient.estimate_tokens(content) + 4  # role 等开销

        if isinstance(content, list):
            total = 4
            for item in content:
                if isinstance(item, dict):
                    if item.get("type") == "text":
                        total += LLMClient.estimate_tokens(item.get("text", ""))
                    elif item.get("type") == "image_url":
                        total += 85  # 图片固定开销
            return total

        return 4

    @staticmethod
    def estimate_messages_tokens(messages: List[Dict[str, Any]]) -> int:
        """估算消息列表的总 token 数量"""
        return sum(LLMClient.estimate_message_tokens(m) for m in messages)


# ==================== 便捷函数 ====================

_default_client: Optional[LLMClient] = None


def get_llm_client(config: Optional[LLMConfig] = None) -> LLMClient:
    """获取默认 LLM 客户端"""
    global _default_client
    if config:
        _default_client = LLMClient(config)
    if _default_client is None:
        raise ValueError("LLM 客户端未初始化，请先传入配置")
    return _default_client


# ==================== 导出 ====================

__all__ = [
    'LLMConfig',
    'LLMResponse',
    'LLMClient',
    'get_llm_client',
]