from ._frame import *
from ._frame import __doc__ as _doc
__doc__ = _doc


class LZ4FrameCompressor(object):
    """Create a LZ4 compressor object, which can be used to compress data
    incrementally.

    Args:
        block_size (int): Sepcifies the maximum blocksize to use.
            Options:
            - lz4.frame.BLOCKSIZE_DEFAULT or 0: the lz4 library default
            - lz4.frame.BLOCKSIZE_MAX64KB or 4: 64 kB
            - lz4.frame.BLOCKSIZE_MAX256KB or 5: 256 kB
            - lz4.frame.BLOCKSIZE_MAX1MB or 6: 1 MB
            - lz4.frame.BLOCKSIZE_MAX4MB or 7: 4 MB
            If unspecified, will default to lz4.frame.BLOCKSIZE_DEFAULT which
            is equal to lz4.frame.BLOCKSIZE_MAX64KB.
        block_linked (bool): Specifies whether to use block-linked
            compression. If `True`, the compression ratio is improved, especially
            for small block sizes. If `False` the blocks are compressed independently.
            The default is `True`.
        compression_level (int): Specifies the level of compression used.
            Values between 0-16 are valid, with 0 (default) being the
            lowest compression (0-2 are the same value), and 16 the highest.
            Values above 16 will be treated as 16.
            Values between 4-9 are recommended. 0 is the default.
            The following module constants are provided as a convenience:
            - lz4.frame.COMPRESSIONLEVEL_MIN: Minimum compression (0)
            - lz4.frame.COMPRESSIONLEVEL_MINHC: Minimum high-compression (3)
            - lz4.frame.COMPRESSIONLEVEL_MAX: Maximum compression (16)
        content_checksum (bool): Specifies whether to enable checksumming of
            the payload content. If `True` a checksum of the uncompressed data
            is stored at the end of the compressed frame which is checked during
            decompression. The default is `False`.
        auto_flush (bool): When `False`, the LZ4 library may buffer data until a
            block is full. When `True` no buffering occurs, and partially full
            blocks may be returned. The default is `True`.
        return_bytearray (bool): When `False` a bytes object is returned from the
            calls to methods of this class. When `True` a bytearray object will be
            returned. The default is `False`.
    """

    def __init__(self,
                 block_size=BLOCKSIZE_DEFAULT,
                 block_linked=True,
                 compression_level=COMPRESSIONLEVEL_MIN,
                 content_checksum=False,
                 auto_flush=True,
                 return_bytearray=False):
        self.block_size = block_size
        self.block_linked = block_linked
        self.compression_level = compression_level
        self.content_checksum = content_checksum
        self.auto_flush = auto_flush
        self.return_bytearray = return_bytearray
        self._context = create_compression_context()
        self._started = False

    def __enter__(self):
        # All necessary initialization is done in __init__
        return self

    def __exit__(self, exception_type, exception, traceback):
        # The compression context is created with an appropriate destructor, so
        # no need to del it here
        pass

    def compress_begin(self, source_size=0):
        """Begin a compression frame. The returned data contains frame header
        information. The data returned from subsequent calls to ``compress()``
        should be concatenated with this header.

        Args:
            source_size (int): Optionally specify the total size of the
                uncompressed data. If specified, will be stored in the
                compressed frame header as an 8-byte field for later use
                during decompression. Default is 0 (no size stored).

        Returns:
            bytes or bytearray: frame header data
        """

        if self._started is False:
            result = compress_begin(
                self._context,
                block_size=self.block_size,
                block_linked=self.block_linked,
                compression_level=self.compression_level,
                content_checksum=self.content_checksum,
                auto_flush=self.auto_flush,
                return_bytearray=self.return_bytearray,
                source_size=source_size
            )
            self._started = True
            return result
        else:
            raise RuntimeError('compress_begin called when not already initialized')

    def compress(self, data):
        """Compress ``data`` (a ``bytes`` object), returning a bytes object
        containing compressed data the input.

        If ``auto_flush`` has been set to ``False``, some of ``data`` may be
        buffered internally, for use in later calls to compress() and flush().

        The returned data should be concatenated with the output of any
        previous calls to ``compress()`` and a single call to
        ``compress_begin()``.

        Args:
            data (str, bytes or buffer-compatible object): data to compress

        Returns:
            bytes or bytearray: compressed data

        """
        if self._context is None:
            raise RuntimeError('compress called after flush()')

        if self._started is False:
            raise RuntimeError('compress called before compress_begin()')

        result = compress_chunk(
            self._context, data,
            return_bytearray=self.return_bytearray
        )

        return result

    def flush(self):
        """Finish the compression process, returning a bytes object containing any data
        stored in the compressor's internal buffers and a frame footer.

        To use the LZ4FrameCompressor instance after this has been called, it
        is necessary to first call the ``reset()`` method.

        Returns:
            bytes or bytearray: any remaining buffered compressed data and frame footer.

        """
        result = compress_end(
            self._context,
            return_bytearray=self.return_bytearray
        )
        self._context = None
        self._started = False
        return result

    def reset(self):
        """Reset the LZ4FrameCompressor instance (after a call to ``flush``) allowing
        it to be re-used

        """
        self._context = create_compression_context()
        self._started = False


class LZ4FrameDecompressor(object):
    """Create a LZ4 frame decompressor object, which can be used to decompress data
    incrementally.

    Args:
        return_bytearray (bool): When ``False`` a bytes object is returned from the
            calls to methods of this class. When ``True`` a bytearray object will be
            returned. The default is ``False``.
        return_bytes_read (bool): When ``True``, calls to ``decompress`` will return
            the number of bytes read from the input data as well as the uncompressed
            data. The default is ``False``.

    """

    def __init__(self, return_bytearray=False, return_bytes_read=False):
        self.return_bytearray = return_bytearray
        self.return_bytes_read = return_bytes_read
        self._context = create_decompression_context()

    def __enter__(self):
        # All necessary initialization is done in __init__
        return self

    def __exit__(self, exception_type, exception, traceback):
        # The decompression context is created with an appropriate destructor,
        # so no need to del it here
        pass

    def decompress(self, data, full_frame=False):
        """Decompresses part of an LZ4 frame of compressed data. The returned data
        should be concatenated with the output of any previous calls to
        ``decompress()``.

        Args:
            data (str, bytes or buffer-compatible object): data to decompress
            full_frame (bool): If ``True``, then the ``data`` argument is expected
                to contain the full LZ4 frame. Default is ``False``.

        Returns:
            bytes or bytearray: decompressed data
            int: Number of bytes consumed from input data. This is only returned if
                ``return_bytes_read`` is ``True`` when ``LZ4FrameDecompressor`` is
                instantiated.

        """

        if full_frame is True:
            return decompress(
                self._context,
                data,
                return_bytearray=self.return_bytearray,
                return_bytes_read=self.return_bytes_read,
            )
        else:
            return decompress_chunk(
                self._context,
                data,
                return_bytearray=self.return_bytearray,
                return_bytes_read=self.return_bytes_read,
            )
