"""Anthropic Messages API bridge — translates OpenAI format ↔ Anthropic format.

Handles:
- Message format conversion (OpenAI ↔ Anthropic)
- Tool calling translation
- Streaming with SSE format conversion
- Prompt caching token tracking (cache_read_input_tokens, cache_creation_input_tokens)
"""

from __future__ import annotations

import json
import logging
from dataclasses import dataclass
from typing import AsyncIterator

import aiohttp

from backend.bridge.base import BaseBridge, BridgeRequest, BridgeResponse, StreamResult, StreamUsage
from backend.config import ModelConfig

logger = logging.getLogger("ecate.bridge.anthropic")


def _convert_content_to_anthropic(content) -> list[dict] | str:
    """Convert OpenAI content format to Anthropic format.
    
    Handles:
    - String content (passthrough)
    - Array content with text/image parts
    """
    if isinstance(content, str):
        return content
    
    if not isinstance(content, list):
        return str(content)
    
    result = []
    for part in content:
        if not isinstance(part, dict):
            continue
        part_type = part.get("type", "")
        
        if part_type == "text":
            result.append({"type": "text", "text": part.get("text", "")})
        
        elif part_type == "image_url":
            # OpenAI: {type: "image_url", image_url: {url: "data:image/png;base64,..."}}
            # Anthropic: {type: "image", source: {type: "base64", media_type: "...", data: "..."}}
            image_url = part.get("image_url", {})
            url = image_url.get("url", "") if isinstance(image_url, dict) else str(image_url)
            
            if url.startswith("data:"):
                # Parse data URL: data:image/png;base64,<data>
                try:
                    header, data = url.split(",", 1)
                    media_type = header.split(":")[1].split(";")[0]
                    result.append({
                        "type": "image",
                        "source": {
                            "type": "base64",
                            "media_type": media_type,
                            "data": data,
                        },
                    })
                except (ValueError, IndexError):
                    logger.warning(f"Failed to parse image data URL")
            else:
                # URL reference — Anthropic requires base64, log warning
                logger.warning(f"Anthropic requires base64 images, got URL: {url[:50]}...")
                # Could fetch and convert, but that's expensive — skip for now
    
    return result if result else ""


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

    Returns (system_prompt, messages) where system is extracted from the messages list.
    
    Handles:
    - System messages → separate system field
    - Assistant messages with tool_calls → tool_use blocks
    - Tool messages → tool_result blocks
    - Multimodal content (images)
    """
    system = None
    anthropic_msgs = []

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

        if role == "system":
            # Anthropic uses a separate system field
            sys_content = content if isinstance(content, str) else str(content)
            if system:
                system += "\n\n" + sys_content
            else:
                system = sys_content
            continue

        if role == "assistant":
            # Check for tool_calls (OpenAI format)
            tool_calls = msg.get("tool_calls", [])
            if tool_calls:
                # Convert to Anthropic tool_use blocks
                blocks = []
                if content:
                    blocks.append({"type": "text", "text": content if isinstance(content, str) else str(content)})
                for tc in tool_calls:
                    if tc.get("type") == "function":
                        fn = tc.get("function", {})
                        args = fn.get("arguments", "{}")
                        if isinstance(args, str):
                            try:
                                args = json.loads(args)
                            except json.JSONDecodeError:
                                args = {}
                        blocks.append({
                            "type": "tool_use",
                            "id": tc.get("id", ""),
                            "name": fn.get("name", ""),
                            "input": args,
                        })
                anthropic_msgs.append({"role": "assistant", "content": blocks})
            else:
                anthropic_msgs.append({"role": "assistant", "content": _convert_content_to_anthropic(content)})
        
        elif role == "tool":
            anthropic_msgs.append({
                "role": "user",
                "content": [{
                    "type": "tool_result",
                    "tool_use_id": msg.get("tool_call_id", ""),
                    "content": content if isinstance(content, str) else str(content),
                }],
            })
        else:
            # User message
            anthropic_msgs.append({"role": "user", "content": _convert_content_to_anthropic(content)})

    return system, anthropic_msgs


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

    anthropic_tools = []
    for tool in tools:
        if tool.get("type") == "function":
            fn = tool["function"]
            anthropic_tools.append({
                "name": fn["name"],
                "description": fn.get("description", ""),
                "input_schema": fn.get("parameters", {"type": "object", "properties": {}}),
            })
    return anthropic_tools or None


def _convert_tool_choice_to_anthropic(tool_choice) -> dict | None:
    """Convert OpenAI tool_choice to Anthropic format.
    
    OpenAI:
    - "auto" | "none" | "required"
    - {type: "function", function: {name: "tool_name"}}
    
    Anthropic:
    - {type: "auto"} | {type: "any"} | {type: "tool", name: "tool_name"}
    """
    if tool_choice is None:
        return None
    
    if isinstance(tool_choice, str):
        if tool_choice == "auto":
            return {"type": "auto"}
        elif tool_choice == "none":
            # Anthropic doesn't have "none" — just don't pass tools
            return None
        elif tool_choice == "required":
            return {"type": "any"}
        else:
            # Assume it's a tool name
            return {"type": "tool", "name": tool_choice}
    
    if isinstance(tool_choice, dict):
        if tool_choice.get("type") == "function":
            fn = tool_choice.get("function", {})
            return {"type": "tool", "name": fn.get("name", "")}
    
    return None


class AnthropicBridge(BaseBridge):
    """Bridge for Anthropic Messages API."""

    @property
    def provider_type(self) -> str:
        return "anthropic"

    def _build_url(self) -> str:
        base = self.provider.base_url.rstrip("/")
        return f"{base}/v1/messages"

    def _build_headers(self) -> dict:
        return {
            "Content-Type": "application/json",
            "x-api-key": self.provider.api_key,
            "anthropic-version": "2023-06-01",
        }

    def _build_body(self, request: BridgeRequest, model: ModelConfig) -> dict:
        system, messages = _convert_messages_to_anthropic(request.messages)

        body: dict = {
            "model": model.id,
            "messages": messages,
            "max_tokens": request.max_tokens or 4096,
        }
        if system:
            body["system"] = system
        if request.temperature is not None:
            body["temperature"] = request.temperature
        if request.top_p is not None:
            body["top_p"] = request.top_p
        if request.stream:
            body["stream"] = True

        tools = _convert_tools_to_anthropic(request.tools)
        if tools:
            body["tools"] = tools
            # Only add tool_choice if we have tools
            tool_choice = _convert_tool_choice_to_anthropic(request.tool_choice)
            if tool_choice:
                body["tool_choice"] = tool_choice

        if request.stop:
            body["stop_sequences"] = request.stop if isinstance(request.stop, list) else [request.stop]

        return body

    async def complete(self, request: BridgeRequest, model: ModelConfig) -> BridgeResponse:
        url = self._build_url()
        headers = self._build_headers()
        body = self._build_body(request, model)
        body.pop("stream", None)

        async with self.session.post(url, json=body, headers=headers) as resp:
            if resp.status != 200:
                error_text = await resp.text()
                logger.error(f"Anthropic API error {resp.status}: {error_text[:500]}")
                raise aiohttp.ClientResponseError(
                    resp.request_info, resp.history,
                    status=resp.status, message=error_text[:200],
                )

            data = await resp.json()
            usage = data.get("usage", {})

            # Extract text content
            content_blocks = data.get("content", [])
            text_parts = [b["text"] for b in content_blocks if b.get("type") == "text"]
            content = "\n".join(text_parts)

            # Extract tool calls and convert to OpenAI format
            tool_calls = None
            tool_blocks = [b for b in content_blocks if b.get("type") == "tool_use"]
            if tool_blocks:
                tool_calls = [
                    {
                        "id": b["id"],
                        "type": "function",
                        "function": {
                            "name": b["name"],
                            "arguments": json.dumps(b.get("input", {})),
                        },
                    }
                    for b in tool_blocks
                ]

            return BridgeResponse(
                content=content,
                model=data.get("model", model.id),
                finish_reason=data.get("stop_reason", "stop"),
                tokens_in=usage.get("input_tokens", 0),
                tokens_out=usage.get("output_tokens", 0),
                # Anthropic prompt caching tokens
                cache_read_tokens=usage.get("cache_read_input_tokens", 0),
                cache_creation_tokens=usage.get("cache_creation_input_tokens", 0),
                tool_calls=tool_calls,
                raw=data,
            )

    async def stream(self, request: BridgeRequest, model: ModelConfig) -> StreamResult:
        """Stream Anthropic SSE and translate to OpenAI SSE format on the fly.
        
        Captures usage from message_start and message_delta events.
        """
        # Create result with shared usage object that the generator will populate
        result = StreamResult.__new__(StreamResult)
        result.usage = StreamUsage()
        result._iterator = self._stream_inner(request, model, result.usage)
        return result

    async def _stream_inner(
        self, request: BridgeRequest, model: ModelConfig, usage: StreamUsage
    ) -> AsyncIterator[str]:
        """Inner streaming generator that captures usage."""
        url = self._build_url()
        headers = self._build_headers()
        body = self._build_body(request, model)
        body["stream"] = True

        async with self.session.post(url, json=body, headers=headers) as resp:
            if resp.status != 200:
                error_text = await resp.text()
                logger.error(f"Anthropic API error {resp.status}: {error_text[:500]}")
                raise aiohttp.ClientResponseError(
                    resp.request_info, resp.history,
                    status=resp.status, message=error_text[:200],
                )

            # Translate Anthropic SSE events to OpenAI SSE format
            async for line in resp.content:
                decoded = line.decode("utf-8", errors="replace").strip()
                if not decoded or not decoded.startswith("data: "):
                    if decoded:
                        yield decoded
                    continue

                data_str = decoded[6:]
                if data_str == "[DONE]":
                    yield "data: [DONE]"
                    continue

                try:
                    event = json.loads(data_str)
                    event_type = event.get("type", "")

                    # message_start contains input token counts
                    if event_type == "message_start":
                        msg = event.get("message", {})
                        msg_usage = msg.get("usage", {})
                        usage.tokens_in = msg_usage.get("input_tokens", 0)
                        usage.cache_read_tokens = msg_usage.get("cache_read_input_tokens", 0)
                        usage.cache_creation_tokens = msg_usage.get("cache_creation_input_tokens", 0)

                    elif event_type == "content_block_delta":
                        delta = event.get("delta", {})
                        if delta.get("type") == "text_delta":
                            openai_chunk = {
                                "choices": [{
                                    "index": 0,
                                    "delta": {"content": delta.get("text", "")},
                                    "finish_reason": None,
                                }],
                                "model": model.id,
                            }
                            yield f"data: {json.dumps(openai_chunk)}"

                    # message_delta contains output token count
                    elif event_type == "message_delta":
                        delta_usage = event.get("usage", {})
                        usage.tokens_out = delta_usage.get("output_tokens", 0)

                    elif event_type == "message_stop":
                        openai_chunk = {
                            "choices": [{
                                "index": 0,
                                "delta": {},
                                "finish_reason": "stop",
                            }],
                            "model": model.id,
                            # Include usage in final chunk (OpenAI style)
                            "usage": {
                                "prompt_tokens": usage.tokens_in,
                                "completion_tokens": usage.tokens_out,
                                "total_tokens": usage.tokens_in + usage.tokens_out,
                            },
                        }
                        yield f"data: {json.dumps(openai_chunk)}"
                        yield "data: [DONE]"

                except json.JSONDecodeError:
                    yield decoded
