"""Anthropic Claude LLM implementation for Fluxibly.

Uses the Anthropic Messages API with support for extended thinking,
tool use, citations, and server-side tools (web search).
"""

from __future__ import annotations

import json
import os
from typing import Any

from fluxibly.exceptions import UnsupportedContentTypeError
from fluxibly.llm.base import BaseLLM, LLMConfig
from fluxibly.schema.response import (
    ContentItem,
    LLMMetadata,
    LLMOutput,
    LLMResponse,
    TokenUsage,
)
from fluxibly.schema.tools import Citation, ToolCall, ToolResult

try:
    import anthropic
except ImportError as e:
    raise ImportError(
        "Anthropic package is required for AnthropicLLM. "
        "Install it with: pip install fluxibly[anthropic]"
    ) from e


class AnthropicLLM(BaseLLM):
    """Anthropic Claude LLM implementation."""

    def __init__(self, config: LLMConfig) -> None:
        super().__init__(config)
        self.client = anthropic.Anthropic(
            api_key=config.api_key or os.environ.get("ANTHROPIC_API_KEY"),
            base_url=config.api_base,
            timeout=config.timeout,
            max_retries=config.max_retries,
        )

    def prepare(self) -> dict[str, Any]:
        """Format for Anthropic Messages API.

        Key differences from OpenAI:
        - System prompt is a top-level "system" param, not a message
        - Messages cannot start with "assistant" role
        - Tool input is "input" not "arguments"
        - Extended thinking uses "thinking" param
        """
        prepared: dict[str, Any] = {
            "model": self.config.model,
            "messages": [],
        }

        # Extract system message as top-level param
        for msg in self.messages:
            if msg["role"] == "system":
                prepared["system"] = msg["content"]
            else:
                prepared["messages"].append(self._convert_message(msg))

        # Max tokens (required for Anthropic)
        prepared["max_tokens"] = self.config.max_output_tokens or 4096

        # Generation params
        if self.config.temperature is not None:
            prepared["temperature"] = self.config.temperature
        if self.config.top_p is not None:
            prepared["top_p"] = self.config.top_p
        if self.config.stop:
            prepared["stop_sequences"] = self.config.stop

        # Extended thinking
        if self.config.reasoning and self.config.reasoning.enabled:
            prepared["thinking"] = {
                "type": "enabled",
                "budget_tokens": self.config.reasoning.budget_tokens or 10000,
            }

        # Tools
        if self.tools:
            prepared["tools"] = self._format_tools()

        # Skills — go into top-level "container" param (Agent Skills standard)
        if self.skills:
            prepared["container"] = {"skills": self._format_skills_anthropic()}

        return prepared

    def forward(self, **kwargs: Any) -> LLMResponse:
        """Execute Anthropic inference."""
        prepared = self.prepare()
        prepared.update(kwargs)

        self.logger.log_input(prepared)

        try:
            raw = self.client.messages.create(**prepared)
            result = self._parse_output(raw)
            self._apply_pricing(result)
            self.logger.log_output(result)
            return result
        except Exception as e:
            self.logger.log_error(e, prepared)
            raise

    # ── Response Parsing ─────────────────────────────────────────────

    def _parse_output(self, raw: Any) -> LLMResponse:
        """Parse Anthropic Messages API output.

        Handles all content block types:
        - "text"              → text with optional citations
        - "thinking"          → extended thinking (reasoning)
        - "redacted_thinking" → redacted thinking (encrypted)
        - "tool_use"          → client-side tool call
        - "server_tool_use"   → server-side tool call (web_search, etc.)
        - "web_search_tool_result" → server-side web search result
        - Unknown types       → preserved with extra dict
        """
        content_items: list[ContentItem] = []
        tool_calls: list[ToolCall] = []
        text_parts: list[str] = []
        reasoning_parts: list[str] = []
        citations: list[Citation] = []

        for block in raw.content or []:
            if block.type == "text":
                text_parts.append(block.text)
                block_citations: list[Citation] | None = None
                if hasattr(block, "citations") and block.citations:
                    block_citations = []
                    for cit in block.citations:
                        c = Citation(
                            type=getattr(cit, "type", "url"),
                            text=getattr(cit, "cited_text", None),
                            url=getattr(cit, "url", None),
                            title=getattr(cit, "title", None),
                        )
                        block_citations.append(c)
                        citations.append(c)
                content_items.append(
                    ContentItem(
                        type="text",
                        text=block.text,
                        annotations=block_citations,
                    )
                )

            elif block.type == "thinking":
                reasoning_parts.append(block.thinking)
                content_items.append(
                    ContentItem(
                        type="reasoning",
                        reasoning_text=block.thinking,
                    )
                )

            elif block.type == "redacted_thinking":
                content_items.append(
                    ContentItem(
                        type="reasoning",
                        extra={"redacted": True, "data": block.data},
                    )
                )

            elif block.type == "tool_use":
                tc = ToolCall(
                    id=block.id,
                    name=block.name,
                    arguments=block.input,
                    type="function",
                )
                tool_calls.append(tc)
                content_items.append(
                    ContentItem(type="tool_call", tool_call=tc)
                )

            elif block.type == "server_tool_use":
                tc = ToolCall(
                    id=block.id,
                    name=block.name,
                    arguments=block.input,
                    type=block.name,
                )
                tool_calls.append(tc)
                content_items.append(
                    ContentItem(type="tool_call", tool_call=tc)
                )

            elif block.type == "web_search_tool_result":
                content_items.append(
                    ContentItem(
                        type="tool_result",
                        tool_result=ToolResult(
                            tool_call_id=block.tool_use_id,
                            content=(
                                block.content
                                if hasattr(block, "content")
                                else str(block)
                            ),
                        ),
                    )
                )

            else:
                content_items.append(
                    ContentItem(
                        type=block.type,
                        extra=(
                            block.model_dump()
                            if hasattr(block, "model_dump")
                            else {"raw": str(block)}
                        ),
                    )
                )

        # Normalize stop reason
        stop_reason_map = {
            "end_turn": "stop",
            "tool_use": "tool_calls",
            "max_tokens": "max_tokens",
            "stop_sequence": "stop",
        }

        return LLMResponse(
            output=LLMOutput(
                output_text="\n".join(text_parts),
                content=content_items,
                tool_calls=tool_calls,
                reasoning=(
                    "\n".join(reasoning_parts) if reasoning_parts else None
                ),
                citations=citations,
            ),
            metadata=LLMMetadata(
                model=raw.model,
                id=raw.id,
                created_at="",
                status=(
                    "completed"
                    if raw.stop_reason in ("end_turn", "stop_sequence")
                    else "incomplete"
                ),
                usage=TokenUsage(
                    input_tokens=raw.usage.input_tokens,
                    output_tokens=raw.usage.output_tokens,
                    total_tokens=raw.usage.input_tokens
                    + raw.usage.output_tokens,
                    cached_tokens=getattr(
                        raw.usage, "cache_read_input_tokens", None
                    ),
                ),
                stop_reason=stop_reason_map.get(
                    raw.stop_reason, raw.stop_reason
                ),
                provider="anthropic",
                raw_response=(
                    raw.model_dump() if hasattr(raw, "model_dump") else None
                ),
            ),
        )

    # ── Message Conversion ───────────────────────────────────────────

    def _convert_message(
        self, msg: dict[str, Any]
    ) -> dict[str, Any] | list[dict[str, Any]]:
        """Convert our unified message format to Anthropic format.

        Returns a single dict for most messages, or a list of dicts when
        an assistant tool_calls message needs to be followed by tool_result messages.
        """
        role = msg["role"]
        content = msg.get("content", "")

        # Assistant message with tool calls → Anthropic tool_use content blocks
        if role == "assistant" and msg.get("tool_calls"):
            blocks: list[dict[str, Any]] = []
            if content:
                blocks.append({"type": "text", "text": str(content)})
            for tc in msg["tool_calls"]:
                args = tc.get("arguments", tc.get("input", {}))
                if isinstance(args, str):
                    try:
                        args = json.loads(args)
                    except json.JSONDecodeError:
                        args = {}
                blocks.append(
                    {
                        "type": "tool_use",
                        "id": tc["id"],
                        "name": tc["name"],
                        "input": args,
                    }
                )
            return {"role": "assistant", "content": blocks}

        # Tool result → Anthropic tool_result content block (role=user)
        if role == "tool":
            return {
                "role": "user",
                "content": [
                    {
                        "type": "tool_result",
                        "tool_use_id": msg["tool_call_id"],
                        "content": str(content),
                    }
                ],
            }

        if isinstance(content, str):
            return {"role": role, "content": content}

        # Multimodal content
        converted: list[dict[str, Any]] = []
        for part in content:
            ptype = part.get("type", "text")
            if ptype == "text":
                converted.append({"type": "text", "text": part["text"]})
            elif ptype == "image":
                source = part.get("source", {})
                converted.append(
                    {
                        "type": "image",
                        "source": {
                            "type": source.get("type", "base64"),
                            "media_type": source.get(
                                "media_type", "image/png"
                            ),
                            **(
                                {"data": source["data"]}
                                if source.get("type") == "base64"
                                else {}
                            ),
                            **(
                                {"url": source["url"]}
                                if source.get("type") == "url"
                                else {}
                            ),
                        },
                    }
                )
            elif ptype == "document":
                source = part.get("source", {})
                converted.append(
                    {
                        "type": "document",
                        "source": {
                            "type": source.get("type", "base64"),
                            "media_type": source.get(
                                "media_type", "application/pdf"
                            ),
                            **(
                                {"data": source["data"]}
                                if source.get("type") == "base64"
                                else {}
                            ),
                            **(
                                {"url": source["url"]}
                                if source.get("type") == "url"
                                else {}
                            ),
                        },
                    }
                )
            elif ptype in ("audio", "video"):
                raise UnsupportedContentTypeError(
                    f"Anthropic does not support {ptype} input"
                )
            else:
                raise UnsupportedContentTypeError(
                    f"Anthropic does not support content type: {ptype}"
                )

        return {"role": role, "content": converted}

    # ── Tool Formatting ──────────────────────────────────────────────

    def _format_tools(self) -> list[dict[str, Any]]:
        """Format tools for Anthropic API."""
        formatted: list[dict[str, Any]] = []
        for tool in self.tools:
            ttype = tool.get("type", "function")
            if ttype == "function":
                func = tool.get("function", {})
                formatted.append(
                    {
                        "name": func.get("name", ""),
                        "description": func.get("description", ""),
                        "input_schema": func.get("parameters", {}),
                    }
                )
            elif ttype == "web_search":
                formatted.append(
                    {
                        "type": "web_search_20250305",
                        "name": "web_search",
                        "max_uses": tool.get("web_search", {}).get(
                            "max_uses", 5
                        ),
                    }
                )
        return formatted

    # ── Skill Formatting ──────────────────────────────────────────────

    def _format_skills_anthropic(self) -> list[dict[str, Any]]:
        """Format Agent Skills for Anthropic container.skills parameter.

        Anthropic supports two skill types:
        - "anthropic": Pre-built skills (pptx, xlsx, docx, pdf, etc.)
        - "custom": User-provided skill bundles

        Each skill entry: {type, skill_id, version?}
        """
        formatted: list[dict[str, Any]] = []
        for skill in self.skills:
            source = skill.get("source", "custom")
            skill_type = "anthropic" if source == "anthropic" else "custom"
            entry: dict[str, Any] = {
                "type": skill_type,
                "skill_id": skill["skill_id"],
            }
            if skill.get("version"):
                entry["version"] = skill["version"]
            formatted.append(entry)
        return formatted
