"""Anthropic (Claude) provider implementation."""

from __future__ import annotations

import asyncio
import time
from typing import Any

from pi_agents.ai.env_api_keys import get_env_api_key
from pi_agents.ai.models import calculate_cost
from pi_agents.ai.types import (
    AssistantMessage,
    Context,
    Cost,
    ImageContent,
    Model,
    SimpleStreamOptions,
    StreamOptions,
    TextContent,
    ThinkingContent,
    Tool,
    ToolCall,
    ToolResultMessage,
    Usage,
    UserMessage,
)
from pi_agents.ai.utils.event_stream import AssistantMessageEventStream
from pi_agents.ai.utils.json_parse import parse_streaming_json


DEFAULT_THINKING_BUDGETS = {"minimal": 1024, "low": 2048, "medium": 8192, "high": 16384}

# Models that support adaptive thinking (effort-based)
ADAPTIVE_MODELS = {"claude-opus-4-6", "claude-sonnet-4-6"}


def _normalize_tool_call_id(id: str) -> str:
    """Normalize tool call ID to match Anthropic's format."""
    import re
    return re.sub(r"[^a-zA-Z0-9_-]", "_", id)[:64]


def _convert_messages(messages: list[Any], tools: list[Tool] | None = None) -> list[dict[str, Any]]:
    """Convert pi-agents messages to Anthropic API format."""
    result: list[dict[str, Any]] = []

    for msg in messages:
        if isinstance(msg, UserMessage) or (isinstance(msg, dict) and msg.get("role") == "user"):
            content = msg.content if isinstance(msg, UserMessage) else msg.get("content", "")
            if isinstance(content, str):
                result.append({"role": "user", "content": content})
            else:
                blocks = []
                for block in content:
                    if isinstance(block, TextContent) or (isinstance(block, dict) and block.get("type") == "text"):
                        text = block.text if isinstance(block, TextContent) else block.get("text", "")
                        blocks.append({"type": "text", "text": text})
                    elif isinstance(block, ImageContent) or (isinstance(block, dict) and block.get("type") == "image"):
                        data = block.data if isinstance(block, ImageContent) else block.get("data", "")
                        mime = block.mime_type if isinstance(block, ImageContent) else block.get("mimeType", "")
                        blocks.append({
                            "type": "image",
                            "source": {"type": "base64", "media_type": mime, "data": data},
                        })
                result.append({"role": "user", "content": blocks})

        elif isinstance(msg, AssistantMessage) or (isinstance(msg, dict) and msg.get("role") == "assistant"):
            content_list = msg.content if isinstance(msg, AssistantMessage) else msg.get("content", [])
            blocks = []
            for block in content_list:
                if isinstance(block, TextContent) or (isinstance(block, dict) and block.get("type") == "text"):
                    text = block.text if isinstance(block, TextContent) else block.get("text", "")
                    sig = block.text_signature if isinstance(block, TextContent) else block.get("textSignature")
                    blocks.append({"type": "text", "text": text})
                elif isinstance(block, ThinkingContent) or (isinstance(block, dict) and block.get("type") == "thinking"):
                    thinking = block.thinking if isinstance(block, ThinkingContent) else block.get("thinking", "")
                    sig = block.thinking_signature if isinstance(block, ThinkingContent) else block.get("thinkingSignature")
                    redacted = block.redacted if isinstance(block, ThinkingContent) else block.get("redacted")
                    if redacted and sig:
                        blocks.append({"type": "thinking", "thinking": "", "signature": sig})
                    else:
                        entry: dict[str, Any] = {"type": "thinking", "thinking": thinking}
                        if sig:
                            entry["signature"] = sig
                        blocks.append(entry)
                elif isinstance(block, ToolCall) or (isinstance(block, dict) and block.get("type") == "toolCall"):
                    tc_id = block.id if isinstance(block, ToolCall) else block.get("id", "")
                    tc_name = block.name if isinstance(block, ToolCall) else block.get("name", "")
                    tc_args = block.arguments if isinstance(block, ToolCall) else block.get("arguments", {})
                    blocks.append({
                        "type": "tool_use",
                        "id": _normalize_tool_call_id(tc_id),
                        "name": tc_name,
                        "input": tc_args,
                    })
            result.append({"role": "assistant", "content": blocks})

        elif isinstance(msg, ToolResultMessage) or (isinstance(msg, dict) and msg.get("role") == "toolResult"):
            tc_id = msg.tool_call_id if isinstance(msg, ToolResultMessage) else msg.get("toolCallId", "")
            content_list = msg.content if isinstance(msg, ToolResultMessage) else msg.get("content", [])
            is_error = msg.is_error if isinstance(msg, ToolResultMessage) else msg.get("isError", False)

            blocks = []
            for block in content_list:
                if isinstance(block, TextContent) or (isinstance(block, dict) and block.get("type") == "text"):
                    text = block.text if isinstance(block, TextContent) else block.get("text", "")
                    blocks.append({"type": "text", "text": text})
                elif isinstance(block, ImageContent) or (isinstance(block, dict) and block.get("type") == "image"):
                    data = block.data if isinstance(block, ImageContent) else block.get("data", "")
                    mime = block.mime_type if isinstance(block, ImageContent) else block.get("mimeType", "")
                    blocks.append({
                        "type": "image",
                        "source": {"type": "base64", "media_type": mime, "data": data},
                    })

            content = blocks if blocks else [{"type": "text", "text": ""}]
            result.append({
                "role": "user",
                "content": [{
                    "type": "tool_result",
                    "tool_use_id": _normalize_tool_call_id(tc_id),
                    "content": content,
                    **({"is_error": True} if is_error else {}),
                }],
            })

    return result


