"""Anthropic model provider — implements ModelProvider using the Anthropic SDK."""

from __future__ import annotations

import json

from hatchdx.agent.providers.base import (
    AgentResponse,
    ProviderAuthError,
    ProviderError,
    ToolCall,
)


# ---------------------------------------------------------------------------
# Pricing (per million tokens) — kept simple for cost estimation
# ---------------------------------------------------------------------------

_PRICING: dict[str, dict[str, float]] = {
    # Claude 4.5 family
    "claude-opus-4-5-20251101": {"input": 5.0, "output": 25.0},
    "claude-sonnet-4-5-20250929": {"input": 3.0, "output": 15.0},
    "claude-haiku-4-5-20251001": {"input": 1.0, "output": 5.0},
    # Claude 4 family (legacy)
    "claude-sonnet-4-20250514": {"input": 3.0, "output": 15.0},
    "claude-opus-4-20250514": {"input": 15.0, "output": 75.0},
    "claude-haiku-4-20250506": {"input": 0.80, "output": 4.0},
    # Claude 3.5 family (legacy)
    "claude-3-5-haiku-20241022": {"input": 0.80, "output": 4.0},
}

_DEFAULT_PRICING = {"input": 3.0, "output": 15.0}


def estimate_cost(model: str, input_tokens: int, output_tokens: int) -> float:
    """Estimate cost in USD for a given model and token counts."""
    pricing = _PRICING.get(model, _DEFAULT_PRICING)
    return (input_tokens * pricing["input"] + output_tokens * pricing["output"]) / 1_000_000


# ---------------------------------------------------------------------------
# Tool format conversion
# ---------------------------------------------------------------------------


def mcp_tools_to_anthropic(tools: list[dict]) -> list[dict]:
    """Convert MCP tool definitions to Anthropic tool format.

    MCP format:
        {"name": "...", "description": "...", "inputSchema": {...}}

    Anthropic format:
        {"name": "...", "description": "...", "input_schema": {...}}
    """
    result = []
    for tool in tools:
        entry: dict = {
            "name": tool["name"],
            "description": tool.get("description", ""),
        }
        schema = tool.get("inputSchema") or tool.get("input_schema")
        if schema:
            entry["input_schema"] = schema
        else:
            entry["input_schema"] = {"type": "object", "properties": {}}
        result.append(entry)
    return result


def _anthropic_messages(messages: list[dict]) -> list[dict]:
    """Convert internal message format to Anthropic API messages.

    Anthropic uses a different structure than OpenAI:
    - System prompt is a top-level parameter (not in messages)
    - Tool results are sent as user messages with tool_result content blocks
    - Assistant messages with tool calls use content blocks
    """
    result: list[dict] = []

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

        if role == "user":
            result.append({"role": "user", "content": msg["content"]})

        elif role == "assistant":
            content_blocks: list[dict] = []

            # Add text block if present
            if msg.get("content"):
                content_blocks.append({"type": "text", "text": msg["content"]})

            # Add tool_use blocks
            if msg.get("tool_calls"):
                for tc in msg["tool_calls"]:
                    tc_id = tc.id if hasattr(tc, "id") else tc["id"]
                    tc_name = tc.name if hasattr(tc, "name") else tc["name"]
                    tc_args = tc.arguments if hasattr(tc, "arguments") else tc["arguments"]
                    content_blocks.append({
                        "type": "tool_use",
                        "id": tc_id,
                        "name": tc_name,
                        "input": tc_args,
                    })

            result.append({"role": "assistant", "content": content_blocks})

        elif role == "tool":
            # Anthropic expects tool results as user messages with tool_result blocks
            # Group consecutive tool results into a single user message
            tool_result_block = {
                "type": "tool_result",
                "tool_use_id": msg["tool_call_id"],
                "content": msg["content"],
            }
            if msg.get("is_error"):
                tool_result_block["is_error"] = True

            # If the last message in result is a user message with tool_result blocks,
            # append to it. Otherwise create a new user message.
            if (
                result
                and result[-1]["role"] == "user"
                and isinstance(result[-1]["content"], list)
                and result[-1]["content"]
                and result[-1]["content"][0].get("type") == "tool_result"
            ):
                result[-1]["content"].append(tool_result_block)
            else:
                result.append({
                    "role": "user",
                    "content": [tool_result_block],
                })

    return result


# ---------------------------------------------------------------------------
# Provider
# ---------------------------------------------------------------------------


class AnthropicProvider:
    """Anthropic model provider using the anthropic SDK."""

    def __init__(
        self,
        model: str,
        api_key: str | None = None,
    ) -> None:
        self.model = model
        self._api_key = api_key
        self._client = None

    def _get_client(self):
        """Lazy-init the Anthropic client."""
        if self._client is not None:
            return self._client

        try:
            import anthropic
        except ImportError:
            raise ProviderError(
                "Anthropic SDK not installed.\n\n"
                "Install it with:\n"
                "  uv add hatchdx --extra anthropic"
            )

        import os

        api_key = self._api_key or os.environ.get("ANTHROPIC_API_KEY")
        if not api_key:
            raise ProviderAuthError(
                "Missing API key for Anthropic.\n\n"
                "Set your API key:\n"
                "  export ANTHROPIC_API_KEY=your-key-here"
            )

        self._client = anthropic.AsyncAnthropic(api_key=api_key)
        return self._client

    async def chat(
        self,
        messages: list[dict],
        tools: list[dict],
        system: str,
        max_tokens: int,
        temperature: float,
    ) -> AgentResponse:
        """Send a chat request to Anthropic and return an AgentResponse."""
        client = self._get_client()
        anthropic_messages = _anthropic_messages(messages)
        anthropic_tools = mcp_tools_to_anthropic(tools) if tools else []

        kwargs: dict = {
            "model": self.model,
            "messages": anthropic_messages,
            "system": system,
            "max_tokens": max_tokens,
            "temperature": temperature,
        }
        if anthropic_tools:
            kwargs["tools"] = anthropic_tools

        try:
            import anthropic as anthropic_module

            response = await client.messages.create(**kwargs)
        except anthropic_module.AuthenticationError as e:
            raise ProviderAuthError(
                f"Anthropic authentication failed: {e}\n\n"
                "Check your API key:\n"
                "  export ANTHROPIC_API_KEY=your-key-here"
            ) from e
        except anthropic_module.APIError as e:
            raise ProviderError(f"Anthropic API error: {e}") from e

        # Extract text and tool calls from content blocks
        text_parts: list[str] = []
        tool_calls: list[ToolCall] = []

        for block in response.content:
            if block.type == "text":
                text_parts.append(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 {},
                    )
                )

        text = "\n".join(text_parts) if text_parts else None

        # Map stop reason
        stop_reason = "end_turn"
        if response.stop_reason == "tool_use":
            stop_reason = "tool_use"
        elif response.stop_reason == "max_tokens":
            stop_reason = "max_tokens"

        # Usage
        usage = {"input": 0, "output": 0}
        if response.usage:
            usage["input"] = response.usage.input_tokens or 0
            usage["output"] = response.usage.output_tokens or 0

        return AgentResponse(
            text=text,
            tool_calls=tool_calls,
            stop_reason=stop_reason,
            usage=usage,
        )
