"""Standard Anthropic response decoding."""

import json
from typing import Any, TypeAlias, cast

from anthropic import types as anthropic_types
from anthropic.lib.streaming import AsyncMessageStreamManager, MessageStreamManager
from anthropic.types.beta import BetaMessageDeltaUsage, BetaUsage

from ....content import (
    AssistantContentPart,
    Text,
    TextChunk,
    TextEndChunk,
    TextStartChunk,
    Thought,
    ThoughtChunk,
    ThoughtEndChunk,
    ThoughtStartChunk,
    ToolCall,
    ToolCallChunk,
    ToolCallEndChunk,
    ToolCallStartChunk,
)
from ....messages import AssistantMessage
from ....responses import (
    AsyncChunkIterator,
    ChunkIterator,
    FinishReason,
    FinishReasonChunk,
    ProviderToolUsage,
    RawMessageChunk,
    RawStreamEventChunk,
    Usage,
    UsageDeltaChunk,
)
from ..model_id import AnthropicModelId, model_name

ANTHROPIC_FINISH_REASON_MAP = {
    "max_tokens": FinishReason.MAX_TOKENS,
    "refusal": FinishReason.REFUSAL,
}


def _decode_assistant_content(
    content: anthropic_types.ContentBlock,
) -> AssistantContentPart | None:
    """Convert Anthropic content block to mirascope AssistantContentPart."""
    if content.type == "text":
        return Text(text=content.text)
    elif content.type == "tool_use":
        return ToolCall(
            id=content.id,
            name=content.name,
            args=json.dumps(content.input),
        )
    elif content.type == "thinking":
        return Thought(thought=content.thinking)
    elif content.type in ("server_tool_use", "web_search_tool_result"):
        return None  # Skip server-side tool content, preserved in raw_message
    else:
        raise NotImplementedError(
            f"Support for content type `{content.type}` is not yet implemented."
        )


def extract_tool_usage(
    usage: (
        anthropic_types.Usage
        | anthropic_types.MessageDeltaUsage
        | BetaUsage
        | BetaMessageDeltaUsage
    ),
) -> list[ProviderToolUsage] | None:
    """Extract provider tool usage from Anthropic usage object."""
    server_tool_use = getattr(usage, "server_tool_use", None)
    if server_tool_use is None:
        return None

    tools: list[ProviderToolUsage] = []

    # Web search
    web_search_requests = getattr(server_tool_use, "web_search_requests", None)
    if web_search_requests and web_search_requests > 0:
        tools.append(
            ProviderToolUsage(name="web_search", call_count=web_search_requests)
        )

    return tools if tools else None


def decode_usage(
    usage: anthropic_types.Usage | BetaUsage,
) -> Usage:
    """Convert Anthropic Usage (or BetaUsage) to Mirascope Usage."""

    cache_read_tokens = usage.cache_read_input_tokens or 0
    cache_write_tokens = usage.cache_creation_input_tokens or 0
    input_tokens = usage.input_tokens + cache_read_tokens + cache_write_tokens
    output_tokens = usage.output_tokens
    return Usage(
        input_tokens=input_tokens,
        output_tokens=output_tokens,
        cache_read_tokens=cache_read_tokens,
        cache_write_tokens=cache_write_tokens,
        reasoning_tokens=0,
        provider_tool_usage=extract_tool_usage(usage),
        raw=usage,
    )


def decode_response(
    response: anthropic_types.Message,
    model_id: AnthropicModelId,
    *,
    include_thoughts: bool,
) -> tuple[AssistantMessage, FinishReason | None, Usage]:
    """Convert Anthropic message to mirascope AssistantMessage and usage."""
    content = [
        part
        for part in (_decode_assistant_content(block) for block in response.content)
        if part is not None
    ]
    if not include_thoughts:
        content = [part for part in content if part.type != "thought"]
    assistant_message = AssistantMessage(
        content=content,
        provider_id="anthropic",
        model_id=model_id,
        provider_model_name=model_name(model_id),
        raw_message={
            "role": response.role,
            "content": [part.model_dump() for part in response.content],
        },
    )
    finish_reason = (
        ANTHROPIC_FINISH_REASON_MAP.get(response.stop_reason)
        if response.stop_reason
        else None
    )
    usage = decode_usage(response.usage)
    return assistant_message, finish_reason, usage


ContentBlock: TypeAlias = (
    anthropic_types.TextBlockParam
    | anthropic_types.ThinkingBlockParam
    | anthropic_types.ToolUseBlockParam
    | anthropic_types.RedactedThinkingBlockParam
)


