"""
Data sending utilities for aiohttp.
"""
from typing import Any, AsyncIterable, Union, Optional

import aiohttp
from aiofiles.threadpool.binary import AsyncBufferedIOBase


class SendEmptyFileError(Exception):
    pass


async def send_open_file(f: AsyncBufferedIOBase, chunk: bytes, until: int, chunk_size: int):
    """
    Buffered file sender for extraction, with maximum size control.
    """
    yield chunk
    pos = len(chunk)
    chunk = await f.read(chunk_size)
    while chunk:
        yield chunk
        pos += len(chunk)
        if pos >= until:
            break
        chunk = await f.read(chunk_size)


async def send_memory_view(buf: memoryview, until: int, chunk_size: int):
    """
    Buffered file sender for extraction, with maximum size control.
    """
    size = len(buf)
    pos = 0
    while pos < size:
        tmp = pos + chunk_size
        yield buf[pos:tmp]
        pos = tmp
        if pos >= until:
            break


class SendFileHandler(object):
    __slots__ = ("filename", "content_type", "iter")

    def __init__(
        self, filename: str, content_type: Optional[str], iter_: AsyncIterable[Union[bytes, bytearray, memoryview]]
    ):
        self.filename = filename
        self.content_type = content_type
        self.iter = iter_


class SendFilePayload(aiohttp.AsyncIterablePayload):
    def __init__(self, value: SendFileHandler, *args: Any, **kwargs: Any) -> None:
        if not isinstance(value, SendFileHandler):
            raise TypeError(f"value argument must be a SendFileHandler, got {type(value)!r}")

        kwargs["content_type"] = value.content_type
        aiohttp.AsyncIterablePayload.__init__(self, value.iter, *args, **kwargs)
        self.set_content_disposition("attachment", filename=value.filename)


aiohttp.PAYLOAD_REGISTRY.register(SendFilePayload, SendFileHandler)
