from __future__ import annotations

import asyncio
from collections import deque
from dataclasses import dataclass
from typing import Any, AsyncIterator, Callable, Deque, Generic, TypeVar, cast

TEvent = TypeVar("TEvent")
TResult = TypeVar("TResult")


@dataclass(frozen=True, slots=True)
class _Done:
    pass


_DONE = _Done()


class EventStream(Generic[TEvent, TResult]):
    """Generic push-based async event stream with a separately awaitable final result."""

    def __init__(
        self,
        is_complete: Callable[[TEvent], bool],
        extract_result: Callable[[TEvent], TResult],
    ) -> None:
        self._is_complete = is_complete
        self._extract_result = extract_result
        self._queue: Deque[TEvent] = deque()
        self._waiting: Deque[asyncio.Future[TEvent | _Done]] = deque()
        self._done = False
        self._final_result: asyncio.Future[TResult] = self._new_future()

    def _new_future(self) -> asyncio.Future[Any]:
        try:
            loop = asyncio.get_running_loop()
        except RuntimeError:
            loop = asyncio.get_event_loop()
        return loop.create_future()

    def push(self, event: TEvent) -> None:
        if self._done:
            return

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

        if self._waiting:
            waiter = self._waiting.popleft()
            if not waiter.done():
                waiter.set_result(event)
            return

        self._queue.append(event)

    def end(self, result: TResult | None = None) -> None:
        self._done = True
        if result is not None and not self._final_result.done():
            self._final_result.set_result(result)
        self._finish_waiters()

    def set_exception(self, exc: BaseException) -> None:
        self._done = True
        if not self._final_result.done():
            self._final_result.set_exception(exc)
        self._finish_waiters()

    def _finish_waiters(self) -> None:
        while self._waiting:
            waiter = self._waiting.popleft()
            if not waiter.done():
                waiter.set_result(_DONE)

    async def result(self) -> TResult:
        return await self._final_result

    def __aiter__(self) -> AsyncIterator[TEvent]:
        return self._iterate()

    async def _iterate(self) -> AsyncIterator[TEvent]:
        while True:
            if self._queue:
                yield self._queue.popleft()
                continue

            if self._done:
                return

            waiter: asyncio.Future[TEvent | _Done] = cast(asyncio.Future[TEvent | _Done], self._new_future())
            self._waiting.append(waiter)
            item = await waiter
            if isinstance(item, _Done):
                return
            yield item


class AssistantMessageEventStream(EventStream[dict[str, Any], dict[str, Any]]):
    """Specialized stream whose final result is extracted from done/error events."""

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


def create_assistant_message_event_stream() -> AssistantMessageEventStream:
    return AssistantMessageEventStream()
