"""
Anthropic adapter — handles Claude models via the native Anthropic SDK.

Converts between OpenAI message format (used internally by PantheonOS)
and Anthropic's native format, and normalizes streaming events.
"""

import json
import time
from typing import Any, Callable

from ..log import logger
from ..misc import run_func
from .base import (
    BaseAdapter,
    ServiceUnavailableError,
    InternalServerError,
    RateLimitError,
    APIConnectionError,
)
from .image_blocks import has_image_content, split_text_and_images


def _content_to_anthropic_tool_result(content: Any) -> Any:
    """Translate a tool message's content into Anthropic tool_result content.

    - Plain string → returned unchanged.
    - Content list with image_url blocks → list of {type:"text"} + {type:"image"} blocks.
    - Content list without images → joined text string.
    """
    if isinstance(content, str):
        return content
    if not has_image_content(content):
        # No images — flatten list of text blocks into a single string for Anthropic.
        text, _inline, _http = split_text_and_images(content)
        return text or ""

    text, inline_images, http_urls = split_text_and_images(content)
    blocks: list[dict] = []
    if text:
        blocks.append({"type": "text", "text": text})
    for mime, data in inline_images:
        blocks.append({
            "type": "image",
            "source": {
                "type": "base64",
                "media_type": mime,
                "data": data,
            },
        })
    for url in http_urls:
        # Anthropic Messages API supports URL sources.
        blocks.append({
            "type": "image",
            "source": {"type": "url", "url": url},
        })
    return blocks if blocks else ""


def _normalize_user_content(content: Any) -> Any:
    """Translate a user message's OpenAI-style content into Anthropic-native blocks.

    Anthropic rejects raw ``image_url`` blocks in user content (400: "Input tag
    'image_url' ... does not match any of the expected tags"), so convert them to
    ``{type: 'image', source: ...}``. Plain strings and image-free content pass
    through unchanged. Mirrors ``_content_to_anthropic_tool_result`` but keeps a
    block LIST (user turns may legitimately mix text + images).
    """
    if isinstance(content, str):
        return content
    if not isinstance(content, list):
        return content
    if not has_image_content(content):
        return content

    text, inline_images, http_urls = split_text_and_images(content)
    blocks: list[dict] = []
    if text:
        blocks.append({"type": "text", "text": text})
    for mime, data in inline_images:
        blocks.append({
            "type": "image",
            "source": {"type": "base64", "media_type": mime, "data": data},
        })
    for url in http_urls:
        blocks.append({"type": "image", "source": {"type": "url", "url": url}})
    return blocks if blocks else content


def _wrap_anthropic_error(e: Exception) -> Exception:
    """Convert anthropic SDK exceptions to unified exception types."""
    try:
        import anthropic as anthropic_mod

        if isinstance(e, anthropic_mod.RateLimitError):
            return RateLimitError(str(e))
        elif isinstance(e, anthropic_mod.APIConnectionError):
            return APIConnectionError(str(e))
        elif isinstance(e, anthropic_mod.InternalServerError):
            return InternalServerError(str(e))
        elif isinstance(e, anthropic_mod.APIStatusError):
            status = getattr(e, "status_code", 0)
            if status == 503:
                return ServiceUnavailableError(str(e))
            elif status == 429:
                return RateLimitError(str(e))
            elif status >= 500:
                return InternalServerError(str(e))
    except ImportError:
        pass
    return e


# ============ Message Format Conversion ============


def _msg_has_content(content: Any) -> bool:
    """Whether a message carries any content. Anthropic 400s on an empty user
    message ('user messages must have non-empty content'), which a bare approval
    / empty user turn (content == [] or '') would otherwise produce."""
    if content is None:
        return False
    if isinstance(content, str):
        return content.strip() != ""
    if isinstance(content, list):
        return len(content) > 0
    return bool(content)


