"""
大模型客户端抽象基类

定义统一的接口，支持：
- 流式输出
- 工具调用
- 消息格式转换

支持的模型类型：
- OpenAI 兼容 API (openai, deepseek, glm, qwen 等)
- Anthropic API (claude)
"""

from abc import ABC, abstractmethod
from typing import AsyncGenerator, List, Dict, Any, Optional, Union
from dataclasses import dataclass, field
from enum import Enum


class MessageRole(str, Enum):
    """消息角色"""
    SYSTEM = "system"
    USER = "user"
    ASSISTANT = "assistant"
    TOOL = "tool"
    # 系统提示词子角色
    AGENT = "agent"
    ENVIRONMENT = "environment"
    CUSTOM = "custom"


# 所有系统提示词角色集合
SYSTEM_ROLES = {
    MessageRole.SYSTEM,
    MessageRole.AGENT,
    MessageRole.ENVIRONMENT,
    MessageRole.CUSTOM,
}


@dataclass
class ToolCall:
    """工具调用"""
    id: str
    name: str
    arguments: Dict[str, Any]
    metadata: Dict[str, Any] = field(default_factory=dict)  # Provider 特定元数据（如 Gemini thought_signature）


@dataclass
class ToolResult:
    """工具执行结果"""
    tool_call_id: str
    name: str
    content: str
    is_error: bool = False
    metadata: Optional[Dict[str, Any]] = None  # 工具返回的额外元数据


@dataclass
class Message:
    """统一消息格式"""
    role: MessageRole
    content: str = ""
    tool_calls: List[ToolCall] = field(default_factory=list)
    tool_result: Optional[ToolResult] = None

    def to_dict(self) -> Dict[str, Any]:
        """转换为字典"""
        result = {
            "role": self.role.value,
            "content": self.content,
        }
        if self.tool_calls:
            result["tool_calls"] = [
                {
                    "id": tc.id,
                    "name": tc.name,
                    "arguments": tc.arguments,
                    **({"metadata": tc.metadata} if tc.metadata else {}),
                }
                for tc in self.tool_calls
            ]
        if self.tool_result:
            result["tool_result"] = {
                "tool_call_id": self.tool_result.tool_call_id,
                "name": self.tool_result.name,
                "content": self.tool_result.content,
                "is_error": self.tool_result.is_error,
            }
        return result

    @classmethod
    def from_dict(cls, data: Dict[str, Any]) -> "Message":
        """从字典恢复"""
        role = MessageRole(data["role"])
        content = data.get("content", "")
        tool_calls = [
            ToolCall(
                id=tc["id"], name=tc["name"], arguments=tc["arguments"],
                metadata=tc.get("metadata", {}),
            )
            for tc in data.get("tool_calls", [])
        ]
        tool_result = None
        if data.get("tool_result"):
            tr = data["tool_result"]
            tool_result = ToolResult(
                tool_call_id=tr["tool_call_id"],
                name=tr["name"],
                content=tr["content"],
                is_error=tr.get("is_error", False),
            )
        return cls(role=role, content=content, tool_calls=tool_calls, tool_result=tool_result)


@dataclass
class ToolDefinition:
    """工具定义"""
    name: str
    description: str
    parameters: Dict[str, Any]  # JSON Schema 格式


@dataclass
class TokenUsage:
    """
    统一的 Token 使用量格式

    各 LLM 客户端负责将各自的字段转换为此统一格式：
    - Anthropic: input_tokens(非缓存), cache_creation_input_tokens, cache_read_input_tokens
    - OpenAI: prompt_tokens - cached_tokens -> input_tokens, cached_tokens -> cache_read_input_tokens

    input_tokens 统一表示非缓存输入 token 数。
    context_tokens = input + cache_creation + cache_read，是真实上下文大小，用于压缩判断。
    """
    input_tokens: int = 0   # 非缓存输入 token 数
    output_tokens: int = 0  # 输出 token 数（生成的内容）
    cache_creation_input_tokens: int = 0  # 缓存写入 token 数（Anthropic 首次缓存）
    cache_read_input_tokens: int = 0      # 缓存命中 token 数（Anthropic / OpenAI）

    @property
    def total_tokens(self) -> int:
        """总 token 数（上下文 + 输出）"""
        return self.input_tokens + self.cache_creation_input_tokens + self.cache_read_input_tokens

    def to_dict(self) -> Dict[str, int]:
        """转换为字典"""
        d = {
            "input_tokens": self.input_tokens,
            "output_tokens": self.output_tokens,
            "total_tokens": self.total_tokens,
        }
        if self.cache_creation_input_tokens > 0:
            d["cache_creation_input_tokens"] = self.cache_creation_input_tokens
        if self.cache_read_input_tokens > 0:
            d["cache_read_input_tokens"] = self.cache_read_input_tokens
        return d


