"""Anthropic OAuth provider - direct API calls with Bearer auth.

This provider bypasses litellm to properly handle OAuth tokens
which require Authorization: Bearer header instead of x-api-key.
"""

import json
from typing import Any

import httpx
from loguru import logger

from nanobot.providers.base import LLMProvider, LLMResponse, ToolCallRequest
from nanobot.providers.oauth_utils import get_auth_headers


class AnthropicOAuthProvider(LLMProvider):
    """
    Anthropic provider using OAuth token authentication.

    Unlike the LiteLLM provider, this calls the Anthropic API directly
    with proper Bearer token authentication for Claude Max/Pro subscriptions.
    """

    ANTHROPIC_API_URL = "https://api.anthropic.com/v1/messages"

    def __init__(
        self,
        oauth_token: str,
        default_model: str = "claude-opus-4-5",
        api_base: str | None = None,
        thinking_budget: int = 0,
    ):
        super().__init__(api_key=None, api_base=api_base)
        self.oauth_token = oauth_token
        self.default_model = default_model
        self.thinking_budget = thinking_budget
        self._client: httpx.AsyncClient | None = None

    def _get_headers(self) -> dict[str, str]:
        """Get request headers with Bearer auth."""
        return get_auth_headers(self.oauth_token, is_oauth=True)

    def _get_api_url(self) -> str:
        """Get API endpoint URL."""
        if self.api_base:
            return f"{self.api_base.rstrip('/')}/v1/messages"
        return self.ANTHROPIC_API_URL

    @staticmethod
    def _normalize_model(model: str) -> str:
        """Normalize model name for the Anthropic API.

        Anthropic model IDs use hyphens (claude-sonnet-4-5), but users often
        write dots (claude-sonnet-4.5). Normalize so both work.
        """
        return model.replace(".", "-")

    async def _get_client(self) -> httpx.AsyncClient:
        """Get or create async HTTP client."""
        if self._client is None:
            self._client = httpx.AsyncClient(timeout=300.0)
        return self._client

    def _prepare_messages(
        self,
        messages: list[dict[str, Any]]
    ) -> tuple[str | None, list[dict[str, Any]]]:
        """Prepare messages: extract system prompt and convert OpenAI format to Anthropic.

        The agent loop produces messages in OpenAI format:
          - assistant msgs with tool_calls [{type:"function", function:{name, arguments}}]
          - tool role msgs with tool_call_id, name, content

        Anthropic API expects:
          - assistant msgs with content blocks [{type:"tool_use", id, name, input}]
          - user msgs with content blocks [{type:"tool_result", tool_use_id, content}]

        Returns (system_prompt, anthropic_messages)
        """
        system_parts = []
        converted: list[dict[str, Any]] = []

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

            if role == "system":
                system_parts.append(msg.get("content", ""))
                continue

            if role == "assistant" and msg.get("tool_calls"):
                # Convert OpenAI tool_calls to Anthropic content blocks
                content_blocks: list[dict[str, Any]] = []
                # Preserve thinking blocks (list=raw API blocks with signatures, str=legacy)
                rc = msg.get("reasoning_content")
                if isinstance(rc, list):
                    content_blocks.extend(rc)
                elif isinstance(rc, str) and rc:
                    content_blocks.append({"type": "thinking", "thinking": rc})
                text = msg.get("content")
                if text:
                    content_blocks.append({"type": "text", "text": text})
                for tc in msg["tool_calls"]:
                    func = tc.get("function", {})
                    args = func.get("arguments", "{}")
                    if isinstance(args, str):
                        try:
                            args = json.loads(args)
                        except (json.JSONDecodeError, TypeError):
                            args = {}
                    content_blocks.append({
                        "type": "tool_use",
                        "id": tc.get("id", ""),
                        "name": func.get("name", ""),
                        "input": args,
                    })
                converted.append({"role": "assistant", "content": content_blocks})
                continue

            if role == "assistant" and msg.get("reasoning_content"):
                # Plain assistant message with thinking (no tool calls)
                rc = msg["reasoning_content"]
                if isinstance(rc, list):
                    content_blocks = list(rc)
                else:
                    content_blocks = [{"type": "thinking", "thinking": rc}]
                text = msg.get("content")
                if text:
                    content_blocks.append({"type": "text", "text": text})
                converted.append({"role": "assistant", "content": content_blocks})
                continue

            if role == "tool":
                # Convert tool result to Anthropic user message with tool_result block
                tool_result_block = {
                    "type": "tool_result",
                    "tool_use_id": msg.get("tool_call_id", ""),
                    "content": msg.get("content", ""),
                }
                # Merge into previous user message if it already has tool_result blocks
                if converted and converted[-1].get("role") == "user":
                    prev_content = converted[-1].get("content")
                    if isinstance(prev_content, list):
                        prev_content.append(tool_result_block)
                        continue
                converted.append({"role": "user", "content": [tool_result_block]})
                continue

            if role == "user":
                content = msg.get("content", "")
                # Convert OpenAI image_url blocks to Anthropic image blocks
                if isinstance(content, list):
                    content = self._convert_image_blocks(content)
                # Merge text into previous user message if it has tool_result blocks
                # (handles the "Reflect on the results" interleaved message)
                if converted and converted[-1].get("role") == "user":
                    prev_content = converted[-1].get("content")
                    if isinstance(prev_content, list):
                        if isinstance(content, str):
                            prev_content.append({"type": "text", "text": content})
                        elif isinstance(content, list):
                            prev_content.extend(content)
                        continue
                converted.append({"role": role, "content": content})
                continue

            # Pass through other messages (assistant without tool_calls, etc.)
            converted.append(msg)

        system_prompt = "\n\n".join(system_parts)
        return system_prompt, converted

    @staticmethod
    def _convert_image_blocks(content: list[dict[str, Any]]) -> list[dict[str, Any]]:
        """Convert OpenAI image_url blocks to Anthropic image blocks.

        OpenAI format:  {"type": "image_url", "image_url": {"url": "data:mime;base64,DATA"}}
        Anthropic format: {"type": "image", "source": {"type": "base64", "media_type": "mime", "data": "DATA"}}
        """
        converted = []
        for block in content:
            if block.get("type") == "image_url":
                url = block.get("image_url", {}).get("url", "")
                if url.startswith("data:") and ";base64," in url:
                    header, data = url.split(";base64,", 1)
                    media_type = header.removeprefix("data:")
                    converted.append({
                        "type": "image",
                        "source": {"type": "base64", "media_type": media_type, "data": data},
                    })
                else:
                    converted.append({
                        "type": "image",
                        "source": {"type": "url", "url": url},
                    })
            else:
                converted.append(block)
        return converted

    def _convert_tools_to_anthropic(
        self,
        tools: list[dict[str, Any]] | None
    ) -> list[dict[str, Any]] | None:
        """Convert OpenAI-format tools to Anthropic format."""
        if not tools:
            return None

        anthropic_tools = []
        for tool in 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": {}})
                })

        return anthropic_tools if anthropic_tools else None

    async def _make_request(
        self,
        messages: list[dict[str, Any]],
        system: str | None = None,
        model: str = "claude-opus-4-5",
        max_tokens: int = 4096,
        temperature: float = 0.7,
        tools: list[dict[str, Any]] | None = None,
        thinking_budget_override: int | None = None,
    ) -> dict[str, Any]:
        """Make request to Anthropic API."""
        client = await self._get_client()

        payload: dict[str, Any] = {
            "model": model,
            "messages": messages,
            "max_tokens": max_tokens,
        }

        # Extended thinking: temperature must be 1 when enabled
        effective_thinking = thinking_budget_override if thinking_budget_override is not None else self.thinking_budget
        if effective_thinking > 0:
            payload["temperature"] = 1
            # max_tokens must exceed budget_tokens
            if max_tokens <= effective_thinking:
                payload["max_tokens"] = effective_thinking + 4096
            payload["thinking"] = {
                "type": "enabled",
                "budget_tokens": effective_thinking,
            }
        else:
            payload["temperature"] = temperature

        if system:
            payload["system"] = system

        if tools:
            payload["tools"] = tools

        logger.info(
            "Anthropic request: model=%s max_tokens=%d thinking=%s",
            payload.get("model"), payload.get("max_tokens"),
            payload.get("thinking", "disabled"),
        )

        response = await client.post(
            self._get_api_url(),
            headers=self._get_headers(),
            json=payload,
        )

        if response.status_code != 200:
            error_text = response.text
            raise Exception(f"Anthropic API error {response.status_code}: {error_text}")

        return response.json()

    async def chat(
        self,
        messages: list[dict[str, Any]],
        tools: list[dict[str, Any]] | None = None,
        model: str | None = None,
        max_tokens: int = 4096,
        temperature: float = 0.7,
        thinking_budget: int | None = None,
    ) -> LLMResponse:
        """Send chat completion request to Anthropic API."""
        model = model or self.default_model

        # Strip provider prefix if present (e.g. "anthropic/claude-opus-4-5" -> "claude-opus-4-5")
        if "/" in model:
            model = model.split("/")[-1]

        # Normalize dots to hyphens (claude-sonnet-4.5 -> claude-sonnet-4-5)
        model = self._normalize_model(model)

        system, prepared_messages = self._prepare_messages(messages)
        anthropic_tools = self._convert_tools_to_anthropic(tools)

        # Per-call thinking override (None = use instance default)
        effective_thinking = self.thinking_budget if thinking_budget is None else thinking_budget

        try:
            response = await self._make_request(
                messages=prepared_messages,
                system=system,
                model=model,
                max_tokens=max_tokens,
                temperature=temperature,
                tools=anthropic_tools,
                thinking_budget_override=effective_thinking,
            )
            return self._parse_response(response)
        except Exception as e:
            return LLMResponse(
                content=f"Error calling LLM: {str(e)}",
                finish_reason="error",
            )

    def _parse_response(self, response: dict[str, Any]) -> LLMResponse:
        """Parse Anthropic API response."""
        content_blocks = response.get("content", [])

        text_content = ""
        thinking_blocks: list[dict[str, Any]] = []
        tool_calls = []

        for block in content_blocks:
            if block.get("type") == "thinking":
                # Preserve full block including signature for multi-turn replay
                thinking_blocks.append(block)
            elif block.get("type") == "text":
                text_content += block.get("text", "")
            elif block.get("type") == "tool_use":
                tool_calls.append(ToolCallRequest(
                    id=block.get("id", ""),
                    name=block.get("name", ""),
                    arguments=block.get("input", {}),
                ))

        usage = {}
        if "usage" in response:
            usage = {
                "prompt_tokens": response["usage"].get("input_tokens", 0),
                "completion_tokens": response["usage"].get("output_tokens", 0),
                "total_tokens": (
                    response["usage"].get("input_tokens", 0) +
                    response["usage"].get("output_tokens", 0)
                ),
            }

        # Log usage and thinking info
        if thinking_blocks:
            thinking_chars = sum(len(b.get("thinking", "")) for b in thinking_blocks)
            logger.info(
                "Anthropic response: %d thinking block(s) (%d chars), "
                "input=%d output=%d tokens",
                len(thinking_blocks), thinking_chars,
                usage.get("prompt_tokens", 0), usage.get("completion_tokens", 0),
            )
        else:
            logger.info(
                "Anthropic response: no thinking blocks, input=%d output=%d tokens",
                usage.get("prompt_tokens", 0), usage.get("completion_tokens", 0),
            )

        return LLMResponse(
            content=text_content or None,
            tool_calls=tool_calls,
            finish_reason=response.get("stop_reason", "end_turn"),
            usage=usage,
            reasoning_content=thinking_blocks or None,
        )

    def get_default_model(self) -> str:
        """Get the default model."""
        return self.default_model

    async def close(self):
        """Close the HTTP client."""
        if self._client:
            await self._client.aclose()
            self._client = None