def _convert_messages_to_anthropic(messages: list[dict]) -> tuple[str | list | None, list[dict]]:
    """Convert OpenAI-format messages to Anthropic format.

    Key differences:
    - System messages become top-level `system` parameter
    - tool_calls in assistant messages become tool_use content blocks
    - tool role messages become tool_result content blocks in user messages

    Returns:
        (system_prompt, converted_messages)
    """
    system_prompt = None
    converted = []
    pending_tool_results = []

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

        if role == "system":
            # First system message becomes the system parameter
            if system_prompt is None:
                # Anthropic's `system` accepts a STRING or a LIST of text blocks.
                # Preserve a block list as-is so the cache_control ephemeral
                # marker on its last block survives — str(content) dropped the
                # marker (→ 0% prompt cache) and turned the prompt into a Python
                # repr blob. A plain string still passes through unchanged.
                if isinstance(content, (str, list)):
                    system_prompt = content
                else:
                    system_prompt = str(content)
            else:
                # Additional system messages become user messages
                converted.append({
                    "role": "user",
                    "content": f"[System]: {content}"
                })
            continue

        if role == "tool":
            # Accumulate tool results to attach to next user message.
            # Image content (via OpenAI-style image_url blocks) is translated
            # into Anthropic image blocks so the agent can see the picture
            # directly in tool_result.
            pending_tool_results.append({
                "type": "tool_result",
                "tool_use_id": msg.get("tool_call_id", ""),
                "content": _content_to_anthropic_tool_result(content or ""),
            })
            continue

        if role == "user":
            # Flush pending tool results first
            if pending_tool_results:
                # Tool results must be in a user message
                result_content = list(pending_tool_results)
                if content:
                    if isinstance(content, str):
                        result_content.append({"type": "text", "text": content})
                    elif isinstance(content, list):
                        normalized = _normalize_user_content(content)
                        if isinstance(normalized, list):
                            result_content.extend(normalized)
                        else:
                            result_content.append({"type": "text", "text": str(normalized)})
                converted.append({"role": "user", "content": result_content})
                pending_tool_results = []
            elif _msg_has_content(content):
                # Skip a bare/empty user turn (e.g. an approval action with no
                # text) — an empty user message is rejected by the API.
                converted.append({"role": "user", "content": _normalize_user_content(content)})
            continue

        if role == "assistant":
            # Flush any pending tool results as a separate user message
            if pending_tool_results:
                converted.append({"role": "user", "content": list(pending_tool_results)})
                pending_tool_results = []

            # Build content blocks for assistant
            content_blocks = []

            # Text content
            if content:
                if isinstance(content, str):
                    content_blocks.append({"type": "text", "text": content})
                elif isinstance(content, list):
                    content_blocks.extend(content)

            # Tool calls → tool_use blocks
            tool_calls = msg.get("tool_calls")
            if tool_calls:
                for tc in tool_calls:
                    func = tc.get("function", {})
                    # Parse arguments from JSON string
                    try:
                        input_data = json.loads(func.get("arguments", "{}"))
                    except (json.JSONDecodeError, TypeError):
                        input_data = {}

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

            if content_blocks:
                converted.append({"role": "assistant", "content": content_blocks})
            elif not content and not tool_calls:
                # Empty assistant message — skip
                pass
            continue

    # Flush remaining tool results
    if pending_tool_results:
        converted.append({"role": "user", "content": list(pending_tool_results)})

    # Belt-and-suspenders: drop any message that still has empty content (a bare
    # approval / empty user turn) before the alternation-merge, so the merge sees
    # only real content and the API never gets an empty user message.
    converted = [m for m in converted if _msg_has_content(m.get("content"))]

    # Anthropic requires alternating user/assistant messages
    # Merge consecutive same-role messages
    merged = []
    for msg in converted:
        if merged and merged[-1]["role"] == msg["role"]:
            prev = merged[-1]["content"]
            curr = msg["content"]
            # Normalize both to lists
            if isinstance(prev, str):
                prev = [{"type": "text", "text": prev}]
            elif not isinstance(prev, list):
                prev = [{"type": "text", "text": str(prev)}]
            if isinstance(curr, str):
                curr = [{"type": "text", "text": curr}]
            elif not isinstance(curr, list):
                curr = [{"type": "text", "text": str(curr)}]
            merged[-1]["content"] = prev + curr
        else:
            merged.append(msg)

    return system_prompt, merged


def _convert_tools_to_anthropic(tools: list[dict] | None) -> list[dict] | None:
    """Convert OpenAI tool format to Anthropic tool format.

    From: {"type": "function", "function": {"name": ..., "description": ..., "parameters": {...}}}
    To:   {"name": ..., "description": ..., "input_schema": {...}}
    """
    if not tools:
        return None

    converted = []
    for tool in tools:
        func = tool.get("function", {})
        anthropic_tool = {
            "name": func.get("name", ""),
        }
        if "description" in func:
            anthropic_tool["description"] = func["description"]
        if "parameters" in func:
            anthropic_tool["input_schema"] = func["parameters"]
        else:
            anthropic_tool["input_schema"] = {"type": "object", "properties": {}}
        converted.append(anthropic_tool)

    return converted


# ============ Adapter ============


