"""Anthropic (Claude) provider client."""

import json
import time
import uuid
from collections.abc import AsyncGenerator

import structlog

from robot_resources.proxy.models import (
    ChatCompletionChunk,
    ChatCompletionRequest,
    ChatCompletionResponse,
    Choice,
    Message,
    StreamChoice,
    StreamDelta,
    Usage,
)
from robot_resources.proxy.providers.base import BaseProvider

logger = structlog.get_logger()

# Model name mappings
MODEL_MAPPING = {
    "claude-opus-4-6": "claude-opus-4-6",
    "claude-sonnet-4-6": "claude-sonnet-4-6",
    "claude-haiku-4-5": "claude-haiku-4-5-20251001",
    # Legacy aliases — route to current models
    "claude-opus-4": "claude-opus-4-6",
    "claude-sonnet-4": "claude-sonnet-4-6",
    "claude-3-opus": "claude-opus-4-6",
    "claude-3-sonnet": "claude-sonnet-4-6",
    "claude-3-haiku": "claude-haiku-4-5-20251001",
}


class AnthropicProvider(BaseProvider):
    """Anthropic Claude API provider."""

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

    @property
    def base_url(self) -> str:
        return "https://api.anthropic.com"

    def _get_headers(self, *, api_key: str | None = None) -> dict:
        key = api_key or self.config.api_key
        return {
            "x-api-key": key,
            "anthropic-version": "2023-06-01",
            "content-type": "application/json",
        }

    def _get_completion_endpoint(self) -> str:
        return "/v1/messages"

    def map_model(self, model: str) -> str:
        """Map model name to Anthropic-specific name."""
        return MODEL_MAPPING.get(model, model)

    def format_request(self, request: ChatCompletionRequest) -> dict:
        """Format request for Anthropic Messages API."""
        # Extract system message if present
        system_content = None
        messages = []

        for msg in request.messages:
            if msg.role == "system":
                system_content = msg.content
            else:
                messages.append(
                    {
                        "role": msg.role,
                        "content": msg.content,
                    }
                )

        formatted = {
            "model": self.map_model(request.model),
            "messages": messages,
            "max_tokens": request.max_tokens or 4096,
        }

        if system_content:
            formatted["system"] = system_content

        if request.temperature != 1.0:
            formatted["temperature"] = request.temperature

        return formatted

    def parse_response(self, response: dict) -> ChatCompletionResponse:
        """Parse Anthropic response to standard format."""
        # Extract text from content blocks
        content = ""
        for block in response.get("content", []):
            if block.get("type") == "text":
                content += block.get("text", "")

        usage = response.get("usage", {})

        return ChatCompletionResponse(
            id=response.get("id", ""),
            model=response.get("model", ""),
            choices=[
                Choice(
                    index=0,
                    message=Message(role="assistant", content=content),
                    finish_reason=self._map_stop_reason(response.get("stop_reason")),
                )
            ],
            usage=Usage(
                prompt_tokens=usage.get("input_tokens", 0),
                completion_tokens=usage.get("output_tokens", 0),
                total_tokens=usage.get("input_tokens", 0) + usage.get("output_tokens", 0),
            ),
        )

    def _map_stop_reason(self, stop_reason: str | None) -> str:
        """Map Anthropic stop reason to OpenAI format."""
        mapping = {
            "end_turn": "stop",
            "max_tokens": "length",
            "stop_sequence": "stop",
        }
        return mapping.get(stop_reason or "", "stop")

    async def stream(
        self, request: ChatCompletionRequest, *, api_key: str | None = None,
    ) -> AsyncGenerator[str, None]:
        """Stream a chat completion from Anthropic.

        Anthropic uses event-based SSE (``event:`` + ``data:`` pairs).
        We convert ``content_block_delta`` events into OpenAI-compatible
        ``ChatCompletionChunk`` SSE lines.
        """
        formatted = self.format_request(request)
        formatted["stream"] = True

        url = f"{self.config.base_url or self.base_url}{self._get_completion_endpoint()}"
        self._last_stream_usage = None

        chunk_id = f"chatcmpl-{uuid.uuid4().hex[:12]}"
        created = int(time.time())
        input_tokens = 0
        output_tokens = 0

        async with self._create_client(api_key=api_key) as client:
            async with client.stream("POST", url, json=formatted) as response:
                response.raise_for_status()

                current_event = ""
                async for line in response.aiter_lines():
                    line = line.strip()

                    if not line:
                        continue

                    if line.startswith("event: "):
                        current_event = line[len("event: "):]
                        continue

                    if not line.startswith("data: "):
                        continue

                    raw_data = line[len("data: "):]
                    try:
                        data = json.loads(raw_data)
                    except json.JSONDecodeError:
                        logger.warning(
                            "stream_json_decode_error",
                            provider="anthropic",
                            raw_line=raw_data[:200],
                        )
                        continue

                    if current_event == "message_start":
                        # Extract input token count and emit role chunk
                        usage = data.get("message", {}).get("usage", {})
                        input_tokens = usage.get("input_tokens", 0)

                        chunk = ChatCompletionChunk(
                            id=chunk_id,
                            created=created,
                            model=request.model,
                            choices=[
                                StreamChoice(
                                    delta=StreamDelta(role="assistant", content=""),
                                    finish_reason=None,
                                )
                            ],
                        )
                        yield f"data: {chunk.model_dump_json()}\n\n"

                    elif current_event == "content_block_delta":
                        delta = data.get("delta", {})
                        text = delta.get("text", "")
                        if text:
                            chunk = ChatCompletionChunk(
                                id=chunk_id,
                                created=created,
                                model=request.model,
                                choices=[
                                    StreamChoice(
                                        delta=StreamDelta(content=text),
                                        finish_reason=None,
                                    )
                                ],
                            )
                            yield f"data: {chunk.model_dump_json()}\n\n"

                    elif current_event == "message_delta":
                        usage = data.get("usage", {})
                        output_tokens = usage.get("output_tokens", 0)
                        stop_reason = data.get("delta", {}).get("stop_reason")

                        chunk = ChatCompletionChunk(
                            id=chunk_id,
                            created=created,
                            model=request.model,
                            choices=[
                                StreamChoice(
                                    delta=StreamDelta(),
                                    finish_reason=self._map_stop_reason(stop_reason),
                                )
                            ],
                        )
                        yield f"data: {chunk.model_dump_json()}\n\n"

                    elif current_event == "message_stop":
                        pass  # Final event — we emit [DONE] below

        self._last_stream_usage = Usage(
            prompt_tokens=input_tokens,
            completion_tokens=output_tokens,
            total_tokens=input_tokens + output_tokens,
        )
        yield "data: [DONE]\n\n"
