from collections.abc import AsyncIterable, AsyncIterator, Buffer, Iterable, Iterator
from io import BytesIO, UnsupportedOperation
from typing import Any, Protocol, cast, overload, runtime_checkable

from storix.types import AsyncDataBuffer, DataBuffer


type Chunk = str | bytes | Buffer


class _IReadableStreamBase(Protocol):
    def seekable(self) -> bool: ...
    def tell(self, *args: Any, **kwargs: Any) -> int: ...
    def seek(self, *args: Any, **kwargs: Any) -> int: ...


@runtime_checkable
class _ReadableStream(_IReadableStreamBase, Protocol):
    def read(self, size: int | None = ..., /) -> bytes: ...


@runtime_checkable
class _AsyncReadableStream(_IReadableStreamBase, Protocol):
    async def read(self, size: int | None = ..., /) -> bytes: ...


@overload
def normalize_data[AnyStr: (str, bytes)](
    data: DataBuffer[AnyStr],
    *,
    encoding: str = 'utf-8',
) -> _ReadableStream: ...
@overload
def normalize_data[AnyStr: (str, bytes)](
    data: AsyncIterable[AnyStr | Buffer],
    *,
    encoding: str = 'utf-8',
) -> _AsyncReadableStream: ...


def normalize_data[AnyStr: (str, bytes)](
    data: AsyncDataBuffer[AnyStr],
    *,
    encoding: str = 'utf-8',
) -> _ReadableStream | _AsyncReadableStream:
    """Normalize data into readable."""
    if isinstance(data, str):
        data = data.encode(encoding=encoding)  # type: ignore[assignment]

    # cover both mutable and immutable buffers
    if isinstance(data, Buffer):
        return BytesIO(data)

    if callable(getattr(data, 'read', None)):
        return data  # type: ignore[return-value]

    if isinstance(data, AsyncIterable):
        return _AsyncIterStreamer(data, encoding=encoding)

    if isinstance(data, Iterable):
        return _IterStreamer(cast(Iterable[Chunk], data), encoding=encoding)

    msg = f'Unsupported data type: {type(data)}'
    raise TypeError(msg)


class _BaseStreamer:
    def seekable(self) -> bool:
        return False

    def tell(self, *_: Any, **__: Any) -> int:
        msg = 'Data generator does not support tell().'
        raise UnsupportedOperation(msg)

    def seek(self, *_: Any, **__: Any) -> int:
        msg = 'Data generator is not seekable().'
        raise UnsupportedOperation(msg)


class _IterStreamer(_BaseStreamer):
    def __init__(self, generator: Iterable[Chunk], *, encoding: str = 'utf-8') -> None:
        self.generator: Iterable[Chunk] = generator
        self.iterator: Iterator[Chunk] = iter(generator)
        self.encoding = encoding
        self.leftover = b''

    def __len__(self) -> int:
        # ignore[attr-defined] because not all iterables implement __len__
        return self.generator.__len__()  # type: ignore[attr-defined]

    def __next__(self) -> Chunk:
        return next(self.iterator)

    def __iter__(self) -> Iterator[Chunk]:
        return self.iterator

    def read(self, size: int | None = None, /) -> bytes:
        size = size or 1024

        data: bytes = self.leftover
        count = len(data)
        try:
            while count < size:
                chunk = self.__next__()
                if isinstance(chunk, str):
                    mv = memoryview(chunk.encode(self.encoding))
                elif isinstance(chunk, Buffer):
                    mv = memoryview(chunk)
                else:
                    msg = f'Unsupported chunk type: {type(chunk)}'
                    raise TypeError(msg)

                data += mv
                count += len(mv)

        except StopIteration:
            self.leftover = b''

        else:
            self.leftover = data[size:]

        return data[:size]


class _AsyncIterStreamer(_BaseStreamer):
    def __init__(
        self, generator: AsyncIterable[Chunk], *, encoding: str = 'utf-8'
    ) -> None:
        self.iterator: AsyncIterator[Chunk] = generator.__aiter__()
        self.encoding = encoding
        self.leftover = b''

    async def read(self, size: int | None = None, /) -> bytes:
        size = size or 1024

        data: bytes = self.leftover
        count = len(data)
        try:
            while count < size:
                chunk = await self.iterator.__anext__()
                if isinstance(chunk, str):
                    mv = memoryview(chunk.encode(self.encoding))
                elif isinstance(chunk, Buffer):
                    mv = memoryview(chunk)
                else:
                    msg = f'Unsupported chunk type: {type(chunk)}'
                    raise TypeError(msg)

                data += mv
                count += len(mv)

        except StopAsyncIteration:
            self.leftover = b''

        else:
            self.leftover = data[size:]

        return data[:size]
