"""Streaming with tool call accumulation."""

from __future__ import annotations

from dataclasses import dataclass, field
from typing import TYPE_CHECKING

if TYPE_CHECKING:
    from collections.abc import AsyncIterator

    from koan.llm.client import LLMClient
    from koan.llm.config import LLMConfig
    from koan.types import Message


@dataclass
class AccumulatedToolCall:
    """A fully accumulated tool call from streaming deltas."""

    id: str
    name: str
    arguments: str


@dataclass
class StreamToolEvent:
    """An event from stream_with_tools — either text or completed tool calls."""

    text: str = ""
    tool_calls: list[AccumulatedToolCall] = field(default_factory=list)
    finish_reason: str | None = None

    @property
    def is_text(self) -> bool:
        return bool(self.text)

    @property
    def has_tool_calls(self) -> bool:
        return bool(self.tool_calls)


async def stream_with_tools(
    client: LLMClient,
    config: LLMConfig,
    messages: list[Message],
) -> AsyncIterator[StreamToolEvent]:
    """Stream LLM response, accumulating tool calls and yielding events.

    Yields ``StreamToolEvent`` objects. Text deltas are yielded immediately.
    Tool calls are accumulated across chunks and yielded as a batch when
    the stream signals ``finish_reason="tool_calls"`` (OpenAI) or
    ``finish_reason="tool_use"`` (Anthropic).

    Usage::

        async for event in stream_with_tools(client, config, messages):
            if event.is_text:
                print(event.text, end="")
            if event.has_tool_calls:
                for tc in event.tool_calls:
                    result = execute_tool(tc.name, tc.arguments)
    """
    pending: dict[str, AccumulatedToolCall] = {}
    current_tool_id: str | None = None

    async for delta in client.stream(config, messages):
        # Text delta — yield immediately
        if delta.text:
            yield StreamToolEvent(text=delta.text)

        # Tool call start (new ID)
        if delta.tool_call_id:
            current_tool_id = delta.tool_call_id
            pending[current_tool_id] = AccumulatedToolCall(
                id=delta.tool_call_id,
                name=delta.tool_name or "",
                arguments="",
            )

        # Tool call name (may arrive separately from ID in some providers)
        if delta.tool_name and current_tool_id and current_tool_id in pending and not pending[current_tool_id].name:
            pending[current_tool_id].name = delta.tool_name

        # Tool call argument delta
        if delta.tool_args_delta and current_tool_id and current_tool_id in pending:
            pending[current_tool_id].arguments += delta.tool_args_delta

        # Finish — yield accumulated tool calls if any
        if delta.finish_reason in ("tool_calls", "tool_use") and pending:
            yield StreamToolEvent(
                tool_calls=list(pending.values()),
                finish_reason=delta.finish_reason,
            )
            pending.clear()
            current_tool_id = None
        elif delta.finish_reason and delta.finish_reason not in ("tool_calls", "tool_use"):
            yield StreamToolEvent(finish_reason=delta.finish_reason)
