from abc import ABC, abstractmethod
from dataclasses import dataclass, field
from typing import Any, Callable, Coroutine, Dict, List, Set, Tuple, Type

import asyncio

from src.common.logger import get_logger
from src.config.config import config_manager
from src.config.model_configs import APIProvider, ModelInfo
from src.llm_models.payload_content.message import Message
from src.llm_models.payload_content.resp_format import RespFormat
from src.llm_models.payload_content.tool_option import ToolCall, ToolOption

logger = get_logger("model_client_registry")


@dataclass
class UsageRecord:
    """
    使用记录类
    """

    model_name: str
    """模型名称"""

    provider_name: str
    """提供商名称"""

    prompt_tokens: int
    """提示token数"""

    completion_tokens: int
    """完成token数"""

    total_tokens: int
    """总token数"""

    prompt_cache_hit_tokens: int = 0
    """输入中缓存命中的 token 数"""

    prompt_cache_miss_tokens: int = 0
    """输入中缓存未命中的 token 数"""


@dataclass
class APIResponse:
    """
    API响应类
    """

    content: str | None = None
    """响应内容"""

    reasoning_content: str | None = None
    """推理内容"""

    tool_calls: List[ToolCall] | None = None
    """工具调用 [(工具名称, 工具参数), ...]"""

    embedding: List[float] | None = None
    """嵌入向量"""

    usage: UsageRecord | None = None
    """使用情况 (prompt_tokens, completion_tokens, total_tokens)"""

    raw_data: Any = None
    """响应原始数据"""


UsageTuple = Tuple[int, ...]
"""统一的使用量元组，顺序为 `(prompt_tokens, completion_tokens, total_tokens, prompt_cache_hit_tokens, prompt_cache_miss_tokens)`。"""

StreamResponseHandler = Callable[
    [Any, asyncio.Event | None],
    Coroutine[Any, Any, Tuple["APIResponse", UsageTuple | None]],
]
"""统一的流式响应处理函数类型。"""

ResponseParser = Callable[[Any], Tuple["APIResponse", UsageTuple | None]]
"""统一的非流式响应解析函数类型。"""


@dataclass(slots=True)
class ResponseRequest:
    """统一的文本/多模态响应请求。"""

    model_info: ModelInfo
    message_list: List[Message]
    tool_options: List[ToolOption] | None = None
    max_tokens: int | None = None
    temperature: float | None = None
    response_format: RespFormat | None = None
    stream_response_handler: StreamResponseHandler | None = None
    async_response_parser: ResponseParser | None = None
    interrupt_flag: asyncio.Event | None = None
    extra_params: Dict[str, Any] = field(default_factory=dict)

    def copy_with(self, **changes: Any) -> "ResponseRequest":
        """基于当前请求创建一个带局部变更的新请求。

        Args:
            **changes: 需要覆盖的字段值。

        Returns:
            ResponseRequest: 复制后的请求对象。
        """
        payload = {
            "model_info": self.model_info,
            "message_list": list(self.message_list),
            "tool_options": None if self.tool_options is None else list(self.tool_options),
            "max_tokens": self.max_tokens,
            "temperature": self.temperature,
            "response_format": self.response_format,
            "stream_response_handler": self.stream_response_handler,
            "async_response_parser": self.async_response_parser,
            "interrupt_flag": self.interrupt_flag,
            "extra_params": dict(self.extra_params),
        }
        payload.update(changes)
        return ResponseRequest(**payload)


@dataclass(slots=True)
class EmbeddingRequest:
    """统一的嵌入请求。"""

    model_info: ModelInfo
    embedding_input: str
    extra_params: Dict[str, Any] = field(default_factory=dict)


@dataclass(slots=True)
class AudioTranscriptionRequest:
    """统一的音频转录请求。"""

    model_info: ModelInfo
    audio_base64: str
    max_tokens: int | None = None
    extra_params: Dict[str, Any] = field(default_factory=dict)


ClientRequest = ResponseRequest | EmbeddingRequest | AudioTranscriptionRequest
"""统一客户端请求类型。"""