class _AnthropicChunkProcessor:
    """Processes Anthropic stream events and maintains state across events."""

    def __init__(self, *, include_thoughts: bool) -> None:
        self.current_block_param: ContentBlock | None = None
        self.accumulated_tool_json: str = ""
        self.accumulated_blocks: list[ContentBlock] = []
        self.include_thoughts = include_thoughts

    def process_event(
        self, event: anthropic_types.RawMessageStreamEvent
    ) -> ChunkIterator:
        """Process a single Anthropic event and yield the appropriate content chunks."""
        yield RawStreamEventChunk(raw_stream_event=event)

        if event.type == "content_block_start":
            content_block = event.content_block

            if content_block.type == "text":
                self.current_block_param = {
                    "type": "text",
                    "text": content_block.text,
                }
                yield TextStartChunk()
            elif content_block.type == "tool_use":
                self.current_block_param = {
                    "type": "tool_use",
                    "id": content_block.id,
                    "name": content_block.name,
                    "input": {},
                }
                self.accumulated_tool_json = ""
                yield ToolCallStartChunk(
                    id=content_block.id,
                    name=content_block.name,
                )
            elif content_block.type == "thinking":
                self.current_block_param = {
                    "type": "thinking",
                    "thinking": "",
                    "signature": "",
                }
                if self.include_thoughts:
                    yield ThoughtStartChunk()
            elif content_block.type == "redacted_thinking":  # pragma: no cover
                self.current_block_param = {
                    "type": "redacted_thinking",
                    "data": content_block.data,
                }
            elif content_block.type in ("server_tool_use", "web_search_tool_result"):
                pass  # Skip server-side tool content
            else:
                raise NotImplementedError

        elif event.type == "content_block_delta":
            if self.current_block_param is None:
                return  # Skip deltas for server-side tool content

            delta = event.delta
            if delta.type == "text_delta":
                if self.current_block_param["type"] != "text":  # pragma: no cover
                    raise RuntimeError(
                        f"Received text_delta for {self.current_block_param['type']} block"
                    )
                self.current_block_param["text"] += delta.text
                yield TextChunk(delta=delta.text)
            elif delta.type == "input_json_delta":
                if self.current_block_param["type"] != "tool_use":  # pragma: no cover
                    raise RuntimeError(
                        f"Received input_json_delta for {self.current_block_param['type']} block"
                    )
                self.accumulated_tool_json += delta.partial_json
                yield ToolCallChunk(
                    id=self.current_block_param["id"], delta=delta.partial_json
                )
            elif delta.type == "thinking_delta":
                if self.current_block_param["type"] != "thinking":  # pragma: no cover
                    raise RuntimeError(
                        f"Received thinking_delta for {self.current_block_param['type']} block"
                    )
                self.current_block_param["thinking"] += delta.thinking
                if self.include_thoughts:
                    yield ThoughtChunk(delta=delta.thinking)
            elif delta.type == "signature_delta":
                if self.current_block_param["type"] != "thinking":  # pragma: no cover
                    raise RuntimeError(
                        f"Received signature_delta for {self.current_block_param['type']} block"
                    )
                self.current_block_param["signature"] += delta.signature
            elif delta.type == "citations_delta":
                pass  # Skip citations delta, preserved in raw_message
            else:
                raise RuntimeError(
                    f"Received unsupported delta type: {delta.type}"
                )  # pragma: no cover

        elif event.type == "content_block_stop":
            if self.current_block_param is None:
                return  # Skip stop for server-side tool content

            block_type = self.current_block_param["type"]

            if block_type == "text":
                yield TextEndChunk()
            elif block_type == "tool_use":
                if self.current_block_param["type"] != "tool_use":  # pragma: no cover
                    raise RuntimeError(
                        f"Block type mismatch: stored {self.current_block_param['type']}, expected tool_use"
                    )
                self.current_block_param["input"] = (
                    json.loads(self.accumulated_tool_json)
                    if self.accumulated_tool_json
                    else {}
                )
                yield ToolCallEndChunk(id=self.current_block_param["id"])
            elif block_type == "thinking":
                if self.include_thoughts:
                    yield ThoughtEndChunk()
            else:
                raise NotImplementedError

            self.accumulated_blocks.append(self.current_block_param)
            self.current_block_param = None

        elif event.type == "message_delta":
            if event.delta.stop_reason:
                finish_reason = ANTHROPIC_FINISH_REASON_MAP.get(event.delta.stop_reason)
                if finish_reason is not None:
                    yield FinishReasonChunk(finish_reason=finish_reason)

            # Emit usage delta
            usage = event.usage
            yield UsageDeltaChunk(
                input_tokens=usage.input_tokens or 0,
                output_tokens=usage.output_tokens,
                cache_read_tokens=usage.cache_read_input_tokens or 0,
                cache_write_tokens=usage.cache_creation_input_tokens or 0,
                reasoning_tokens=0,
                provider_tool_usage=extract_tool_usage(usage),
            )

    def raw_message_chunk(self) -> RawMessageChunk:
        return RawMessageChunk(
            raw_message=cast(
                dict[str, Any],
                {
                    "role": "assistant",
                    "content": self.accumulated_blocks,
                },
            )
        )


def decode_stream(
    anthropic_stream_manager: MessageStreamManager, *, include_thoughts: bool
) -> ChunkIterator:
    """Returns a ChunkIterator converted from an Anthropic MessageStreamManager."""
    processor = _AnthropicChunkProcessor(include_thoughts=include_thoughts)
    with anthropic_stream_manager as stream:
        for event in stream._raw_stream:  # pyright: ignore[reportPrivateUsage]
            yield from processor.process_event(event)
    yield processor.raw_message_chunk()


async def decode_async_stream(
    anthropic_stream_manager: AsyncMessageStreamManager, *, include_thoughts: bool
) -> AsyncChunkIterator:
    """Returns an AsyncChunkIterator converted from an Anthropic MessageStreamManager."""
    processor = _AnthropicChunkProcessor(include_thoughts=include_thoughts)
    async with anthropic_stream_manager as stream:
        async for event in stream._raw_stream:  # pyright: ignore[reportPrivateUsage]
            for item in processor.process_event(event):
                yield item
    yield processor.raw_message_chunk()
