"""EventStream — async iterator-based streaming for LLM responses."""

from __future__ import annotations

import asyncio
from typing import Any, Callable, Generic, TypeVar

T = TypeVar("T")
R = TypeVar("R")


class EventStream(Generic[T, R]):
    """Generic event stream class for async iteration.

    Mirrors the TypeScript EventStream<T, R> with asyncio.Queue.
    """

    def __init__(
        self,
        is_complete: Callable[[T], bool],
        extract_result: Callable[[T], R],
    ) -> None:
        self._is_complete = is_complete
        self._extract_result = extract_result
        self._queue: asyncio.Queue[T | None] = asyncio.Queue()
        self._done = False
        self._final_result_future: asyncio.Future[R] = asyncio.get_event_loop().create_future()

    def push(self, event: T) -> None:
        """Push an event to the stream."""
        if self._done:
            return

        if self._is_complete(event):
            self._done = True
            if not self._final_result_future.done():
                self._final_result_future.set_result(self._extract_result(event))

        self._queue.put_nowait(event)

    def end(self, result: R | None = None) -> None:
        """End the stream, optionally providing a final result."""
        self._done = True
        if result is not None and not self._final_result_future.done():
            self._final_result_future.set_result(result)
        # Signal end with None sentinel
        self._queue.put_nowait(None)

    async def result(self) -> R:
        """Get the final result (waits for stream completion)."""
        return await self._final_result_future

    def __aiter__(self):
        return self

    async def __anext__(self) -> T:
        while True:
            item = await self._queue.get()
            if item is None:
                raise StopAsyncIteration
            return item


class AssistantMessageEventStream(EventStream[dict[str, Any], dict[str, Any]]):
    """Typed event stream for assistant message events."""

    def __init__(self) -> None:
        super().__init__(
            is_complete=lambda event: event.get("type") in ("done", "error"),
            extract_result=lambda event: (
                event.get("message") if event.get("type") == "done" else event.get("error", event)
            ),
        )


def create_assistant_message_event_stream() -> AssistantMessageEventStream:
    """Factory function for AssistantMessageEventStream."""
    return AssistantMessageEventStream()