class BaseClient(ABC):
    """
    基础客户端
    """

    api_provider: APIProvider

    def __init__(self, api_provider: APIProvider) -> None:
        """初始化基础客户端。

        Args:
            api_provider: API 提供商配置。
        """
        self.api_provider = api_provider

    @abstractmethod
    async def get_response(self, request: ResponseRequest) -> APIResponse:
        """获取对话响应。

        Args:
            request: 统一响应请求对象。

        Returns:
            APIResponse: 统一响应对象。
        """
        raise NotImplementedError("'get_response' method should be overridden in subclasses")

    @abstractmethod
    async def get_embedding(self, request: EmbeddingRequest) -> APIResponse:
        """获取文本嵌入。

        Args:
            request: 统一嵌入请求对象。

        Returns:
            APIResponse: 嵌入响应。
        """
        raise NotImplementedError("'get_embedding' method should be overridden in subclasses")

    @abstractmethod
    async def get_audio_transcriptions(self, request: AudioTranscriptionRequest) -> APIResponse:
        """获取音频转录。

        Args:
            request: 统一音频转录请求对象。

        Returns:
            APIResponse: 音频转录响应。
        """
        raise NotImplementedError("'get_audio_transcriptions' method should be overridden in subclasses")

    @abstractmethod
    def get_support_image_formats(self) -> List[str]:
        """获取支持的图片格式。

        Returns:
            List[str]: 支持的图片格式列表。
        """
        raise NotImplementedError("'get_support_image_formats' method should be overridden in subclasses")


ClientFactory = Callable[[APIProvider], BaseClient]
"""根据 APIProvider 创建客户端实例的工厂函数。"""


@dataclass(slots=True)
class ClientProviderRegistration:
    """LLM Provider 客户端类型注册信息。"""

    client_type: str
    """客户端类型标识，对应模型配置中的 `api_providers[].client_type`。"""

    factory: ClientFactory
    """客户端实例工厂。"""

    owner_plugin_id: str | None = None
    """拥有该客户端类型的插件 ID；主程序内置类型为 ``None``。"""

    version: str = "1.0.0"
    """Provider 实现版本。"""

    description: str = ""
    """Provider 描述文本。"""

    builtin: bool = False
    """是否为主程序内置 Provider。"""


