"""
Anthropic Provider Module

Implements LLM interface for Anthropic API:
- Claude 3, Claude 3.5 models
- Tool use
- Streaming support
"""

import json
import time
from typing import Optional

from anthropic import Anthropic, APIError, RateLimitError as AnthropicRateLimitError, AuthenticationError as AnthropicAuthError

from ai_coder.llm.interface import (
    LLMProvider,
    LLMConfig,
    Message,
    ToolCall,
    CompletionResponse,
    LLMError,
    RateLimitError,
    AuthenticationError,
    register_provider,
)


@register_provider("anthropic")
class AnthropicProvider(LLMProvider):
    """
    Anthropic API provider implementation.
    
    Supports:
    - Claude 3.5 Sonnet, Claude 3 Opus, Claude 3 Sonnet, Claude 3 Haiku
    - Tool use
    """

    # Model name mapping
    MODEL_MAP = {
        "claude-3-opus": "claude-3-opus-20240229",
        "claude-3-sonnet": "claude-3-sonnet-20240229",
        "claude-3-haiku": "claude-3-haiku-20240307",
        "claude-3.5-sonnet": "claude-3-5-sonnet-20241022",
        "claude-3-5-sonnet": "claude-3-5-sonnet-20241022",
    }

    def __init__(self, config: LLMConfig):
        super().__init__(config)
        self.client = Anthropic(
            api_key=config.api_key,
            timeout=config.timeout,
        )

    def complete(
        self,
        messages: list[Message],
        tools: Optional[list[dict]] = None,
    ) -> CompletionResponse:
        """Generate completion using Anthropic API."""
        # Extract system message
        system_content = ""
        conversation_messages = []

        for msg in messages:
            if msg.role == "system":
                system_content = msg.content
            else:
                conversation_messages.append(msg)

        formatted_messages = self._format_messages(conversation_messages)

        # Map model name if needed
        model = self.MODEL_MAP.get(self.config.model, self.config.model)

        kwargs = {
            "model": model,
            "messages": formatted_messages,
            "max_tokens": self.config.max_tokens,
            "temperature": self.config.temperature,
        }

        if system_content:
            kwargs["system"] = system_content

        if tools:
            # Convert OpenAI tool format to Anthropic format
            kwargs["tools"] = self._convert_tools(tools)

        for attempt in range(self.config.max_retries):
            try:
                response = self.client.messages.create(**kwargs)
                return self._parse_response(response)

            except AnthropicRateLimitError as e:
                if attempt < self.config.max_retries - 1:
                    wait_time = 2 ** attempt
                    time.sleep(wait_time)
                    continue
                raise RateLimitError(str(e))

            except AnthropicAuthError as e:
                raise AuthenticationError(f"Anthropic authentication failed: {e}")

            except APIError as e:
                if attempt < self.config.max_retries - 1:
                    time.sleep(1)
                    continue
                raise LLMError(f"Anthropic API error: {e}")

        raise LLMError("Max retries exceeded")

    def _convert_tools(self, openai_tools: list[dict]) -> list[dict]:
        """Convert OpenAI tool format to Anthropic format."""
        anthropic_tools = []

        for tool in openai_tools:
            if tool.get("type") == "function":
                func = tool["function"]
                anthropic_tools.append({
                    "name": func["name"],
                    "description": func.get("description", ""),
                    "input_schema": func.get("parameters", {"type": "object", "properties": {}}),
                })
            else:
                # Already in Anthropic format
                anthropic_tools.append(tool)

        return anthropic_tools

    def _format_messages(self, messages: list[Message]) -> list[dict]:
        """Format messages for Anthropic API."""
        formatted = []

        for msg in messages:
            if msg.role == "tool":
                # Tool result message
                formatted.append({
                    "role": "user",
                    "content": [{
                        "type": "tool_result",
                        "tool_use_id": msg.tool_call_id or "",
                        "content": msg.content,
                    }],
                })
            elif msg.tool_calls:
                # Assistant message with tool use
                content = []
                if msg.content:
                    content.append({"type": "text", "text": msg.content})

                for tc in msg.tool_calls:
                    content.append({
                        "type": "tool_use",
                        "id": tc.get("id", ""),
                        "name": tc.get("function", {}).get("name", tc.get("name", "")),
                        "input": tc.get("function", {}).get("arguments", tc.get("arguments", {})),
                    })

                formatted.append({
                    "role": "assistant",
                    "content": content,
                })
            else:
                # Regular message
                formatted.append({
                    "role": msg.role if msg.role in ["user", "assistant"] else "user",
                    "content": msg.content,
                })

        return formatted

    def _parse_response(self, response) -> CompletionResponse:
        """Parse Anthropic API response."""
        content_parts = []
        tool_calls = []

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

        content = "\n".join(content_parts) if content_parts else None

        usage = None
        if response.usage:
            usage = {
                "prompt_tokens": response.usage.input_tokens,
                "completion_tokens": response.usage.output_tokens,
                "total_tokens": response.usage.input_tokens + response.usage.output_tokens,
            }

        return CompletionResponse(
            content=content,
            tool_calls=tool_calls,
            finish_reason=response.stop_reason or "stop",
            usage=usage,
            raw_response=response,
        )

    def count_tokens(self, text: str) -> int:
        """
        Estimate token count for Anthropic models.
        Uses approximate character ratio.
        """
        # Claude models: ~3.5 chars per token
        return len(text) // 3
