# Copyright 1999-2020 Alibaba Group Holding Ltd.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#      http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

from io import BytesIO

from ..serialize.dataserializer import SerialType, CompressType, get_compressobj, get_decompressobj, \
    HEADER_LENGTH, file_header, read_file_header, write_file_header, SERIAL_VERSION


class WorkerBufferIO(object):
    """
    File-like object handling data with file header
    """
    def __init__(self, mode='r', compress_in=None, compress_out=None, block_size=8192):
        """
        :param mode: 'r' indicates read, or 'w' indicates write
        :param compress_in: compression type inside
        :param compress_out: compression type outside
        :param block_size: size of data block when copying
        """
        self._mode = mode
        self._block_size = block_size
        self._compress_type_in = compress_in or CompressType.NONE
        self._compress_type_out = compress_out or CompressType.NONE

        if 'w' in mode:
            self._compressor_in = None
            self._decompressor_out = None
        else:
            self._remain_buf = None
            self._remain_offset = None
            self._block_iterator = self._iter_blocks()

    def __enter__(self):
        return self

    def __exit__(self, exc_type, exc_val, exc_tb):
        self.close()

    def __del__(self):
        self.close()

    def _read_block(self, size):
        """
        Read a data block from data source with given size.
        """
        raise NotImplementedError

    def _read_header(self):
        """
        Read the data header of data source.
        """
        raise NotImplementedError

    def _iter_blocks(self):
        """
        Returns a generator providing data blocks. The sizes of data blocks can vary.
        """
        bio = BytesIO()
        compressor_out = get_compressobj(self._compress_type_out)

        header = self._read_header()
        self._compress_type_in = header.compress
        decompressor_in = get_decompressobj(header.compress)

        new_header = file_header(header.type, header.version, header.nbytes, self._compress_type_out)
        write_file_header(bio, new_header)

        if self._compress_type_in == self._compress_type_out:
            handle_block = lambda v: v

            compressor_out = get_compressobj(CompressType.NONE)
            decompressor_in = get_decompressobj(CompressType.NONE)
        else:
            handle_block = lambda v: compressor_out.compress(decompressor_in.decompress(v))
            if hasattr(compressor_out, 'begin'):
                bio.write(compressor_out.begin())

        # yield file header and compress header
        yield bio.getvalue()

        copy_size = self._block_size
        while True:
            block = self._read_block(copy_size)
            if not block:
                break
            buf = handle_block(block)
            if buf:
                # only yield when some data are produced by the compressor
                yield buf

        if hasattr(compressor_out, 'flush'):
            yield compressor_out.flush()

    def read(self, size=-1):
        bio = BytesIO()
        if self._remain_buf is not None:
            if size < 0 or len(self._remain_buf) <= self._remain_offset + size:
                bio.write(self._remain_buf[self._remain_offset:])
                self._remain_buf = None
            else:
                right = self._remain_offset + size
                bio.write(self._remain_buf[self._remain_offset:right])
                self._remain_offset = right
            if bio.tell() == size:
                return bio.getvalue()
        while True:
            try:
                block = next(self._block_iterator)
            except StopIteration:
                break
            if size < 0 or bio.tell() + len(block) <= size:
                bio.write(block)
                self._remain_buf = None
            else:
                offset = self._remain_offset = size - bio.tell()
                self._remain_buf = buf = memoryview(block)
                bio.write(buf[:offset])
            if bio.tell() == size:
                break
        return bio.getvalue()

    def _write_header(self, header):
        """
        Write a data header to the output file.
        """
        raise NotImplementedError

    def _write_block(self, d):
        """
        Write a data block to the output file.
        """
        raise NotImplementedError

    def _write_with_compression(self, mv):
        data_len = len(mv)
        copy_size = self._block_size
        offset = 0

        if self._compressor_in is not None:
            compressor_in = self._compressor_in
        else:
            compressor_in = self._compressor_in = get_compressobj(self._compress_type_in)
            if hasattr(compressor_in, 'begin') and self._compress_type_in != self._compress_type_out:
                self._write_block(compressor_in.begin())

        if self._compress_type_in == self._compress_type_out:
            self._compressor_in = get_compressobj(CompressType.NONE)
            self._decompressor_out = get_decompressobj(CompressType.NONE)

            handle_block = lambda v: v
        else:
            handle_block = lambda v: compressor_in.compress(decompressor_out.decompress(v))

        decompressor_out = self._decompressor_out

        while offset < data_len:
            right = min(data_len, offset + copy_size)
            cblock = handle_block(mv[offset:right])
            if cblock:
                self._write_block(cblock)
            offset = right

    def write(self, d):
        size = len(d)
        mv = memoryview(d)
        if self._decompressor_out is not None:
            # header already processed, we can continue
            # with decompression
            self._write_with_compression(mv)
        else:
            # header not processed, we need to read header
            # to get compression method and build decompressor
            if size < HEADER_LENGTH:
                raise IOError('Block size too small')

            header = read_file_header(mv)
            self._compress_type_out = header.compress
            new_header = file_header(header.type, header.version, header.nbytes, self._compress_type_in)
            self._write_header(new_header)

            self._decompressor_out = get_decompressobj(header.compress)
            if size > HEADER_LENGTH:
                self._write_with_compression(mv[HEADER_LENGTH:])

    def close(self):
        if 'w' in self._mode and hasattr(self._compressor_in, 'flush'):
            self._write_block(self._compressor_in.flush())

        self._compressor_in = self._decompressor_out = None
        self._remain_buf = self._block_iterator = None