class AnthropicAdapter(BaseAdapter):
    """Adapter for Anthropic Claude API."""

    def _make_client(
        self,
        base_url: str | None = None,
        api_key: str | None = None,
    ):
        from anthropic import AsyncAnthropic
        import httpx

        # Bound the request so a stalled endpoint fails fast (~120s) and lets
        # pantheon's own num_retries loop retry, instead of hanging on the SDK
        # default (~600s). max_retries=0: don't also retry inside the SDK.
        kwargs = {
            "timeout": httpx.Timeout(120.0, connect=10.0),
            "max_retries": 0,
        }
        if base_url:
            kwargs["base_url"] = base_url
        if api_key:
            kwargs["api_key"] = api_key
        return AsyncAnthropic(**kwargs)

    async def acompletion(
        self,
        *,
        model: str,
        messages: list[dict],
        tools: list[dict] | None = None,
        response_format: Any | None = None,
        stream: bool = True,
        process_chunk: Callable | None = None,
        base_url: str | None = None,
        api_key: str | None = None,
        num_retries: int = 3,
        **kwargs,
    ):
        """Streaming chat completion using the Anthropic SDK.

        Converts OpenAI messages to Anthropic format, streams events,
        normalizes them to OpenAI-compatible chunk dicts, and returns
        collected chunks.
        """
        client = self._make_client(base_url, api_key)

        # Convert messages and tools
        system_prompt, anthropic_messages = _convert_messages_to_anthropic(messages)
        anthropic_tools = _convert_tools_to_anthropic(tools)

        # Build call kwargs (stream() method implies streaming, don't pass stream=True)
        call_kwargs = {
            "model": model,
            "messages": anthropic_messages,
            "max_tokens": kwargs.pop("max_tokens", None) or kwargs.pop("max_output_tokens", 8192),
        }

        if system_prompt:
            call_kwargs["system"] = system_prompt

        if anthropic_tools:
            call_kwargs["tools"] = anthropic_tools

        # Handle thinking parameter
        thinking = kwargs.pop("thinking", None)
        reasoning_effort = kwargs.pop("reasoning_effort", None)
        if thinking:
            call_kwargs["thinking"] = thinking
        elif reasoning_effort:
            # Map reasoning_effort to Anthropic thinking with effort-based budget
            effort_budgets = {"low": 5000, "medium": 10000, "high": 30000}
            budget = effort_budgets.get(reasoning_effort, 10000)
            call_kwargs["thinking"] = {"type": "enabled", "budget_tokens": budget}

        # Temperature
        temperature = kwargs.pop("temperature", None)
        if temperature is not None:
            call_kwargs["temperature"] = temperature

        # Top-p
        top_p = kwargs.pop("top_p", None)
        if top_p is not None:
            call_kwargs["top_p"] = top_p

        # Extra headers
        extra_headers = kwargs.pop("extra_headers", None)

        try:
            stream_start_time = time.time()
            first_chunk_time = None
            chunk_count = 0
            collected_chunks = []

            # Track state for building OpenAI-compatible chunks
            current_text = ""
            current_tool_calls = []
            tool_call_index = -1
            tool_call_json_accum = ""
            usage_info = {}

            async with client.messages.stream(
                **call_kwargs,
                extra_headers=extra_headers or {},
            ) as stream_resp:
                async for event in stream_resp:
                    event_type = event.type

                    if event_type == "message_start":
                        # Extract initial usage. Anthropic splits input into three
                        # NON-overlapping buckets: input_tokens (freshly prefilled),
                        # cache_creation_input_tokens (written to the prompt cache
                        # this call) and cache_read_input_tokens (served from cache).
                        # prompt_tokens must be their SUM to match the OpenAI path
                        # (where prompt_tokens already includes the cached portion) —
                        # otherwise cost and the [cache] hit-rate are computed against
                        # just the handful of fresh tokens (e.g. prompt_tokens=2 for a
                        # fully-cached 57k prompt). The two cache_* fields flow through
                        # stream_chunk_builder into usage so _extract_cost_and_usage
                        # can surface prompt-cache reuse.
                        msg = getattr(event, "message", None)
                        if msg and hasattr(msg, "usage"):
                            u = msg.usage
                            fresh = getattr(u, "input_tokens", 0) or 0
                            cache_write = getattr(u, "cache_creation_input_tokens", 0) or 0
                            cache_read = getattr(u, "cache_read_input_tokens", 0) or 0
                            usage_info["prompt_tokens"] = fresh + cache_write + cache_read
                            if cache_write:
                                usage_info["cache_creation_input_tokens"] = cache_write
                            if cache_read:
                                usage_info["cache_read_input_tokens"] = cache_read

                    elif event_type == "content_block_start":
                        block = event.content_block
                        if block.type == "tool_use":
                            tool_call_index += 1
                            tool_call_json_accum = ""
                            current_tool_calls.append({
                                "index": tool_call_index,
                                "id": block.id,
                                "type": "function",
                                "function": {
                                    "name": block.name,
                                    "arguments": "",
                                },
                            })
                            # Emit initial chunk with id and name
                            chunk_dict = {
                                "choices": [{
                                    "index": 0,
                                    "delta": {
                                        "role": "assistant",
                                        "tool_calls": [{
                                            "index": tool_call_index,
                                            "id": block.id,
                                            "type": "function",
                                            "function": {
                                                "name": block.name,
                                                "arguments": "",
                                            },
                                        }],
                                    },
                                    "finish_reason": None,
                                }],
                            }
                            collected_chunks.append(chunk_dict)

                    elif event_type == "content_block_delta":
                        delta_obj = event.delta

                        if delta_obj.type == "text_delta":
                            text = delta_obj.text
                            current_text += text

                            if first_chunk_time is None:
                                first_chunk_time = time.time()
                                ttfb = first_chunk_time - stream_start_time
                                logger.info(f"⚡ First chunk received: {ttfb:.3f}s (TTFB) [{model}]")

                            # Build OpenAI-compatible chunk
                            chunk_dict = {
                                "choices": [{
                                    "index": 0,
                                    "delta": {
                                        "role": "assistant",
                                        "content": text,
                                    },
                                    "finish_reason": None,
                                }],
                            }
                            collected_chunks.append(chunk_dict)

                            if process_chunk:
                                chunk_count += 1
                                await run_func(process_chunk, {
                                    "role": "assistant",
                                    "content": text,
                                })

                        elif delta_obj.type == "input_json_delta":
                            # Accumulate tool call arguments
                            partial = delta_obj.partial_json
                            tool_call_json_accum += partial
                            if current_tool_calls:
                                current_tool_calls[-1]["function"]["arguments"] += partial

                            chunk_dict = {
                                "choices": [{
                                    "index": 0,
                                    "delta": {
                                        "role": "assistant",
                                        "tool_calls": [{
                                            "index": tool_call_index,
                                            "function": {
                                                "arguments": partial,
                                            },
                                        }],
                                    },
                                    "finish_reason": None,
                                }],
                            }
                            collected_chunks.append(chunk_dict)

                        elif delta_obj.type == "thinking_delta":
                            thinking_text = delta_obj.thinking

                            # Write into chunks so stream_chunk_builder captures it
                            chunk_dict = {
                                "choices": [{
                                    "index": 0,
                                    "delta": {
                                        "role": "assistant",
                                        "reasoning_content": thinking_text,
                                    },
                                    "finish_reason": None,
                                }],
                            }
                            collected_chunks.append(chunk_dict)

                            if process_chunk:
                                await run_func(process_chunk, {
                                    "role": "assistant",
                                    "reasoning_content": thinking_text,
                                })

                    elif event_type == "message_delta":
                        delta = event.delta
                        stop_reason = getattr(delta, "stop_reason", None)

                        # Extract usage from message_delta
                        usage = getattr(event, "usage", None)
                        if usage:
                            usage_info["completion_tokens"] = getattr(usage, "output_tokens", 0)

                        # Map Anthropic stop reasons to OpenAI finish reasons
                        finish_reason = None
                        if stop_reason == "end_turn":
                            finish_reason = "stop"
                        elif stop_reason == "tool_use":
                            finish_reason = "tool_calls"
                        elif stop_reason == "max_tokens":
                            finish_reason = "length"

                        if finish_reason:
                            chunk_dict = {
                                "choices": [{
                                    "index": 0,
                                    "delta": {},
                                    "finish_reason": finish_reason,
                                }],
                            }
                            collected_chunks.append(chunk_dict)

                            if process_chunk and finish_reason == "stop":
                                await run_func(process_chunk, {"stop": True})

                    elif event_type == "message_stop":
                        pass

            # Add usage chunk at the end (OpenAI stream_options style)
            total_tokens = usage_info.get("prompt_tokens", 0) + usage_info.get("completion_tokens", 0)
            usage_info["total_tokens"] = total_tokens
            collected_chunks.append({
                "usage": usage_info,
                "choices": [],
            })

            total_time = time.time() - stream_start_time
            logger.info(f"✅ Stream completed: {total_time:.3f}s, {chunk_count} chunks [{model}]")

            return collected_chunks

        except Exception as e:
            raise _wrap_anthropic_error(e) from e
