from math import ceil
import mmap
from tempfile import TemporaryFile
from asyncio import StreamReader, StreamWriter
from io import BufferedRWPair, RawIOBase, BufferedReader, BytesIO, UnsupportedOperation
from json import JSONEncoder
from typing import IO, Union, Dict, BinaryIO
from dataclasses import is_dataclass, asdict
import json
from .meta import Meta


Binary = Union[bytes, bytearray, memoryview, mmap.mmap]


class MetaEncoder(JSONEncoder):
    def default(self, obj: Meta) -> Dict:
        if is_dataclass(obj):
            return asdict(obj)
        elif isinstance(obj, dict):
            return obj
        else:
            raise TypeError
class Envelope:
    _MAX_SIZE = 128*1024*1024 # 128 Mb

    def __init__(self, meta: Meta, data : Binary = b''):
        self.meta = meta
        self.data = data

    def __str__(self):
        return str(self.meta)

    @staticmethod
    def read(input: BufferedReader | BytesIO) -> "Envelope":
        assert b'#~'   == input.read(2), "Envelope byte sequence doesn't start with b'#~'"
        assert b'DF02' == input.read(4), 'Envelope type (version) is not DF02'
        input.read(2) # MetaType: XML or YAML

        metaLength = int.from_bytes(input.read(4))
        dataLength = int.from_bytes(input.read(4))

        meta = json.loads(input.read(metaLength))

        if dataLength < Envelope._MAX_SIZE:
            data = input.read(dataLength)
        else:
            data = mmap.mmap(input.fileno(), dataLength, offset = input.tell())
            # will raise an exception if 'input' is BytesIO and dataLength > _MAX_SIZE

        assert b'~#' == input.read(2), "Envelope byte sequence doesn't end with b'~#'"

        return Envelope(meta, data)


    @staticmethod
    def from_bytes(buffer: bytes) -> "Envelope":
        return Envelope.read(BytesIO(buffer))

    def to_bytes(self) -> bytes:
        output = BytesIO()
        self.write_to(output)
        output.seek(0)
        return output.read()
    def write_to(self, output: RawIOBase | BytesIO):
        output.write(b'#~')
        output.write(b'DF02')
        output.write(b'..')
        meta_str = bytes(json.dumps(self.meta), 'utf8')

        output.write(len(meta_str ).to_bytes(4))
        output.write(len(self.data).to_bytes(4))
        output.write(meta_str)
        output.write(self.data)

        output.write(b'~#')

    @staticmethod
    async def async_read(reader: StreamReader) -> "Envelope":
        assert b'#~' == await reader.read(2), "Envelope header doesn't start with b'#~'"
        assert b'DF02' == await reader.read(4), "Envelope type (version) is not DF02"
        await reader.read(2)  # MetaType: XML or YAML

        metaLength = int.from_bytes(await reader.read(4))
        dataLength = int.from_bytes(await reader.read(4))

        assert b'~#\r\n' == await reader.read(4), r"Envelope header doesn't end with b'~#\r\n'"

        meta = json.loads(await reader.read(metaLength))
        tmp_file = None

        if dataLength < Envelope._MAX_SIZE:
            data = await reader.read(dataLength)
        else:
            tmp_file = TemporaryFile('rwb')
            # write input to tmp in chuncks:
            for _ in range(ceil(dataLength / Envelope._MAX_SIZE)):
                tmp_file.write(await reader.read(Envelope._MAX_SIZE))
            tmp_file.flush()

            data = mmap.mmap(tmp_file.fileno(), dataLength)

        return Envelope(meta, data, tmp_file)

    async def async_write_to(self, writer: StreamWriter):
        self.write_to(writer)