def _convert_tools(tools: list[Tool] | None) -> list[dict[str, Any]] | None:
    """Convert tools to Anthropic format."""
    if not tools:
        return None
    return [
        {
            "name": t.name,
            "description": t.description,
            "input_schema": t.parameters,
        }
        for t in tools
    ]


def stream_anthropic(
    model: Model,
    context: Context,
    options: dict[str, Any] | None = None,
) -> AssistantMessageEventStream:
    """Stream responses from the Anthropic API."""
    event_stream = AssistantMessageEventStream()
    opts = options or {}

    async def _run():
        try:
            import anthropic

            api_key = opts.get("api_key") or opts.get("apiKey") or get_env_api_key("anthropic")
            if not api_key:
                raise ValueError("Anthropic API key not found. Set ANTHROPIC_API_KEY environment variable.")

            client = anthropic.AsyncAnthropic(
                api_key=api_key,
                base_url=model.base_url or None,
            )

            messages = _convert_messages(context.messages, context.tools)
            tools = _convert_tools(context.tools)

            params: dict[str, Any] = {
                "model": model.id,
                "messages": messages,
                "max_tokens": opts.get("max_tokens") or opts.get("maxTokens") or model.max_tokens or 4096,
                "stream": True,
            }

            if context.system_prompt:
                params["system"] = context.system_prompt

            if tools:
                params["tools"] = tools

            if opts.get("temperature") is not None:
                params["temperature"] = opts["temperature"]

            # Thinking/reasoning support
            thinking_enabled = opts.get("thinking_enabled") or opts.get("thinkingEnabled", False)
            thinking_budget = opts.get("thinking_budget_tokens") or opts.get("thinkingBudgetTokens")

            if thinking_enabled and thinking_budget:
                # Check if model supports adaptive thinking
                is_adaptive = any(m in model.id for m in ADAPTIVE_MODELS)
                if is_adaptive:
                    effort = opts.get("effort", "medium")
                    params["thinking"] = {"type": "adaptive", "output_config": {"effort": effort}}
                else:
                    params["thinking"] = {"type": "enabled", "budget_tokens": thinking_budget}

            # Build partial assistant message
            partial = AssistantMessage(
                role="assistant",
                content=[],
                api=model.api,
                provider=model.provider,
                model=model.id,
                usage=Usage(cost=Cost()),
                stop_reason="stop",
                timestamp=int(time.time() * 1000),
            )

            event_stream.push({"type": "start", "partial": partial})

            content_index = 0
            async with client.messages.stream(**params) as stream:
                async for event in stream:
                    event_type = getattr(event, "type", None)

                    if event_type == "content_block_start":
                        block = event.content_block
                        if block.type == "text":
                            event_stream.push({"type": "text_start", "contentIndex": content_index, "partial": partial})
                        elif block.type == "thinking":
                            event_stream.push({"type": "thinking_start", "contentIndex": content_index, "partial": partial})
                        elif block.type == "tool_use":
                            event_stream.push({"type": "toolcall_start", "contentIndex": content_index, "partial": partial})

                    elif event_type == "content_block_delta":
                        delta = event.delta
                        if hasattr(delta, "text"):
                            event_stream.push({
                                "type": "text_delta",
                                "contentIndex": content_index,
                                "delta": delta.text,
                                "partial": partial,
                            })
                        elif hasattr(delta, "thinking"):
                            event_stream.push({
                                "type": "thinking_delta",
                                "contentIndex": content_index,
                                "delta": delta.thinking,
                                "partial": partial,
                            })
                        elif hasattr(delta, "partial_json"):
                            event_stream.push({
                                "type": "toolcall_delta",
                                "contentIndex": content_index,
                                "delta": delta.partial_json,
                                "partial": partial,
                            })

                    elif event_type == "content_block_stop":
                        content_index += 1

                    elif event_type == "message_delta":
                        if hasattr(event, "usage"):
                            u = event.usage
                            partial.usage.output = getattr(u, "output_tokens", 0)

                    elif event_type == "message_start":
                        if hasattr(event, "message") and hasattr(event.message, "usage"):
                            u = event.message.usage
                            partial.usage.input = getattr(u, "input_tokens", 0)
                            partial.usage.cache_read = getattr(u, "cache_read_input_tokens", 0) or 0
                            partial.usage.cache_write = getattr(u, "cache_creation_input_tokens", 0) or 0

                # Final message
                final_message = stream.get_final_message()
                stop_reason_map = {
                    "end_turn": "stop",
                    "max_tokens": "length",
                    "tool_use": "toolUse",
                    "stop_sequence": "stop",
                }
                stop_reason = stop_reason_map.get(final_message.stop_reason, "stop")

                # Build final content
                final_content: list[Any] = []
                for block in final_message.content:
                    if block.type == "text":
                        final_content.append(TextContent(type="text", text=block.text))
                    elif block.type == "thinking":
                        final_content.append(ThinkingContent(
                            type="thinking",
                            thinking=block.thinking,
                            thinking_signature=getattr(block, "signature", None),
                        ))
                    elif block.type == "tool_use":
                        final_content.append(ToolCall(
                            type="toolCall",
                            id=block.id,
                            name=block.name,
                            arguments=dict(block.input) if block.input else {},
                        ))

                partial.content = final_content
                partial.stop_reason = stop_reason
                partial.usage.total_tokens = partial.usage.input + partial.usage.output + partial.usage.cache_read + partial.usage.cache_write

                calculate_cost(model, partial.usage)

                if stop_reason in ("stop", "length", "toolUse"):
                    event_stream.push({"type": "done", "reason": stop_reason, "message": partial})
                else:
                    event_stream.push({"type": "error", "reason": "error", "error": partial})

        except Exception as e:
            error_msg = AssistantMessage(
                role="assistant",
                content=[],
                api=model.api,
                provider=model.provider,
                model=model.id,
                usage=Usage(cost=Cost()),
                stop_reason="error",
                error_message=str(e),
                timestamp=int(time.time() * 1000),
            )
            event_stream.push({"type": "error", "reason": "error", "error": error_msg})

        event_stream.end()

    asyncio.ensure_future(_run())
    return event_stream


