"""
Anthropic provider for Claude models.

This provider connects to Anthropic's API for accessing Claude models.
"""

import json
from typing import Any, AsyncIterator, Dict, Iterator, List, Optional, Union

import anthropic
from anthropic.types import Message, MessageStreamEvent

from oai.constants import ANTHROPIC_BASE_URL
from oai.providers.base import (
    AIProvider,
    ChatMessage,
    ChatResponse,
    ChatResponseChoice,
    ModelInfo,
    ProviderCapabilities,
    StreamChunk,
    ToolCall,
    ToolFunction,
    UsageStats,
)
from oai.utils.logging import get_logger

logger = get_logger()


# Model name aliases
MODEL_ALIASES = {
    "claude-sonnet": "claude-sonnet-4-5-20250929",
    "claude-haiku": "claude-haiku-4-5-20251001",
    "claude-opus": "claude-opus-4-5-20251101",
    # Legacy aliases
    "claude-3-haiku": "claude-3-haiku-20240307",
    "claude-3-7-sonnet": "claude-3-7-sonnet-20250219",
}


class AnthropicProvider(AIProvider):
    """
    Anthropic API provider.

    Provides access to Claude 3.5 Sonnet, Claude 3 Opus, and other Anthropic models.
    """

    def __init__(
        self,
        api_key: str,
        base_url: Optional[str] = None,
        app_name: str = "oAI",
        app_url: str = "",
        **kwargs: Any,
    ):
        """
        Initialize Anthropic provider.

        Args:
            api_key: Anthropic API key
            base_url: Optional custom base URL
            app_name: Application name (for headers)
            app_url: Application URL (for headers)
            **kwargs: Additional arguments
        """
        super().__init__(api_key, base_url or ANTHROPIC_BASE_URL)
        self.client = anthropic.Anthropic(api_key=api_key)
        self.async_client = anthropic.AsyncAnthropic(api_key=api_key)
        self._models_cache: Optional[List[ModelInfo]] = None

    def _create_web_search_tool(self, config: Dict[str, Any]) -> Dict[str, Any]:
        """
        Create Anthropic native web search tool definition.

        Args:
            config: Optional configuration for web search (max_uses, allowed_domains, etc.)

        Returns:
            Tool definition dict
        """
        tool: Dict[str, Any] = {
            "type": "web_search_20250305",
            "name": "web_search",
        }

        # Add optional parameters if provided
        if "max_uses" in config:
            tool["max_uses"] = config["max_uses"]
        else:
            tool["max_uses"] = 5  # Default

        if "allowed_domains" in config:
            tool["allowed_domains"] = config["allowed_domains"]

        if "blocked_domains" in config:
            tool["blocked_domains"] = config["blocked_domains"]

        if "user_location" in config:
            tool["user_location"] = config["user_location"]

        return tool

    @property
    def name(self) -> str:
        """Get provider name."""
        return "Anthropic"

    @property
    def capabilities(self) -> ProviderCapabilities:
        """Get provider capabilities."""
        return ProviderCapabilities(
            streaming=True,
            tools=True,
            images=True,
            online=True,  # Web search via DuckDuckGo/Google
            max_context=200000,
        )

    def list_models(self, filter_text_only: bool = True) -> List[ModelInfo]:
        """
        List available Anthropic models.

        Args:
            filter_text_only: Whether to filter for text models only

        Returns:
            List of ModelInfo objects
        """
        if self._models_cache:
            return self._models_cache

        # Anthropic doesn't have a models list API, so we hardcode the available models
        models = [
            # Current Claude 4.5 models
            ModelInfo(
                id="claude-sonnet-4-5-20250929",
                name="Claude Sonnet 4.5",
                description="Smart model for complex agents and coding (recommended)",
                context_length=200000,
                pricing={"input": 3.0, "output": 15.0},
                supported_parameters=["temperature", "max_tokens", "stream", "tools"],
                input_modalities=["text", "image"],
            ),
            ModelInfo(
                id="claude-haiku-4-5-20251001",
                name="Claude Haiku 4.5",
                description="Fastest model with near-frontier intelligence",
                context_length=200000,
                pricing={"input": 1.0, "output": 5.0},
                supported_parameters=["temperature", "max_tokens", "stream", "tools"],
                input_modalities=["text", "image"],
            ),
            ModelInfo(
                id="claude-opus-4-5-20251101",
                name="Claude Opus 4.5",
                description="Premium model with maximum intelligence",
                context_length=200000,
                pricing={"input": 5.0, "output": 25.0},
                supported_parameters=["temperature", "max_tokens", "stream", "tools"],
                input_modalities=["text", "image"],
            ),
            # Legacy models (still available)
            ModelInfo(
                id="claude-3-7-sonnet-20250219",
                name="Claude Sonnet 3.7",
                description="Legacy model - recommend migrating to 4.5",
                context_length=200000,
                pricing={"input": 3.0, "output": 15.0},
                supported_parameters=["temperature", "max_tokens", "stream", "tools"],
                input_modalities=["text", "image"],
            ),
            ModelInfo(
                id="claude-3-haiku-20240307",
                name="Claude 3 Haiku",
                description="Legacy fast model - recommend migrating to 4.5",
                context_length=200000,
                pricing={"input": 0.25, "output": 1.25},
                supported_parameters=["temperature", "max_tokens", "stream", "tools"],
                input_modalities=["text", "image"],
            ),
        ]

        self._models_cache = models
        logger.info(f"Loaded {len(models)} Anthropic models")
        return models

    def get_model(self, model_id: str) -> Optional[ModelInfo]:
        """
        Get information about a specific model.

        Args:
            model_id: Model identifier

        Returns:
            ModelInfo or None
        """
        # Resolve alias
        resolved_id = MODEL_ALIASES.get(model_id, model_id)

        models = self.list_models()
        for model in models:
            if model.id == resolved_id or model.id == model_id:
                return model
        return None

    def chat(
        self,
        model: str,
        messages: List[ChatMessage],
        stream: bool = False,
        max_tokens: Optional[int] = None,
        temperature: Optional[float] = None,
        tools: Optional[List[Dict[str, Any]]] = None,
        tool_choice: Optional[str] = None,
        **kwargs: Any,
    ) -> Union[ChatResponse, Iterator[StreamChunk]]:
        """
        Send chat completion request to Anthropic.

        Args:
            model: Model ID
            messages: Chat messages
            stream: Whether to stream response
            max_tokens: Maximum tokens
            temperature: Sampling temperature
            tools: Tool definitions
            tool_choice: Tool selection mode
            **kwargs: Additional parameters (including enable_web_search)

        Returns:
            ChatResponse or Iterator[StreamChunk]
        """
        # Resolve model alias
        model_id = MODEL_ALIASES.get(model, model)

        # Extract system message (Anthropic requires it separate from messages)
        system_prompt, anthropic_messages = self._convert_messages(messages)

        # Build request parameters
        params: Dict[str, Any] = {
            "model": model_id,
            "messages": anthropic_messages,
            "max_tokens": max_tokens or 4096,
        }

        if system_prompt:
            params["system"] = system_prompt

        if temperature is not None:
            params["temperature"] = temperature

        # Prepare tools list
        tools_list = []

        # Add web search tool if requested via kwargs
        if kwargs.get("enable_web_search", False):
            web_search_config = kwargs.get("web_search_config", {})
            tools_list.append(self._create_web_search_tool(web_search_config))
            logger.info("Added Anthropic native web search tool")

        # Add user-provided tools
        if tools:
            # Convert tools to Anthropic format
            converted_tools = self._convert_tools(tools)
            tools_list.extend(converted_tools)

        if tools_list:
            params["tools"] = tools_list

        if tool_choice and tool_choice != "auto":
            # Anthropic uses different tool_choice format
            if tool_choice == "none":
                pass  # Don't include tools
            elif tool_choice == "required":
                params["tool_choice"] = {"type": "any"}
            else:
                params["tool_choice"] = {"type": "tool", "name": tool_choice}

        logger.debug(f"Anthropic request: model={model_id}, messages={len(anthropic_messages)}")

        try:
            if stream:
                return self._stream_chat(params)
            else:
                return self._sync_chat(params)

        except Exception as e:
            logger.error(f"Anthropic request failed: {e}")
            return ChatResponse(
                id="error",
                choices=[
                    ChatResponseChoice(
                        index=0,
                        message=ChatMessage(role="assistant", content=f"Error: {str(e)}"),
                        finish_reason="error",
                    )
                ],
            )

    def _convert_messages(self, messages: List[ChatMessage]) -> tuple[str, List[Dict[str, Any]]]:
        """
        Convert messages to Anthropic format.

        Anthropic requires system messages to be separate from the conversation.

        Args:
            messages: List of ChatMessage objects

        Returns:
            Tuple of (system_prompt, anthropic_messages)
        """
        system_prompt = ""
        anthropic_messages = []

        for msg in messages:
            if msg.role == "system":
                # Accumulate system messages
                if system_prompt:
                    system_prompt += "\n\n"
                system_prompt += msg.content or ""
            else:
                # Convert to Anthropic format
                message_dict: Dict[str, Any] = {"role": msg.role}

                # Handle content
                if msg.content:
                    message_dict["content"] = msg.content

                # Handle tool calls (assistant messages)
                if msg.tool_calls:
                    # Anthropic format for tool use
                    content_blocks = []
                    if msg.content:
                        content_blocks.append({"type": "text", "text": msg.content})

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

                    message_dict["content"] = content_blocks

                # Handle tool results (tool messages)
                if msg.role == "tool" and msg.tool_call_id:
                    # Convert to Anthropic's tool_result format
                    anthropic_messages.append({
                        "role": "user",
                        "content": [{
                            "type": "tool_result",
                            "tool_use_id": msg.tool_call_id,
                            "content": msg.content or "",
                        }]
                    })
                    continue

                anthropic_messages.append(message_dict)

        return system_prompt, anthropic_messages

    def _convert_tools(self, tools: List[Dict[str, Any]]) -> List[Dict[str, Any]]:
        """
        Convert OpenAI-style tools to Anthropic format.

        Args:
            tools: OpenAI tool definitions

        Returns:
            Anthropic tool definitions
        """
        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", {}),
                })

        return anthropic_tools

    def _sync_chat(self, params: Dict[str, Any]) -> ChatResponse:
        """
        Send synchronous chat request.

        Args:
            params: Request parameters

        Returns:
            ChatResponse
        """
        message: Message = self.client.messages.create(**params)

        # Extract content
        content = ""
        tool_calls = []

        for block in message.content:
            if block.type == "text":
                content += block.text
            elif block.type == "tool_use":
                # Convert to ToolCall format
                tool_calls.append(
                    ToolCall(
                        id=block.id,
                        type="function",
                        function=ToolFunction(
                            name=block.name,
                            arguments=json.dumps(block.input),
                        ),
                    )
                )

        # Build ChatMessage
        chat_message = ChatMessage(
            role="assistant",
            content=content if content else None,
            tool_calls=tool_calls if tool_calls else None,
        )

        # Extract usage
        usage = None
        if message.usage:
            usage = UsageStats(
                prompt_tokens=message.usage.input_tokens,
                completion_tokens=message.usage.output_tokens,
                total_tokens=message.usage.input_tokens + message.usage.output_tokens,
            )

        return ChatResponse(
            id=message.id,
            choices=[
                ChatResponseChoice(
                    index=0,
                    message=chat_message,
                    finish_reason=message.stop_reason,
                )
            ],
            usage=usage,
            model=message.model,
        )

    def _stream_chat(self, params: Dict[str, Any]) -> Iterator[StreamChunk]:
        """
        Stream chat response from Anthropic.

        Args:
            params: Request parameters

        Yields:
            StreamChunk objects
        """
        stream = self.client.messages.stream(**params)

        with stream as event_stream:
            for event in event_stream:
                event_data: MessageStreamEvent = event

                # Handle different event types
                if event_data.type == "content_block_delta":
                    delta = event_data.delta
                    if hasattr(delta, "text"):
                        yield StreamChunk(
                            id="stream",
                            delta_content=delta.text,
                        )

                elif event_data.type == "message_stop":
                    # Final event with usage
                    pass

                elif event_data.type == "message_delta":
                    # Contains stop reason and usage
                    usage = None
                    if hasattr(event_data, "usage"):
                        usage_data = event_data.usage
                        usage = UsageStats(
                            prompt_tokens=0,
                            completion_tokens=usage_data.output_tokens,
                            total_tokens=usage_data.output_tokens,
                        )

                    yield StreamChunk(
                        id="stream",
                        finish_reason=event_data.delta.stop_reason if hasattr(event_data.delta, "stop_reason") else None,
                        usage=usage,
                    )

    async def chat_async(
        self,
        model: str,
        messages: List[ChatMessage],
        stream: bool = False,
        max_tokens: Optional[int] = None,
        temperature: Optional[float] = None,
        tools: Optional[List[Dict[str, Any]]] = None,
        tool_choice: Optional[str] = None,
        **kwargs: Any,
    ) -> Union[ChatResponse, AsyncIterator[StreamChunk]]:
        """
        Send async chat request to Anthropic.

        Args:
            model: Model ID
            messages: Chat messages
            stream: Whether to stream
            max_tokens: Max tokens
            temperature: Temperature
            tools: Tool definitions
            tool_choice: Tool choice
            **kwargs: Additional args

        Returns:
            ChatResponse or AsyncIterator[StreamChunk]
        """
        # Resolve model alias
        model_id = MODEL_ALIASES.get(model, model)

        # Convert messages
        system_prompt, anthropic_messages = self._convert_messages(messages)

        # Build params
        params: Dict[str, Any] = {
            "model": model_id,
            "messages": anthropic_messages,
            "max_tokens": max_tokens or 4096,
        }

        if system_prompt:
            params["system"] = system_prompt
        if temperature is not None:
            params["temperature"] = temperature
        if tools:
            params["tools"] = self._convert_tools(tools)

        if stream:
            return self._stream_chat_async(params)
        else:
            message = await self.async_client.messages.create(**params)
            return self._convert_message(message)

    async def _stream_chat_async(self, params: Dict[str, Any]) -> AsyncIterator[StreamChunk]:
        """
        Stream async chat response.

        Args:
            params: Request parameters

        Yields:
            StreamChunk objects
        """
        stream = await self.async_client.messages.stream(**params)

        async with stream as event_stream:
            async for event in event_stream:
                if event.type == "content_block_delta":
                    delta = event.delta
                    if hasattr(delta, "text"):
                        yield StreamChunk(
                            id="stream",
                            delta_content=delta.text,
                        )

    def _convert_message(self, message: Message) -> ChatResponse:
        """Helper to convert Anthropic message to ChatResponse."""
        content = ""
        tool_calls = []

        for block in message.content:
            if block.type == "text":
                content += block.text
            elif block.type == "tool_use":
                tool_calls.append(
                    ToolCall(
                        id=block.id,
                        type="function",
                        function=ToolFunction(
                            name=block.name,
                            arguments=json.dumps(block.input),
                        ),
                    )
                )

        chat_message = ChatMessage(
            role="assistant",
            content=content if content else None,
            tool_calls=tool_calls if tool_calls else None,
        )

        usage = None
        if message.usage:
            usage = UsageStats(
                prompt_tokens=message.usage.input_tokens,
                completion_tokens=message.usage.output_tokens,
                total_tokens=message.usage.input_tokens + message.usage.output_tokens,
            )

        return ChatResponse(
            id=message.id,
            choices=[
                ChatResponseChoice(
                    index=0,
                    message=chat_message,
                    finish_reason=message.stop_reason,
                )
            ],
            usage=usage,
            model=message.model,
        )

    def get_credits(self) -> Optional[Dict[str, Any]]:
        """
        Get account credits from Anthropic.

        Note: Anthropic does not currently provide a public API endpoint
        for checking account credits/balance. This information is only
        available through the Anthropic Console web interface.

        Returns:
            None (credits API not available)
        """
        # Anthropic doesn't provide a public credits API endpoint
        # Users must check their balance at console.anthropic.com
        return None

    def clear_cache(self) -> None:
        """Clear model cache."""
        self._models_cache = None

    def get_raw_models(self) -> List[Dict[str, Any]]:
        """
        Get raw model data as dictionaries.

        Returns:
            List of model dictionaries
        """
        models = self.list_models()
        return [
            {
                "id": model.id,
                "name": model.name,
                "description": model.description,
                "context_length": model.context_length,
                "pricing": model.pricing,
                "supported_parameters": model.supported_parameters,
                "input_modalities": model.input_modalities,
                "output_modalities": model.output_modalities,
            }
            for model in models
        ]

    def get_raw_model(self, model_id: str) -> Optional[Dict[str, Any]]:
        """
        Get raw model data for a specific model.

        Args:
            model_id: Model identifier

        Returns:
            Model dictionary or None
        """
        model = self.get_model(model_id)
        if model:
            return {
                "id": model.id,
                "name": model.name,
                "description": model.description,
                "context_length": model.context_length,
                "pricing": model.pricing,
                "supported_parameters": model.supported_parameters,
                "input_modalities": model.input_modalities,
                "output_modalities": model.output_modalities,
            }
        return None