from __future__ import annotations

import contextlib
from typing import (
    TYPE_CHECKING,
    Any,
    Generator,
    Literal,
    Protocol,
    overload,
    runtime_checkable,
)

if TYPE_CHECKING:
    from pyarrow import RecordBatch

    from .block import SeriesBlock
    from .channel import Channel


@runtime_checkable
class StreamReader(contextlib.AbstractContextManager, Protocol):
    """Non-blocking Arrow RecordBatch stream reader

    May produce multiple independent streams.  Streams should be
    buffered in a queue, and the `read` method should return all
    RecordBatches currently in the queue.

    """

    @property
    def streams(self) -> dict[str, list[Channel]]:
        """stream channels dictionary

        A StreamReader may produce multiple independent streams.  This
        property includes a key for each stream name, who value is the
        list of channel names in each individual stream.

        """
        ...

    def __repr__(self):
        return f"<{self.__class__.__name__}: {self.streams}>"

    def enter(self):
        """open the streams"""
        pass

    def __enter__(self) -> StreamReader:
        self.enter()
        return self

    @overload
    def read(
        self,
        *,
        convert_blocks: Literal[False] = False,
    ) -> Generator[tuple[str, RecordBatch], None, None]: ...

    @overload
    def read(
        self,
        *,
        convert_blocks: Literal[True] = True,
    ) -> Generator[tuple[str, SeriesBlock], None, None]: ...

    def read(
        self,
        *,
        convert_blocks: bool = False,
    ) -> Generator[tuple[str, RecordBatch | SeriesBlock], None, None]:
        """Read all available elements.

        Parameters
        ----------
        convert_blocks : bool
            Convert RecordBatches to SeriesBlocks

        Yields
        ------
        tuple of (stream name, RecordBatch | SeriesBlock)

        """
        ...

    def done(self) -> bool:
        """return True if the collector is done"""
        return False

    def unpack(self) -> Generator[Any, None, None]:
        """synchronously unpack all batch streams into individual elements"""
        while not self.done():
            for _, batch in self.read(convert_blocks=False):
                yield from batch.to_pylist()

    def close(self):
        """close the stream"""
        pass

    def __exit__(self, exc_type, exc_value, traceback) -> Literal[False]:
        self.close()
        return False
