"""Collapse an Anthropic event stream into one ``Message``.

Mirrors the Anthropic SDK's own accumulator: folds content-block deltas
into finalized blocks, indexed by their ``content_block`` index.
"""

from __future__ import annotations

import json
from typing import TYPE_CHECKING, Any, Literal, cast

from anthropic.types import (
    InputJSONDelta,
    Message,
    RawContentBlockDeltaEvent,
    RawContentBlockStartEvent,
    RawContentBlockStopEvent,
    RawMessageDeltaEvent,
    RawMessageStartEvent,
    SignatureDelta,
    TextBlock,
    TextDelta,
    ThinkingBlock,
    ThinkingDelta,
    ToolUseBlock,
    Usage,
)

if TYPE_CHECKING:
    from collections.abc import AsyncIterator

    from beaver_gateway.events.stream import MessageStreamEvent, StopReason


__all__ = ["StreamAccumulator", "accumulate"]


class StreamAccumulator:
    """Folds a stream of events into one ``Message``, incrementally.

    Call :meth:`feed` for each event, then :meth:`finalize` once.
    """

    __slots__ = (
        "_blocks",
        "_json_buffers",
        "_message_id",
        "_role",
        "_stop_reason",
        "_stop_sequence",
        "_usage",
    )

    def __init__(self) -> None:
        self._message_id: str = ""
        self._role: str = "assistant"
        self._usage: Usage = Usage(input_tokens=0, output_tokens=0)
        self._blocks: dict[int, dict[str, Any]] = {}
        self._json_buffers: dict[int, str] = {}
        self._stop_reason: str | None = None
        self._stop_sequence: str | None = None

    def feed(self, ev: MessageStreamEvent) -> None:
        """Fold one stream event into the in-progress message state.

        Uses ``isinstance`` rather than ``ev.type == ...`` so ty narrows
        the event type from the class, not a string comparison.
        """
        if isinstance(ev, RawMessageStartEvent):
            self._message_id = ev.message.id
            self._role = ev.message.role
            self._usage = ev.message.usage
        elif isinstance(ev, RawContentBlockStartEvent):
            self._blocks[ev.index] = ev.content_block.model_dump()
            if self._blocks[ev.index].get("type") == "tool_use":
                self._json_buffers[ev.index] = ""
        elif isinstance(ev, RawContentBlockDeltaEvent):
            blk = self._blocks.setdefault(ev.index, {})
            delta = ev.delta
            if isinstance(delta, TextDelta):
                blk["text"] = blk.get("text", "") + delta.text
            elif isinstance(delta, InputJSONDelta):
                self._json_buffers[ev.index] = (
                    self._json_buffers.get(ev.index, "") + delta.partial_json
                )
            elif isinstance(delta, ThinkingDelta):
                blk["thinking"] = blk.get("thinking", "") + delta.thinking
            elif isinstance(delta, SignatureDelta):
                blk["signature"] = delta.signature
        elif isinstance(ev, RawContentBlockStopEvent):
            blk = self._blocks.get(ev.index, {})
            if blk.get("type") == "tool_use":
                raw = self._json_buffers.pop(ev.index, "")
                blk["input"] = json.loads(raw) if raw.strip() else {}
        elif isinstance(ev, RawMessageDeltaEvent):
            self._stop_reason = ev.delta.stop_reason
            self._stop_sequence = ev.delta.stop_sequence
            if ev.usage.output_tokens:
                self._usage = Usage.model_validate(
                    {
                        **self._usage.model_dump(),
                        "output_tokens": ev.usage.output_tokens,
                    }
                )

    def finalize(self, *, model: str) -> Message:
        """Build the finalized ``Message`` from accumulated block state.

        ``role`` is always ``"assistant"`` at the wire level; the cast
        avoids a runtime check ty would otherwise require.
        """
        content: list[Any] = []
        for idx in sorted(self._blocks):
            bd = self._blocks[idx]
            btype = bd.get("type")
            if btype == "text":
                content.append(TextBlock.model_validate(bd))
            elif btype == "tool_use":
                content.append(ToolUseBlock.model_validate(bd))
            elif btype == "thinking":
                content.append(ThinkingBlock.model_validate(bd))

        return Message(
            id=self._message_id or "msg_unknown",
            type="message",
            role=cast("Literal['assistant']", self._role),
            model=model,
            content=content,
            stop_reason=cast("StopReason | None", self._stop_reason),
            stop_sequence=self._stop_sequence,
            usage=self._usage,
        )


async def accumulate(
    events: AsyncIterator[MessageStreamEvent], *, model: str
) -> Message:
    """Drain ``events`` into one ``Message`` (used by ``stream=false`` paths)."""
    acc = StreamAccumulator()
    async for ev in events:
        acc.feed(ev)
    return acc.finalize(model=model)