"""Anthropic LLM plugin — implements BaseLLMProvider for Claude models."""

from __future__ import annotations

import json
import logging

from vastian.config import Settings
from vastian.providers.llm_provider import BaseLLMProvider

logger = logging.getLogger(__name__)


class AnthropicProvider(BaseLLMProvider):
    """Anthropic Claude API provider.

    The orchestrator builds messages in OpenAI format. This provider converts
    them to Anthropic format before sending, and converts Anthropic responses
    back to the common format.
    """

    def __init__(self, settings: Settings) -> None:
        super().__init__(settings)
        from anthropic import AsyncAnthropic

        self.client = AsyncAnthropic(api_key=settings.anthropic_api_key)
        self.last_usage: dict[str, int] = {}  # populated after each call
        # Remap OpenAI model names to Anthropic equivalents
        from vastian.config import ANTHROPIC_MODELS

        if "gpt" in self.default_model:
            self.default_model = ANTHROPIC_MODELS["default"]  # claude-opus-4-5
        if "gpt" in self.fast_model:
            self.fast_model = ANTHROPIC_MODELS["fast"]  # claude-sonnet-4

    def _extract_usage(self, response: object, model: str) -> None:
        """Extract token usage from Anthropic response."""
        usage = getattr(response, "usage", None)
        if usage:
            self.last_usage = {
                "provider": "anthropic",
                "model": model,
                "input_tokens": getattr(usage, "input_tokens", 0) or 0,
                "output_tokens": getattr(usage, "output_tokens", 0) or 0,
            }
        else:
            self.last_usage = {}

    @staticmethod
    def _convert_messages_to_anthropic(messages: list[dict]) -> list[dict]:
        """
        Convert messages from OpenAI format to Anthropic format.

        Key differences:
        - OpenAI assistant messages have 'tool_calls' as a top-level list
          -> Anthropic uses content blocks with type='tool_use'
        - OpenAI tool results use role='tool' with tool_call_id
          -> Anthropic uses role='user' with content blocks of type='tool_result'
        - Anthropic requires strict role alternation (user/assistant)
          so consecutive same-role messages must be merged
        """
        converted: list[dict] = []

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

            # -- Skip system messages (handled separately in Anthropic) --
            if role == "system":
                continue

            # -- User messages -- pass through --
            if role == "user":
                content = msg.get("content", "")
                if isinstance(content, str):
                    converted.append({"role": "user", "content": content})
                else:
                    converted.append({"role": "user", "content": content})

            # -- Assistant messages -- convert tool_calls to content blocks --
            elif role == "assistant":
                content_blocks: list[dict] = []

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

                # Convert OpenAI tool_calls to Anthropic tool_use blocks
                for tc in msg.get("tool_calls", []):
                    fn = tc.get("function", {})
                    fn_name = fn.get("name", "") if isinstance(fn, dict) else tc.get("function", "")
                    fn_args = (
                        fn.get("arguments", "{}")
                        if isinstance(fn, dict)
                        else tc.get("arguments", "{}")
                    )

                    # Parse arguments -- may be string or dict
                    if isinstance(fn_args, str):
                        try:
                            parsed_args = json.loads(fn_args)
                        except json.JSONDecodeError:
                            parsed_args = {}
                    else:
                        parsed_args = fn_args

                    content_blocks.append(
                        {
                            "type": "tool_use",
                            "id": tc.get("id", ""),
                            "name": fn_name,
                            "input": parsed_args,
                        }
                    )

                # Anthropic requires at least one content block
                if not content_blocks:
                    content_blocks.append({"type": "text", "text": ""})

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

            # -- Tool results -> Anthropic 'user' message with tool_result blocks --
            elif role == "tool":
                tool_result_block = {
                    "type": "tool_result",
                    "tool_use_id": msg.get("tool_call_id", ""),
                    "content": msg.get("content", ""),
                }
                converted.append(
                    {
                        "role": "user",
                        "content": [tool_result_block],
                    }
                )

        # -- Merge consecutive same-role messages --
        # Anthropic requires strict role alternation
        merged: list[dict] = []
        for msg in converted:
            if merged and merged[-1]["role"] == msg["role"]:
                prev_content = merged[-1]["content"]
                new_content = msg["content"]

                # Convert strings to content block lists for merging
                if isinstance(prev_content, str):
                    prev_content = [{"type": "text", "text": prev_content}]
                if isinstance(new_content, str):
                    new_content = [{"type": "text", "text": new_content}]

                merged[-1]["content"] = prev_content + new_content
            else:
                merged.append(msg)

        return merged

    async def generate(
        self,
        system_prompt: str,
        messages: list[dict[str, str]],
        model: str | None = None,
        temperature: float = 0.7,
        max_tokens: int = 4096,
    ) -> str:
        model = model or self.default_model
        anthropic_messages = self._convert_messages_to_anthropic(messages)

        response = await self.client.messages.create(
            model=model,
            system=system_prompt,
            messages=anthropic_messages,
            temperature=temperature,
            max_tokens=max_tokens,
        )
        self._extract_usage(response, model)
        # Extract text from content blocks
        return "".join(block.text for block in response.content if hasattr(block, "text"))

    async def generate_with_tools(
        self,
        system_prompt: str,
        messages: list[dict[str, str]],
        tools: list[dict],
        model: str | None = None,
        temperature: float = 0.7,
    ) -> dict:
        model = model or self.default_model

        # Convert OpenAI tool definitions to Anthropic format
        anthropic_tools = []
        for tool in tools:
            fn = tool.get("function", {})
            anthropic_tools.append(
                {
                    "name": fn.get("name", ""),
                    "description": fn.get("description", ""),
                    "input_schema": fn.get("parameters", {}),
                }
            )

        # Convert messages from OpenAI format to Anthropic format
        anthropic_messages = self._convert_messages_to_anthropic(messages)

        kwargs: dict = {
            "model": model,
            "system": system_prompt,
            "messages": anthropic_messages,
            "temperature": temperature,
            "max_tokens": 4096,
        }
        # Only include tools if there are any (Anthropic errors on empty list)
        if anthropic_tools:
            kwargs["tools"] = anthropic_tools

        response = await self.client.messages.create(**kwargs)
        self._extract_usage(response, model)

        result: dict = {"content": ""}
        tool_calls = []
        for block in response.content:
            if hasattr(block, "text"):
                result["content"] += block.text
            elif block.type == "tool_use":
                tool_calls.append(
                    {
                        "id": block.id,
                        "function": block.name,
                        "arguments": block.input,
                    }
                )

        if tool_calls:
            result["tool_calls"] = tool_calls
        return result


def register() -> None:
    """Register the Anthropic LLM plugin with the loader."""
    from vastian.loader import register as loader_register

    loader_register("anthropic_llm", AnthropicProvider)