@dataclass
class StreamChunk:
    """流式输出块"""
    # 文本内容（增量）
    text: str = ""
    # 文本流结束标记（本轮文本输出完毕）
    text_complete: bool = False
    # 工具调用（可能是增量或完整）
    tool_calls: List[ToolCall] = field(default_factory=list)
    # 是否完成
    is_finished: bool = False
    # 完成原因
    finish_reason: Optional[str] = None
    # 使用统计（统一格式）
    usage: Optional[TokenUsage] = None


@dataclass
class CompletionResponse:
    """完整响应"""
    content: str = ""
    tool_calls: List[ToolCall] = field(default_factory=list)
    finish_reason: Optional[str] = None
    # 使用统计（统一格式）
    usage: Optional[TokenUsage] = None


class BaseModelClient(ABC):
    """
    大模型客户端抽象基类

    所有模型客户端必须实现此接口
    """

    # 移至 provider/models.py - 结构化模型定义

    def __init__(
        self,
        model: str,
        api_key: str,
        base_url: Optional[str] = None,
        **kwargs
    ):
        """
        初始化客户端

        Args:
            model: 模型名称
            api_key: API 密钥
            base_url: API 基础 URL（可选）
            **kwargs: 其他参数（temperature, max_tokens 等会提升为实例属性）
        """
        self.model = model
        self.api_key = api_key
        self.base_url = base_url
        # 从 kwargs 提取标准参数为实例属性（config 配置的值在这里生效）
        self.temperature: float = kwargs.pop("temperature", 0.7)
        self.max_tokens: Optional[int] = kwargs.pop("max_tokens", None)
        # 上下文窗口：始终从 models.py 查表获取
        kwargs.pop("context_window", None)  # 忽略旧配置传入的值
        self.context_window: int = self.get_context_window(model)
        self.extra_params = kwargs

    @staticmethod
    def get_context_window(model: str) -> int:
        """
        根据模型名称返回上下文窗口大小

        匹配策略：先精确匹配，再前缀匹配（取最长前缀）。

        Args:
            model: 模型名称

        Returns:
            上下文窗口大小（token 数）
        """
        from provider.models import get_context_window as get_ctx_window
        return get_ctx_window(model)

    @abstractmethod
    async def chat_stream(
        self,
        messages: List[Message],
        tools: Optional[List[ToolDefinition]] = None,
        system_prompts: Optional[List[Message]] = None,
        temperature: Optional[float] = None,
        max_tokens: Optional[int] = None,
        **kwargs
    ) -> AsyncGenerator[StreamChunk, None]:
        """
        流式对话

        Args:
            messages: 消息列表
            tools: 工具定义列表
            system_prompts: 系统提示词列表（role 为 AGENT/ENVIRONMENT/CUSTOM）
            temperature: 温度参数（None 时使用实例默认值）
            max_tokens: 最大 token 数（None 时使用实例默认值）
            **kwargs: 其他参数

        Yields:
            StreamChunk: 流式输出块
        """
        pass

    @abstractmethod
    async def chat(
        self,
        messages: List[Message],
        tools: Optional[List[ToolDefinition]] = None,
        system_prompts: Optional[List[Message]] = None,
        temperature: Optional[float] = None,
        max_tokens: Optional[int] = None,
        **kwargs
    ) -> CompletionResponse:
        """
        非流式对话

        Args:
            messages: 消息列表
            tools: 工具定义列表
            system_prompts: 系统提示词列表（role 为 AGENT/ENVIRONMENT/CUSTOM）
            temperature: 温度参数（None 时使用实例默认值）
            max_tokens: 最大 token 数（None 时使用实例默认值）
            **kwargs: 其他参数

        Returns:
            CompletionResponse: 完整响应
        """
        pass

    @property
    @abstractmethod
    def provider(self) -> str:
        """返回提供商名称"""
        pass

    async def count_tokens(
        self,
        messages: List[Message],
        system_prompts: Optional[List[Message]] = None,
        tools: Optional[List[ToolDefinition]] = None,
    ) -> int:
        """
        估算 token 数量

        默认实现：字符数 / 4。子类可覆盖以使用 SDK 精确计算。

        Args:
            messages: 对话消息列表
            system_prompts: 系统提示词列表
            tools: 工具定义列表

        Returns:
            估算的 token 数
        """
        total_chars = 0
        if system_prompts:
            total_chars += sum(len(m.content) for m in system_prompts)
        for msg in messages:
            if msg.content:
                total_chars += len(msg.content)
            for tc in (msg.tool_calls or []):
                total_chars += len(str(tc.arguments))
            if msg.tool_result:
                total_chars += len(msg.tool_result.content)
        if tools:
            for tool in tools:
                total_chars += len(tool.name) + len(tool.description)
                total_chars += len(str(tool.parameters))
        return total_chars // 4

    @property
    def model_name(self) -> str:
        """返回模型名称"""
        return self.model
