"""
Anthropic API 客户端（标准基类）

支持 Claude 系列模型，也作为 Anthropic SDK 兼容 provider 的基类。
子类只需覆盖 provider 属性和 _clean_env() 即可。
"""

import json
import os
import httpx
from typing import AsyncGenerator, List, Dict, Any, Optional
from anthropic import AsyncAnthropic

from .base import (
    BaseModelClient,
    Message,
    MessageRole,
    ToolCall,
    ToolDefinition,
    StreamChunk,
    CompletionResponse,
    TokenUsage,
)
from core.logger import log_llm_request, log_llm_response, log_error, log_token_usage


class AnthropicClient(BaseModelClient):
    """
    Anthropic API 客户端（标准基类）

    使用官方 anthropic SDK 实现。
    兼容 Anthropic API 的 provider（如 MiniMax）可继承此类，
    只需覆盖 provider 属性和 _clean_env() 方法。
    """

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

        Args:
            model: 模型名称
            api_key: API 密钥
            base_url: API 基础 URL（可选）
            **kwargs: 其他参数（temperature, max_tokens 等）
        """
        super().__init__(model, api_key, base_url, **kwargs)
        self._clean_env()
        self.client = self._create_sdk_client(api_key, base_url)

    def _clean_env(self):
        """清理可能干扰 SDK 的环境变量，子类可覆盖"""
        if "ANTHROPIC_AUTH_TOKEN" in os.environ:
            del os.environ["ANTHROPIC_AUTH_TOKEN"]

    def _create_sdk_client(self, api_key: str, base_url: Optional[str] = None) -> AsyncAnthropic:
        """创建 SDK 客户端实例，子类可覆盖以自定义参数"""
        client_kwargs = {
            "api_key": api_key,
            "timeout": httpx.Timeout(300.0, connect=30.0),
            "max_retries": 5,
        }
        if base_url:
            client_kwargs["base_url"] = base_url
        return AsyncAnthropic(**client_kwargs)

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

    def _convert_messages(self, messages: List[Message]) -> List[Dict[str, Any]]:
        """
        将统一消息格式转换为 Anthropic 格式

        系统提示词通过 system_prompts 参数单独处理，不在此方法中处理。

        Args:
            messages: 统一消息列表（不含系统提示词）

        Returns:
            Anthropic 格式的消息列表
        """
        result = []

        for msg in messages:
            if msg.role == MessageRole.USER:
                # 用户消息
                result.append({
                    "role": "user",
                    "content": msg.content,
                })
            elif msg.role == MessageRole.ASSISTANT:
                if msg.tool_calls:
                    # 带工具调用的助手消息
                    content = []
                    if msg.content:
                        content.append({
                            "type": "text",
                            "text": msg.content,
                        })
                    for tc in msg.tool_calls:
                        content.append({
                            "type": "tool_use",
                            "id": tc.id,
                            "name": tc.name,
                            "input": tc.arguments,
                        })
                    result.append({
                        "role": "assistant",
                        "content": content,
                    })
                else:
                    # 普通助手消息（确保 content 不为空）
                    content = msg.content if msg.content else " "
                    result.append({
                        "role": "assistant",
                        "content": content,
                    })
            elif msg.role == MessageRole.TOOL:
                # 工具结果消息
                if msg.tool_result:
                    result.append({
                        "role": "user",
                        "content": [{
                            "type": "tool_result",
                            "tool_use_id": msg.tool_result.tool_call_id,
                            "content": msg.tool_result.content,
                            "is_error": msg.tool_result.is_error,
                        }],
                    })

        return result

    def _convert_tools(self, tools: List[ToolDefinition]) -> List[Dict[str, Any]]:
        """
        将工具定义转换为 Anthropic 格式

        Args:
            tools: 工具定义列表

        Returns:
            Anthropic 格式的工具列表
        """
        result = []
        for tool in tools:
            result.append({
                "name": tool.name,
                "description": tool.description,
                "input_schema": tool.parameters,
            })
        return result

    def _build_system_param(self, system_prompts: Optional[List[Message]]) -> Optional[str]:
        """
        将系统提示词列表组装为 Anthropic system 参数

        规则：
        - AGENT → <instructions>
        - ENVIRONMENT → <environment>
        - CUSTOM → <custom_instructions>
        - 其他 → 用 role.value 作为标签

        Args:
            system_prompts: 系统提示词 Message 列表

        Returns:
            拼接后的系统提示词字符串，无内容时返回 None
        """
        if not system_prompts:
            return None

        sections = []
        for m in system_prompts:
            # AGENT 用 system_instructions 标签
            if m.role == MessageRole.AGENT:
                tag = "system_instructions"
            elif m.role == MessageRole.ENVIRONMENT:
                tag = "environment"
            elif m.role == MessageRole.CUSTOM:
                tag = "custom_instructions"
            else:
                tag = m.role.value

            sections.append(f"<{tag}>\n{m.content}\n</{tag}>")

        return "\n\n".join(sections) if sections else None

    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: 工具定义列表
            temperature: 温度参数（None 时使用实例默认值）
            max_tokens: 最大 token 数（None 时使用实例默认值）
            **kwargs: 其他参数

        Yields:
            StreamChunk: 流式输出块
        """
        # 三级优先级: 入参 > 实例属性 > 兜底
        effective_temperature = temperature if temperature is not None else self.temperature
        effective_max_tokens = max_tokens if max_tokens is not None else (self.max_tokens or 4096)

        # 转换消息格式
        anthropic_messages = self._convert_messages(messages)

        # 构建请求参数
        request_params = {
            "model": self.model,
            "messages": anthropic_messages,
            "temperature": effective_temperature,
            "max_tokens": effective_max_tokens,
        }

        # 系统提示词：XML 包裹后传给 system 参数
        system_text = self._build_system_param(system_prompts)
        if system_text:
            request_params["system"] = system_text

        if tools:
            request_params["tools"] = self._convert_tools(tools)

        # 合并额外参数
        request_params.update(kwargs)

        # 记录 LLM 请求（完整请求体）
        log_llm_request(
            provider="anthropic",
            model=self.model,
            request_body=request_params,
        )

        # 用于累积工具调用
        current_tool_call: Optional[Dict[str, Any]] = None
        tool_calls: List[ToolCall] = []
        # 用于累积文本内容（用于日志记录）
        accumulated_text: str = ""
        # 用于记录 finish_reason
        finish_reason: Optional[str] = None
        # 是否正在接收文本块
        _in_text_block = False

        try:
            async with self.client.messages.stream(**request_params) as stream:
                async for event in stream:
                    stream_chunk = StreamChunk()

                    if event.type == "content_block_start":
                        block = event.content_block
                        if block.type == "tool_use":
                            # 开始新的工具调用
                            current_tool_call = {
                                "id": block.id,
                                "name": block.name,
                                "arguments": "",
                            }
                        elif block.type == "text":
                            _in_text_block = True

                    elif event.type == "content_block_delta":
                        delta = event.delta
                        if delta.type == "text_delta":
                            # 文本内容
                            stream_chunk.text = delta.text
                            accumulated_text += delta.text  # 累积文本
                        elif delta.type == "input_json_delta":
                            # 工具调用参数（增量）
                            if current_tool_call:
                                current_tool_call["arguments"] += delta.partial_json

                    elif event.type == "content_block_stop":
                        # 内容块结束
                        if current_tool_call:
                            # 解析工具调用参数
                            try:
                                arguments = json.loads(current_tool_call["arguments"]) if current_tool_call["arguments"] else {}
                            except json.JSONDecodeError:
                                arguments = {}

                            tool_calls.append(ToolCall(
                                id=current_tool_call["id"],
                                name=current_tool_call["name"],
                                arguments=arguments,
                            ))
                            current_tool_call = None
                        elif _in_text_block:
                            # 文本块结束
                            _in_text_block = False
                            if accumulated_text:
                                yield StreamChunk(text_complete=True)

                    elif event.type == "message_stop":
                        # 消息结束，但不在这里 yield is_finished
                        # 等待获取 final_message 后再统一 yield，确保包含完整的 usage
                        pass

                    elif event.type == "message_delta":
                        # 消息增量（包含 stop_reason）
                        if hasattr(event, "delta") and hasattr(event.delta, "stop_reason"):
                            finish_reason = event.delta.stop_reason

                    # 只有有文本内容时才 yield（不在这里处理 is_finished）
                    if stream_chunk.text:
                        yield stream_chunk

                # 获取最终 usage
                final_message = await stream.get_final_message()

                # 提取 token 使用量（含 Anthropic 缓存字段）
                usage_obj = None
                if final_message.usage:
                    raw_input = final_message.usage.input_tokens or 0
                    output = final_message.usage.output_tokens or 0
                    cache_creation = getattr(final_message.usage, 'cache_creation_input_tokens', 0) or 0
                    cache_read = getattr(final_message.usage, 'cache_read_input_tokens', 0) or 0
                    usage_obj = TokenUsage(
                        input_tokens=raw_input,
                        output_tokens=output,
                        cache_creation_input_tokens=cache_creation,
                        cache_read_input_tokens=cache_read,
                    )

                # 打印 Token 使用量日志
                if usage_obj:
                    log_token_usage(
                        provider="anthropic",
                        model=self.model,
                        input_tokens=usage_obj.input_tokens,
                        output_tokens=usage_obj.output_tokens,
                        cache_creation_input_tokens=usage_obj.cache_creation_input_tokens,
                        cache_read_input_tokens=usage_obj.cache_read_input_tokens,
                    )

                # 记录 LLM 响应（完整内容）
                log_llm_response(
                    provider="anthropic",
                    model=self.model,
                    response_body={
                        "content": accumulated_text,
                        "tool_calls": [tc.__dict__ for tc in tool_calls] if tool_calls else None,
                        "usage": usage_obj.to_dict() if usage_obj else None,
                    },
                )

                # 统一在这里 yield is_finished=True，确保包含完整的 usage 信息
                yield StreamChunk(
                    is_finished=True,
                    finish_reason=finish_reason or final_message.stop_reason,
                    tool_calls=tool_calls,
                    usage=usage_obj,
                )

        except Exception as e:
            import traceback
            log_error(
                f"anthropic_stream_error: [{type(e).__name__}] {str(e) or repr(e)}\n{traceback.format_exc()}",
                model=self.model,
            )
            raise

    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: 系统提示词列表
            temperature: 温度参数（None 时使用实例默认值）
            max_tokens: 最大 token 数（None 时使用实例默认值）
            **kwargs: 其他参数

        Returns:
            CompletionResponse: 完整响应
        """
        # 三级优先级: 入参 > 实例属性 > 兜底
        effective_temperature = temperature if temperature is not None else self.temperature
        effective_max_tokens = max_tokens if max_tokens is not None else (self.max_tokens or 4096)

        # 转换消息格式
        anthropic_messages = self._convert_messages(messages)

        # 构建请求参数
        request_params = {
            "model": self.model,
            "messages": anthropic_messages,
            "temperature": effective_temperature,
            "max_tokens": effective_max_tokens,
        }

        # 系统提示词：XML 包裹后传给 system 参数
        system_text = self._build_system_param(system_prompts)
        if system_text:
            request_params["system"] = system_text

        if tools:
            request_params["tools"] = self._convert_tools(tools)

        # 合并额外参数
        request_params.update(kwargs)

        log_llm_request(
            provider="anthropic",
            model=self.model,
            request_body=request_params,
        )

        try:
            response = await self.client.messages.create(**request_params)

            result = CompletionResponse(
                finish_reason=response.stop_reason,
            )

            # 解析内容
            text_content = ""
            tool_calls = []

            for block in response.content:
                if block.type == "text":
                    text_content += block.text
                elif block.type == "tool_use":
                    tool_calls.append(ToolCall(
                        id=block.id,
                        name=block.name,
                        arguments=block.input if isinstance(block.input, dict) else {},
                    ))

            result.content = text_content
            result.tool_calls = tool_calls

            # 获取 usage（使用统一的 TokenUsage 格式，含缓存）
            if response.usage:
                raw_input = response.usage.input_tokens or 0
                output = response.usage.output_tokens or 0
                cache_creation = getattr(response.usage, 'cache_creation_input_tokens', 0) or 0
                cache_read = getattr(response.usage, 'cache_read_input_tokens', 0) or 0
                result.usage = TokenUsage(
                    input_tokens=raw_input,
                    output_tokens=output,
                    cache_creation_input_tokens=cache_creation,
                    cache_read_input_tokens=cache_read,
                )

                log_token_usage(
                    provider="anthropic",
                    model=self.model,
                    input_tokens=result.usage.input_tokens,
                    output_tokens=result.usage.output_tokens,
                    cache_creation_input_tokens=result.usage.cache_creation_input_tokens,
                    cache_read_input_tokens=result.usage.cache_read_input_tokens,
                )

            log_llm_response(
                provider="anthropic",
                model=self.model,
                response_body={
                    "content": result.content,
                    "tool_calls": [tc.__dict__ for tc in result.tool_calls] if result.tool_calls else None,
                    "usage": result.usage.to_dict() if result.usage else None,
                },
            )

            return result

        except Exception as e:
            import traceback
            log_error(
                f"anthropic_chat_error: [{type(e).__name__}] {str(e) or repr(e)}\n{traceback.format_exc()}",
                model=self.model,
            )
            raise

    async def count_tokens(
        self,
        messages: List[Message],
        system_prompts: Optional[List[Message]] = None,
        tools: Optional[List[ToolDefinition]] = None,
    ) -> int:
        """
        使用 Anthropic SDK 精确计算 token 数

        调用 client.messages.count_tokens() API。
        失败时 fallback 到基类的字符估算。
        """
        try:
            params: Dict[str, Any] = {
                "model": self.model,
                "messages": self._convert_messages(messages),
            }

            system_text = self._build_system_param(system_prompts)
            if system_text:
                params["system"] = system_text

            if tools:
                params["tools"] = self._convert_tools(tools)

            response = await self.client.messages.count_tokens(**params)
            return response.input_tokens

        except Exception as e:
            log_error(f"anthropic_count_tokens_error: {e}")
            return await super().count_tokens(messages, system_prompts, tools)
