"""
QwenMTClient: OpenAI Client 封装，支持多 Key 智能轮换。

业务代码只需将 `OpenAI(...)` 替换为 `QwenMTClient(...)`，其他完全不变。
"""

import logging
from pathlib import Path
from typing import Optional

from openai import AuthenticationError, OpenAI, RateLimitError

from qwen_proxy.key_manager import KeyManager
from qwen_proxy.logging_config import setup_logging

logger = logging.getLogger("qwen_proxy")

_DEFAULT_BASE_URL = "https://dashscope.aliyuncs.com/compatible-mode/v1"


class QwenMTClient:
    """
    多 Key 轮换的 Qwen-MT 客户端。

    用法:
        client = QwenMTClient(keys_file=".keys")
        # 或
        client = QwenMTClient(api_keys=["sk-aaa", "sk-bbb"])
        # 或 (自动从 $CWD/.keys 或环境变量发现)
        client = QwenMTClient()

        # 调用方式与 OpenAI SDK 完全一致
        result = client.chat.completions.create(
            model="qwen-mt-plus",
            messages=[{"role": "user", "content": "你好"}],
            extra_body={"translation_options": {...}}
        )
    """

    def __init__(
        self,
        api_keys: Optional[list[str]] = None,
        keys_file: Optional[str | Path] = None,
        base_url: str = _DEFAULT_BASE_URL,
        cooldown_seconds: float = 60.0,
        max_retries: int = 3,
        **openai_kwargs,
    ):
        # 初始化 KeyManager
        km_kwargs = {"cooldown_seconds": cooldown_seconds, "max_retries": max_retries}

        if api_keys:
            self._key_manager = KeyManager(keys=api_keys, **km_kwargs)
        else:
            self._key_manager = KeyManager.auto_discover(
                keys_file=keys_file, **km_kwargs
            )

        self._base_url = base_url
        self._max_retries = max_retries
        self._openai_kwargs = openai_kwargs

        # 自动初始化日志
        setup_logging()

        # 创建一个初始 client 来暴露 .chat 等属性的类型
        # 实际每次请求会动态切换 key
        first_key = self._key_manager.get_next_key()
        self._current_client = self._make_client(first_key)
        self._current_key = first_key

    def _make_client(self, api_key: str) -> OpenAI:
        """创建一个绑定特定 Key 的 OpenAI Client"""
        return OpenAI(
            api_key=api_key,
            base_url=self._base_url,
            max_retries=0,  # 我们自己管理重试
            **self._openai_kwargs,
        )

    @property
    def key_manager(self) -> KeyManager:
        """暴露 KeyManager 以便查看统计"""
        return self._key_manager

    @property
    def chat(self) -> "_ChatNamespace":
        return _ChatNamespace(self)


class _ChatNamespace:
    """模拟 OpenAI client.chat 命名空间"""

    def __init__(self, client: QwenMTClient):
        self._client = client

    @property
    def completions(self) -> "_CompletionsNamespace":
        return _CompletionsNamespace(self._client)


class _CompletionsNamespace:
    """模拟 OpenAI client.chat.completions 命名空间，核心重试逻辑在此"""

    def __init__(self, client: QwenMTClient):
        self._client = client

    def create(self, **kwargs):
        """
        与 OpenAI SDK 的 chat.completions.create() 完全兼容。
        自动轮换 Key，遇 429 切换重试，遇 401 禁用 Key。
        """
        km = self._client._key_manager
        max_retries = self._client._max_retries
        last_error = None

        for attempt in range(max_retries):
            key = km.get_next_key()
            client = self._client._make_client(key)

            try:
                is_stream = kwargs.get("stream", False)

                if is_stream:
                    # 流式请求：包装 generator 以便在流中捕获错误
                    return self._stream_with_key(client, key, km, **kwargs)
                else:
                    result = client.chat.completions.create(**kwargs)
                    km.report_success(key)
                    return result

            except RateLimitError as e:
                km.report_rate_limited(key)
                last_error = e
                logger.warning(
                    f"[尝试 {attempt + 1}/{max_retries}] "
                    f"Key ...{key[-6:]} 限流，切换下一个 Key"
                )
                continue

            except AuthenticationError as e:
                km.report_invalid(key)
                last_error = e
                logger.error(
                    f"[尝试 {attempt + 1}/{max_retries}] "
                    f"Key ...{key[-6:]} 鉴权失败，已禁用"
                )
                continue

        # 所有重试均失败
        raise RuntimeError(
            f"所有重试均失败 ({max_retries} 次)。"
            f"可用 Key: {km.available_count} 个。"
            f"最后错误: {last_error}"
        ) from last_error

    def _stream_with_key(self, client: OpenAI, key: str, km: KeyManager, **kwargs):
        """
        流式请求的 Key 感知包装。
        流式场景下，错误可能在迭代时才抛出，所以不在此处重试，
        而是直接透传异常，让调用方感知。
        """
        try:
            stream = client.chat.completions.create(**kwargs)
            for chunk in stream:
                yield chunk
            km.report_success(key)
        except RateLimitError:
            km.report_rate_limited(key)
            raise
        except AuthenticationError:
            km.report_invalid(key)
            raise