class ClientRegistry:
    """客户端注册表。"""

    def __init__(self) -> None:
        """初始化注册表并绑定配置重载回调。"""
        self.client_registry: Dict[str, ClientProviderRegistration] = {}
        """APIProvider.client_type -> Provider 注册信息映射表。"""
        self.client_instance_cache: Dict[str, BaseClient] = {}
        """APIProvider.name -> BaseClient的映射表"""
        self._owner_client_types: Dict[str, Set[str]] = {}
        """插件 ID -> 该插件拥有的 client_type 集合。"""
        config_manager.register_reload_callback(self.clear_client_instance_cache)

    def register_client_class(self, client_type: str) -> Callable[[Type[BaseClient]], Type[BaseClient]]:
        """注册主程序内置 API 客户端类。

        Args:
            client_type: 客户端类型标识。

        Returns:
            Callable[[Type[BaseClient]], Type[BaseClient]]: 装饰器函数。
        """

        def decorator(cls: Type[BaseClient]) -> Type[BaseClient]:
            """将内置客户端类注册到全局客户端注册表。

            Args:
                cls: 待注册的客户端类。

            Returns:
                Type[BaseClient]: 原始客户端类。
            """
            if not issubclass(cls, BaseClient):
                raise TypeError(f"{cls.__name__} is not a subclass of BaseClient")
            self.register_provider(
                ClientProviderRegistration(
                    client_type=client_type,
                    factory=cls,
                    builtin=True,
                )
            )
            return cls

        return decorator

    @staticmethod
    def _normalize_client_type(client_type: str) -> str:
        """规范化客户端类型标识。

        Args:
            client_type: 原始客户端类型标识。

        Returns:
            str: 去除首尾空白后的客户端类型标识。

        Raises:
            ValueError: 当客户端类型为空时抛出。
        """
        normalized_client_type = str(client_type or "").strip()
        if not normalized_client_type:
            raise ValueError("client_type 不能为空")
        return normalized_client_type

    def register_provider(self, registration: ClientProviderRegistration) -> None:
        """注册单个客户端类型。

        Args:
            registration: Provider 注册信息。

        Raises:
            ValueError: 当客户端类型冲突时抛出。
        """
        client_type = self._normalize_client_type(registration.client_type)
        existing = self.client_registry.get(client_type)
        if existing is not None and existing.owner_plugin_id != registration.owner_plugin_id:
            raise ValueError(
                f"LLM Provider client_type 冲突: {client_type} 已由 {existing.owner_plugin_id or 'host'} 注册"
            )

        self.client_registry[client_type] = ClientProviderRegistration(
            client_type=client_type,
            factory=registration.factory,
            owner_plugin_id=registration.owner_plugin_id,
            version=registration.version,
            description=registration.description,
            builtin=registration.builtin,
        )
        if registration.owner_plugin_id:
            self._owner_client_types.setdefault(registration.owner_plugin_id, set()).add(client_type)
        self.clear_client_instance_cache_by_client_type(client_type)

    def validate_plugin_provider_replacement(self, plugin_id: str, client_types: List[str]) -> None:
        """校验插件 Provider 替换是否会造成运行时冲突。

        Args:
            plugin_id: 目标插件 ID。
            client_types: 插件即将注册的客户端类型列表。

        Raises:
            ValueError: 当客户端类型为空、重复或与其他 owner 冲突时抛出。
        """
        normalized_plugin_id = str(plugin_id or "").strip()
        if not normalized_plugin_id:
            raise ValueError("plugin_id 不能为空")

        normalized_client_types = [self._normalize_client_type(client_type) for client_type in client_types]
        duplicate_client_types = sorted(
            {
                client_type
                for client_type in normalized_client_types
                if normalized_client_types.count(client_type) > 1
            }
        )
        if duplicate_client_types:
            raise ValueError(f"插件 {normalized_plugin_id} 重复声明 LLM Provider: {', '.join(duplicate_client_types)}")

        for client_type in normalized_client_types:
            existing = self.client_registry.get(client_type)
            if existing is None or existing.owner_plugin_id == normalized_plugin_id:
                continue
            raise ValueError(
                f"LLM Provider client_type 冲突: {client_type} 已由 {existing.owner_plugin_id or 'host'} 注册"
            )

    def replace_plugin_providers(
        self,
        plugin_id: str,
        registrations: List[ClientProviderRegistration],
    ) -> None:
        """原子替换一个插件拥有的全部 Provider 注册。

        Args:
            plugin_id: 目标插件 ID。
            registrations: 插件当前上报的 Provider 注册列表。

        Raises:
            ValueError: 当注册信息不合法或存在冲突时抛出。
        """
        normalized_plugin_id = str(plugin_id or "").strip()
        self.validate_plugin_provider_replacement(
            normalized_plugin_id,
            [registration.client_type for registration in registrations],
        )
        self.unregister_plugin_providers(normalized_plugin_id)
        for registration in registrations:
            self.register_provider(
                ClientProviderRegistration(
                    client_type=registration.client_type,
                    factory=registration.factory,
                    owner_plugin_id=normalized_plugin_id,
                    version=registration.version,
                    description=registration.description,
                    builtin=False,
                )
            )

    def unregister_plugin_providers(self, plugin_id: str) -> int:
        """注销一个插件拥有的全部 Provider 注册。

        Args:
            plugin_id: 目标插件 ID。

        Returns:
            int: 被注销的客户端类型数量。
        """
        normalized_plugin_id = str(plugin_id or "").strip()
        if not normalized_plugin_id:
            return 0

        client_types = self._owner_client_types.pop(normalized_plugin_id, set())
        removed_count = 0
        for client_type in client_types:
            registration = self.client_registry.get(client_type)
            if registration is None or registration.owner_plugin_id != normalized_plugin_id:
                continue
            self.client_registry.pop(client_type, None)
            self.clear_client_instance_cache_by_client_type(client_type)
            removed_count += 1
        return removed_count

    def clear_client_instance_cache_by_client_type(self, client_type: str) -> None:
        """清理指定客户端类型对应的客户端实例缓存。

        Args:
            client_type: 需要清理缓存的客户端类型。
        """
        normalized_client_type = str(client_type or "").strip()
        if not normalized_client_type:
            return

        stale_provider_names = [
            provider_name
            for provider_name, client in self.client_instance_cache.items()
            if client.api_provider.client_type == normalized_client_type
        ]
        for provider_name in stale_provider_names:
            self.client_instance_cache.pop(provider_name, None)

    def get_client_class_instance(self, api_provider: APIProvider, force_new: bool = False) -> BaseClient:
        """获取注册的 API 客户端实例。

        Args:
            api_provider: APIProvider 实例。
            force_new: 是否强制创建新实例。

        Returns:
            BaseClient: 注册的 API 客户端实例。
        """
        from . import ensure_client_type_loaded

        ensure_client_type_loaded(api_provider.client_type)

        # 如果强制创建新实例，直接创建不使用缓存
        if force_new:
            if registration := self.client_registry.get(api_provider.client_type):
                return registration.factory(api_provider)
            raise KeyError(f"'{api_provider.client_type}' 类型的 Client 未注册")

        # 正常的缓存逻辑
        if api_provider.name not in self.client_instance_cache:
            if registration := self.client_registry.get(api_provider.client_type):
                self.client_instance_cache[api_provider.name] = registration.factory(api_provider)
            else:
                raise KeyError(f"'{api_provider.client_type}' 类型的 Client 未注册")
        return self.client_instance_cache[api_provider.name]

    def clear_client_instance_cache(self) -> None:
        """清空客户端实例缓存。"""
        self.client_instance_cache.clear()
        logger.info("检测到配置重载，已清空LLM客户端实例缓存")


client_registry = ClientRegistry()