class ArrowBufferIO(WorkerBufferIO):
    """
    File-like object mocking object stored in shared memory as file with header
    """
    def __init__(self, buf, mode='r', compress_in=None, compress_out=None, block_size=8192):
        super().__init__(
            mode=mode, compress_in=compress_in, compress_out=compress_out, block_size=block_size)

        self._buf = buf
        if 'r' in mode:
            self._mv = memoryview(buf)
            self._nbytes = len(buf)
            self._buf_offset = 0
            self._writer = None
        else:
            import pyarrow
            self._writer = pyarrow.FixedSizeBufferWriter(buf)
            self._writer.set_memcopy_threads(6)

    def _read_header(self):
        return file_header(SerialType.ARROW, SERIAL_VERSION, self._nbytes, self._compress_type_in)

    def _read_block(self, size):
        right = min(self._nbytes, self._buf_offset + size)
        ret = self._mv[self._buf_offset:right]
        self._buf_offset = right
        return ret

    def _write_header(self, header):
        pass

    def _write_block(self, d):
        self._writer.write(d)

    def close(self):
        if self._writer is not None:
            self._writer.close()
        self._writer = self._mv = self._buf = None

        super().close()


class FileBufferIO(WorkerBufferIO):
    """
    File-like object handling input of one compression type and outputs another
    """
    def __init__(self, file, mode='r', compress_in=None, compress_out=None,
                 block_size=8192, managed=True, header=None, file_end=None):
        super().__init__(
            mode=mode, compress_in=compress_in, compress_out=compress_out, block_size=block_size)

        self._managed = managed
        self._file = file
        self._header = header
        self._file_end = file_end

    def _read_header(self):
        if self._header is not None:
            return self._header
        return read_file_header(self._file)

    def _read_block(self, size):
        if self._file_end is not None:
            max_size = self._file_end - self._file.tell()
            if size == -1:
                size = max_size
            size = max(min(size, max_size), 0)
        return self._file.read(size)

    def _write_header(self, header):
        return write_file_header(self._file, header)

    def _write_block(self, d):
        self._file.write(d)

    def close(self):
        super().close()
        if self._file:
            if self._managed:
                self._file.close()
            else:
                self._file.flush()
        self._file = None