def stream_simple_anthropic(
    model: Model,
    context: Context,
    options: SimpleStreamOptions | None = None,
) -> AssistantMessageEventStream:
    """Stream responses using simplified options interface."""
    opts: dict[str, Any] = {}
    if options:
        if options.temperature is not None:
            opts["temperature"] = options.temperature
        if options.max_tokens is not None:
            opts["maxTokens"] = options.max_tokens
        if options.api_key:
            opts["apiKey"] = options.api_key

        # Convert reasoning level to thinking config
        if options.reasoning:
            level = options.reasoning
            budgets = DEFAULT_THINKING_BUDGETS.copy()
            if options.thinking_budgets:
                if options.thinking_budgets.minimal is not None:
                    budgets["minimal"] = options.thinking_budgets.minimal
                if options.thinking_budgets.low is not None:
                    budgets["low"] = options.thinking_budgets.low
                if options.thinking_budgets.medium is not None:
                    budgets["medium"] = options.thinking_budgets.medium
                if options.thinking_budgets.high is not None:
                    budgets["high"] = options.thinking_budgets.high

            opts["thinkingEnabled"] = True
            opts["thinkingBudgetTokens"] = budgets.get(level, 8192)

            # Map level to effort for adaptive models
            effort_map = {"minimal": "low", "low": "low", "medium": "medium", "high": "high", "xhigh": "high"}
            opts["effort"] = effort_map.get(level, "medium")

    return stream_anthropic(model, context, opts)
