"""Anthropic Provider - 原生 Anthropic Messages API 支持"""

import json
import logging
import time
from typing import Any

import httpx

from agentkit.core.exceptions import LLMProviderError
from agentkit.llm.protocol import (
    LLMProvider,
    LLMRequest,
    LLMResponse,
    StreamChunk,
    TokenUsage,
    ToolCall,
)
from agentkit.llm.retry import (
    CircuitBreaker,
    CircuitBreakerConfig,
    RetryConfig,
    RetryPolicy,
)

logger = logging.getLogger(__name__)

# Anthropic API 常量
_ANTHROPIC_VERSION = "2023-06-01"


class _AnthropicStreamContext:
    """Wraps an httpx streaming response context manager for use with retry/circuit breaker."""

    def __init__(self, response_ctx, response):
        self._response_ctx = response_ctx
        self._response = response

    async def __aenter__(self):
        return self._response

    async def __aexit__(self, exc_type, exc_val, exc_tb):
        return await self._response_ctx.__aexit__(exc_type, exc_val, exc_tb)


class AnthropicProvider(LLMProvider):
    """Anthropic Messages API 原生 Provider"""

    def __init__(
        self,
        api_key: str,
        model: str = "claude-sonnet-4-20250514",
        max_tokens: int = 4096,
        base_url: str = "https://api.anthropic.com",
        timeout: float = 120.0,
        thinking_enabled: bool = False,
        retry_config: RetryConfig | None = None,
        circuit_breaker_config: CircuitBreakerConfig | None = None,
    ):
        self._api_key = api_key
        self._model = model
        self._max_tokens = max_tokens
        self._base_url = base_url.rstrip("/")
        self._timeout = timeout
        self._thinking_enabled = thinking_enabled
        self._client: httpx.AsyncClient | None = None
        self._retry_policy = RetryPolicy(retry_config) if retry_config else None
        self._circuit_breaker = (
            CircuitBreaker(circuit_breaker_config, provider="anthropic")
            if circuit_breaker_config
            else None
        )

    def _get_client(self) -> httpx.AsyncClient:
        """Lazy client initialization"""
        if self._client is None:
            self._client = httpx.AsyncClient(timeout=self._timeout)
        return self._client

    async def close(self) -> None:
        """关闭 HTTP 客户端连接池"""
        if self._client is not None:
            await self._client.aclose()
            self._client = None

    def _build_headers(self) -> dict[str, str]:
        """构建 Anthropic API 请求头"""
        return {
            "x-api-key": self._api_key,
            "anthropic-version": _ANTHROPIC_VERSION,
            "content-type": "application/json",
        }

    def _convert_messages(self, messages: list[dict[str, str]]) -> tuple[str | None, list[dict[str, Any]]]:
        """将 OpenAI 风格消息转换为 Anthropic 格式

        Returns:
            (system_prompt, anthropic_messages)
        """
        system_prompt: str | None = None
        anthropic_messages: list[dict[str, Any]] = []

        for msg in messages:
            role = msg.get("role", "")
            content = msg.get("content", "")

            if role == "system":
                system_prompt = content
                continue

            if role == "assistant":
                # 检查是否有 tool_calls (OpenAI 格式)
                tool_calls = msg.get("tool_calls")
                if tool_calls:
                    blocks: list[dict[str, Any]] = []
                    # 如果有文本内容，先添加文本块
                    if content:
                        blocks.append({"type": "text", "text": content})
                    for tc in tool_calls:
                        func = tc.get("function", {})
                        arguments = func.get("arguments", "{}")
                        if isinstance(arguments, str):
                            try:
                                arguments = json.loads(arguments)
                            except json.JSONDecodeError:
                                arguments = {"raw": arguments}
                        blocks.append({
                            "type": "tool_use",
                            "id": tc.get("id", ""),
                            "name": func.get("name", ""),
                            "input": arguments,
                        })
                    anthropic_messages.append({"role": "assistant", "content": blocks})
                else:
                    anthropic_messages.append({
                        "role": "assistant",
                        "content": [{"type": "text", "text": content}],
                    })
                continue

            if role == "user":
                # 检查是否是 tool_result 消息 (OpenAI 格式中 tool 角色的结果)
                # OpenAI 格式: {"role": "tool", "tool_call_id": "...", "content": "..."}
                if msg.get("tool_call_id"):
                    tool_result_blocks: list[dict[str, Any]] = []
                    tool_content = msg.get("content", "")
                    # tool_result 的 content 可以是字符串或内容块列表
                    if isinstance(tool_content, str):
                        tool_result_blocks.append({"type": "text", "text": tool_content})
                    elif isinstance(tool_content, list):
                        tool_result_blocks = tool_content  # type: ignore[assignment]
                    else:
                        tool_result_blocks.append({"type": "text", "text": str(tool_content)})

                    anthropic_messages.append({
                        "role": "user",
                        "content": [{
                            "type": "tool_result",
                            "tool_use_id": msg.get("tool_call_id", ""),
                            "content": tool_result_blocks,
                        }],
                    })
                else:
                    anthropic_messages.append({
                        "role": "user",
                        "content": [{"type": "text", "text": content}],
                    })
                continue

            if role == "tool":
                # OpenAI 格式中独立的 tool 消息
                tool_content = msg.get("content", "")
                if isinstance(tool_content, str):
                    result_content: list[dict[str, Any]] | str = [{"type": "text", "text": tool_content}]
                elif isinstance(tool_content, list):
                    result_content = tool_content
                else:
                    result_content = [{"type": "text", "text": str(tool_content)}]

                anthropic_messages.append({
                    "role": "user",
                    "content": [{
                        "type": "tool_result",
                        "tool_use_id": msg.get("tool_call_id", ""),
                        "content": result_content,
                    }],
                })

        return system_prompt, anthropic_messages

    def _convert_tools(self, tools: list[dict[str, Any]]) -> list[dict[str, Any]]:
        """将 OpenAI function 格式转换为 Anthropic tool 格式"""
        anthropic_tools = []
        for tool in tools:
            if tool.get("type") == "function":
                func = tool.get("function", {})
                anthropic_tools.append({
                    "name": func.get("name", ""),
                    "description": func.get("description", ""),
                    "input_schema": func.get("parameters", {"type": "object", "properties": {}}),
                })
        return anthropic_tools

    def _convert_tool_choice(self, tool_choice: str) -> dict[str, Any] | None:
        """将 OpenAI tool_choice 格式转换为 Anthropic 格式"""
        if tool_choice == "auto":
            return {"type": "auto"}
        elif tool_choice == "required":
            return {"type": "any"}
        elif tool_choice and tool_choice not in ("none",):
            # 如果指定了具体工具名
            return {"type": "tool", "name": tool_choice}
        return None

    def _parse_response(self, data: dict[str, Any], model: str) -> LLMResponse:
        """将 Anthropic 响应转换为 LLMResponse"""
        content_blocks = data.get("content", [])
        text_parts: list[str] = []
        tool_calls: list[ToolCall] = []

        for block in content_blocks:
            block_type = block.get("type", "")
            if block_type == "text":
                text_parts.append(block.get("text", ""))
            elif block_type == "tool_use":
                tool_calls.append(ToolCall(
                    id=block.get("id", ""),
                    name=block.get("name", ""),
                    arguments=block.get("input", {}),
                ))

        usage_data = data.get("usage", {})
        usage = TokenUsage(
            prompt_tokens=usage_data.get("input_tokens", 0),
            completion_tokens=usage_data.get("output_tokens", 0),
        )

        return LLMResponse(
            content="".join(text_parts),
            model=data.get("model", model),
            usage=usage,
            tool_calls=tool_calls,
        )

    def _handle_error(self, status_code: int, resp_body: bytes) -> None:
        """处理 Anthropic API 错误响应"""
        try:
            error_data = json.loads(resp_body)
            error_info = error_data.get("error", {})
            error_msg = error_info.get("message", f"HTTP {status_code}")
        except (json.JSONDecodeError, AttributeError):
            error_msg = f"HTTP {status_code}"

        raise LLMProviderError("anthropic", f"HTTP {status_code}: {error_msg}")

    async def chat(self, request: LLMRequest) -> LLMResponse:
        """发送 chat 请求（带 retry + circuit breaker）"""
        if self._circuit_breaker and self._retry_policy:
            return await self._circuit_breaker.execute(
                self._retry_policy.execute, self._chat_impl, request
            )
        if self._retry_policy:
            return await self._retry_policy.execute(self._chat_impl, request)
        if self._circuit_breaker:
            return await self._circuit_breaker.execute(self._chat_impl, request)
        return await self._chat_impl(request)

    async def _chat_impl(self, request: LLMRequest) -> LLMResponse:
        client = self._get_client()
        url = f"{self._base_url}/v1/messages"
        headers = self._build_headers()

        system_prompt, anthropic_messages = self._convert_messages(request.messages)

        payload: dict[str, Any] = {
            "model": request.model,
            "max_tokens": request.max_tokens or self._max_tokens,
            "messages": anthropic_messages,
        }

        if system_prompt is not None:
            payload["system"] = system_prompt

        if request.tools:
            payload["tools"] = self._convert_tools(request.tools)
            tool_choice = self._convert_tool_choice(request.tool_choice)
            if tool_choice is not None:
                payload["tool_choice"] = tool_choice

        start = time.monotonic()

        try:
            resp = await client.post(url, json=payload, headers=headers)
        except httpx.HTTPError as e:
            raise LLMProviderError("anthropic", str(e)) from e

        latency_ms = (time.monotonic() - start) * 1000

        if resp.status_code != 200:
            self._handle_error(resp.status_code, resp.content)

        data = resp.json()
        response = self._parse_response(data, request.model)
        response.latency_ms = latency_ms

        return response

    async def chat_stream(self, request: LLMRequest):
        """Stream chat response using SSE（带 retry + circuit breaker）"""
        # For streaming, retry/circuit breaker only protect the connection phase.
        if self._circuit_breaker and self._retry_policy:
            ctx = await self._circuit_breaker.execute(
                self._retry_policy.execute, self._open_stream, request
            )
        elif self._retry_policy:
            ctx = await self._retry_policy.execute(self._open_stream, request)
        elif self._circuit_breaker:
            ctx = await self._circuit_breaker.execute(self._open_stream, request)
        else:
            ctx = await self._open_stream(request)

        async with ctx as response:
            async for chunk in self._iterate_stream(response, request):
                yield chunk

    async def _open_stream(self, request: LLMRequest):
        """Open the streaming HTTP connection; returns an async context manager."""
        client = self._get_client()
        url = f"{self._base_url}/v1/messages"
        headers = self._build_headers()

        system_prompt, anthropic_messages = self._convert_messages(request.messages)

        payload: dict[str, Any] = {
            "model": request.model,
            "max_tokens": request.max_tokens or self._max_tokens,
            "messages": anthropic_messages,
            "stream": True,
        }

        if system_prompt is not None:
            payload["system"] = system_prompt

        if request.tools:
            payload["tools"] = self._convert_tools(request.tools)
            tool_choice = self._convert_tool_choice(request.tool_choice)
            if tool_choice is not None:
                payload["tool_choice"] = tool_choice

        response_ctx = client.stream("POST", url, json=payload, headers=headers)
        response = await response_ctx.__aenter__()

        if response.status_code != 200:
            error_body = await response.aread()
            await response_ctx.__aexit__(None, None, None)
            self._handle_error(response.status_code, error_body)

        return _AnthropicStreamContext(response_ctx, response)

    async def _iterate_stream(self, response, request: LLMRequest):
        """Iterate over an already-open SSE stream and yield StreamChunks."""
        # Accumulated tool calls: tool_use_id -> {id, name, input_json_str}
        accumulated_tool_calls: dict[str, dict[str, Any]] = {}
        current_tool_id: str | None = None
        current_tool_name: str | None = None
        current_tool_input_json: str = ""

        async for line in response.aiter_lines():
            line = line.strip()
            if not line:
                continue

            # Anthropic SSE format: "event: <type>" then "data: <json>"
            if line.startswith("event: "):
                event_type = line[7:]
                continue

            if not line.startswith("data: "):
                continue

            data_str = line[6:]
            try:
                data = json.loads(data_str)
            except json.JSONDecodeError:
                continue

            event_type = data.get("type", "")

            if event_type == "message_start":
                # Message started, no content yet
                continue

            elif event_type == "content_block_start":
                content_block = data.get("content_block", {})
                if content_block.get("type") == "tool_use":
                    current_tool_id = content_block.get("id", "")
                    current_tool_name = content_block.get("name", "")
                    current_tool_input_json = ""

            elif event_type == "content_block_delta":
                delta = data.get("delta", {})
                delta_type = delta.get("type", "")

                if delta_type == "text_delta":
                    text = delta.get("text", "")
                    if text:
                        yield StreamChunk(
                            content=text,
                            model=request.model,
                        )

                elif delta_type == "input_json_delta":
                    partial_json = delta.get("partial_json", "")
                    if partial_json:
                        current_tool_input_json += partial_json

            elif event_type == "content_block_stop":
                # Finalize current tool call if any
                if current_tool_id is not None:
                    try:
                        arguments = json.loads(current_tool_input_json) if current_tool_input_json else {}
                    except json.JSONDecodeError:
                        arguments = {"raw": current_tool_input_json}

                    accumulated_tool_calls[current_tool_id] = {
                        "id": current_tool_id,
                        "name": current_tool_name or "",
                        "arguments": arguments,
                    }
                    current_tool_id = None
                    current_tool_name = None
                    current_tool_input_json = ""

            elif event_type == "message_delta":
                # Message delta may contain usage and stop_reason
                usage_data = data.get("usage", {})

                if usage_data:
                    usage = TokenUsage(
                        prompt_tokens=usage_data.get("input_tokens", 0),
                        completion_tokens=usage_data.get("output_tokens", 0),
                    )

                    # Yield accumulated tool calls if any
                    if accumulated_tool_calls:
                        tool_calls = [
                            ToolCall(
                                id=tc["id"],
                                name=tc["name"],
                                arguments=tc["arguments"],
                            )
                            for tc in accumulated_tool_calls.values()
                        ]
                        yield StreamChunk(
                            content="",
                            model=request.model,
                            tool_calls=tool_calls,
                            usage=usage,
                            is_final=True,
                        )
                        accumulated_tool_calls = {}
                    else:
                        yield StreamChunk(
                            content="",
                            model=request.model,
                            usage=usage,
                            is_final=True,
                        )

            elif event_type == "message_stop":
                # Message ended
                # If we have accumulated tool calls but haven't yielded them yet
                if accumulated_tool_calls:
                    tool_calls = [
                        ToolCall(
                            id=tc["id"],
                            name=tc["name"],
                            arguments=tc["arguments"],
                        )
                        for tc in accumulated_tool_calls.values()
                    ]
                    yield StreamChunk(
                        content="",
                        model=request.model,
                        tool_calls=tool_calls,
                        is_final=True,
                    )
                    accumulated_tool_calls = {}

            elif event_type == "ping":
                continue

            elif event_type == "error":
                error_info = data.get("error", {})
                error_msg = error_info.get("message", "Stream error")
                raise LLMProviderError("anthropic", error_msg)

    def get_model_info(self) -> dict[str, Any]:
        """返回 Provider 和模型信息"""
        return {
            "provider": "anthropic",
            "model": self._model,
            "max_tokens": self._max_tokens,
            "thinking_enabled": self._thinking_enabled,
        }
