import sys

if sys.version_info[0] == 2:

    bytes_type = str

    int_type = long

else:

    bytes_type = bytes

    int_type = int

COMPRESSION_RECOMMENDED_INPUT_SIZE = lib.ZSTD_CStreamInSize()

COMPRESSION_RECOMMENDED_OUTPUT_SIZE = lib.ZSTD_CStreamOutSize()

DECOMPRESSION_RECOMMENDED_INPUT_SIZE = lib.ZSTD_DStreamInSize()

DECOMPRESSION_RECOMMENDED_OUTPUT_SIZE = lib.ZSTD_DStreamOutSize()

new_nonzero = ffi.new_allocator(should_clear_after_alloc=False)

MAX_COMPRESSION_LEVEL = lib.ZSTD_maxCLevel()

MAGIC_NUMBER = lib.ZSTD_MAGICNUMBER

FRAME_HEADER = b'\x28\xb5\x2f\xfd'

ZSTD_VERSION = (lib.ZSTD_VERSION_MAJOR, lib.ZSTD_VERSION_MINOR, lib.ZSTD_VERSION_RELEASE)

WINDOWLOG_MIN = lib.ZSTD_WINDOWLOG_MIN

WINDOWLOG_MAX = lib.ZSTD_WINDOWLOG_MAX

CHAINLOG_MIN = lib.ZSTD_CHAINLOG_MIN

CHAINLOG_MAX = lib.ZSTD_CHAINLOG_MAX

HASHLOG_MIN = lib.ZSTD_HASHLOG_MIN

HASHLOG_MAX = lib.ZSTD_HASHLOG_MAX

HASHLOG3_MAX = lib.ZSTD_HASHLOG3_MAX

SEARCHLOG_MIN = lib.ZSTD_SEARCHLOG_MIN

SEARCHLOG_MAX = lib.ZSTD_SEARCHLOG_MAX

SEARCHLENGTH_MIN = lib.ZSTD_SEARCHLENGTH_MIN

SEARCHLENGTH_MAX = lib.ZSTD_SEARCHLENGTH_MAX

TARGETLENGTH_MIN = lib.ZSTD_TARGETLENGTH_MIN

TARGETLENGTH_MAX = lib.ZSTD_TARGETLENGTH_MAX

STRATEGY_FAST = lib.ZSTD_fast

STRATEGY_DFAST = lib.ZSTD_dfast

STRATEGY_GREEDY = lib.ZSTD_greedy

STRATEGY_LAZY = lib.ZSTD_lazy

STRATEGY_LAZY2 = lib.ZSTD_lazy2

STRATEGY_BTLAZY2 = lib.ZSTD_btlazy2

STRATEGY_BTOPT = lib.ZSTD_btopt

COMPRESSOBJ_FLUSH_FINISH = 0

COMPRESSOBJ_FLUSH_BLOCK = 1

class ZstdError(Exception):

    pass

class CompressionParameters(object):

    def __init__(self, window_log, chain_log, hash_log, search_log,

                 search_length, target_length, strategy):

        if window_log < WINDOWLOG_MIN or window_log > WINDOWLOG_MAX:

            raise ValueError('invalid window log value')

        if chain_log < CHAINLOG_MIN or chain_log > CHAINLOG_MAX:

            raise ValueError('invalid chain log value')

        if hash_log < HASHLOG_MIN or hash_log > HASHLOG_MAX:

            raise ValueError('invalid hash log value')

        if search_log < SEARCHLOG_MIN or search_log > SEARCHLOG_MAX:

            raise ValueError('invalid search log value')

        if search_length < SEARCHLENGTH_MIN or search_length > SEARCHLENGTH_MAX:

            raise ValueError('invalid search length value')

        if target_length < TARGETLENGTH_MIN or target_length > TARGETLENGTH_MAX:

            raise ValueError('invalid target length value')

        if strategy < STRATEGY_FAST or strategy > STRATEGY_BTOPT:

            raise ValueError('invalid strategy value')

        self.window_log = window_log

        self.chain_log = chain_log

        self.hash_log = hash_log

        self.search_log = search_log

        self.search_length = search_length

        self.target_length = target_length

        self.strategy = strategy

    def as_compression_parameters(self):

        p = ffi.new('ZSTD_compressionParameters *')[0]

        p.windowLog = self.window_log

        p.chainLog = self.chain_log

        p.hashLog = self.hash_log

        p.searchLog = self.search_log

        p.searchLength = self.search_length

        p.targetLength = self.target_length

        p.strategy = self.strategy

        return p

def get_compression_parameters(level, source_size=0, dict_size=0):

    params = lib.ZSTD_getCParams(level, source_size, dict_size)

    return CompressionParameters(window_log=params.windowLog,

                                 chain_log=params.chainLog,

                                 hash_log=params.hashLog,

                                 search_log=params.searchLog,

                                 search_length=params.searchLength,

                                 target_length=params.targetLength,

                                 strategy=params.strategy)

def estimate_compression_context_size(params):

    if not isinstance(params, CompressionParameters):

        raise ValueError('argument must be a CompressionParameters')

    cparams = params.as_compression_parameters()

    return lib.ZSTD_estimateCCtxSize(cparams)

def estimate_decompression_context_size():

    return lib.ZSTD_estimateDCtxSize()

class ZstdCompressionWriter(object):

    def __init__(self, compressor, writer, source_size, write_size):

        self._compressor = compressor

        self._source_size = source_size

        self._write_size = write_size

        self._entered = False

        if self._entered:

            raise ZstdError('cannot __enter__ multiple times')

        self._cstream = self._compressor._get_cstream(self._source_size)

        self._entered = True

        self._entered = False

            dst_buffer = ffi.new('char[]', self._write_size)

            out_buffer.dst = dst_buffer

            out_buffer.size = self._write_size

                zresult = lib.ZSTD_endStream(self._cstream, out_buffer)

                if lib.ZSTD_isError(zresult):

                    raise ZstdError('error ending compression stream: %s' %

                                    ffi.string(lib.ZSTD_getErrorName(zresult)))

                    self._writer.write(ffi.buffer(out_buffer.dst, out_buffer.pos)[:])

                if zresult == 0:

        self._cstream = None

        self._compressor = None

    def memory_size(self):

        if not self._entered:

            raise ZstdError('cannot determine size of an inactive compressor; '

                            'call when a context manager is active')

        return lib.ZSTD_sizeof_CStream(self._cstream)

        if not self._entered:

            raise ZstdError('write() must be called from an active context '

                            'manager')

        total_write = 0

        data_buffer = ffi.from_buffer(data)

        in_buffer = ffi.new('ZSTD_inBuffer *')

        in_buffer.src = data_buffer

        in_buffer.size = len(data_buffer)

        in_buffer.pos = 0

        dst_buffer = ffi.new('char[]', self._write_size)

        out_buffer.dst = dst_buffer

        out_buffer.size = self._write_size

        out_buffer.pos = 0

        while in_buffer.pos < in_buffer.size:

            zresult = lib.ZSTD_compressStream(self._cstream, out_buffer, in_buffer)

            if lib.ZSTD_isError(zresult):

                raise ZstdError('zstd compress error: %s' %

                                ffi.string(lib.ZSTD_getErrorName(zresult)))

            if out_buffer.pos:

                self._writer.write(ffi.buffer(out_buffer.dst, out_buffer.pos)[:])

                total_write += out_buffer.pos

                out_buffer.pos = 0

        return total_write

    def flush(self):

        if not self._entered:

            raise ZstdError('flush must be called from an active context manager')

        total_write = 0

        out_buffer = ffi.new('ZSTD_outBuffer *')

        dst_buffer = ffi.new('char[]', self._write_size)

        out_buffer.dst = dst_buffer

        out_buffer.size = self._write_size

        while True:

            zresult = lib.ZSTD_flushStream(self._cstream, out_buffer)

            if lib.ZSTD_isError(zresult):

                raise ZstdError('zstd compress error: %s' %

                                ffi.string(lib.ZSTD_getErrorName(zresult)))

            if not out_buffer.pos:

                break

            self._writer.write(ffi.buffer(out_buffer.dst, out_buffer.pos)[:])

            total_write += out_buffer.pos

            out_buffer.pos = 0

        return total_write

class ZstdCompressionObj(object):

    def compress(self, data):

        if self._finished:

            raise ZstdError('cannot call compress() after compressor finished')

        data_buffer = ffi.from_buffer(data)

        source = ffi.new('ZSTD_inBuffer *')

        source.src = data_buffer

        source.size = len(data_buffer)

        source.pos = 0

        chunks = []

        while source.pos < len(data):

            zresult = lib.ZSTD_compressStream(self._cstream, self._out, source)

            if lib.ZSTD_isError(zresult):

                raise ZstdError('zstd compress error: %s' %

                                ffi.string(lib.ZSTD_getErrorName(zresult)))

            if self._out.pos:

                chunks.append(ffi.buffer(self._out.dst, self._out.pos)[:])

                self._out.pos = 0

        return b''.join(chunks)

    def flush(self, flush_mode=COMPRESSOBJ_FLUSH_FINISH):

        if flush_mode not in (COMPRESSOBJ_FLUSH_FINISH, COMPRESSOBJ_FLUSH_BLOCK):

            raise ValueError('flush mode not recognized')

        if self._finished:

            raise ZstdError('compressor object already finished')

        assert self._out.pos == 0

        if flush_mode == COMPRESSOBJ_FLUSH_BLOCK:

            zresult = lib.ZSTD_flushStream(self._cstream, self._out)

            if lib.ZSTD_isError(zresult):

                raise ZstdError('zstd compress error: %s' %

                                ffi.string(lib.ZSTD_getErrorName(zresult)))

            # Output buffer is guaranteed to hold full block.

            assert zresult == 0

            if self._out.pos:

                result = ffi.buffer(self._out.dst, self._out.pos)[:]

                self._out.pos = 0

                return result

            else:

                return b''

        assert flush_mode == COMPRESSOBJ_FLUSH_FINISH

        self._finished = True

        chunks = []

        while True:

            zresult = lib.ZSTD_endStream(self._cstream, self._out)

            if lib.ZSTD_isError(zresult):

                raise ZstdError('error ending compression stream: %s' %

                                ffi.string(lib.ZSTD_getErroName(zresult)))

            if self._out.pos:

                chunks.append(ffi.buffer(self._out.dst, self._out.pos)[:])

                self._out.pos = 0

            if not zresult:

                break

        # GC compression stream immediately.

        self._cstream = None

        return b''.join(chunks)

    def __init__(self, level=3, dict_data=None, compression_params=None,

                 write_checksum=False, write_content_size=False,

                 write_dict_id=True):

        if level < 1:

            raise ValueError('level must be greater than 0')

        elif level > lib.ZSTD_maxCLevel():

            raise ValueError('level must be less than %d' % lib.ZSTD_maxCLevel())

        self._dict_data = dict_data

        self._cparams = compression_params

        self._fparams = ffi.new('ZSTD_frameParameters *')[0]

        self._fparams.checksumFlag = write_checksum

        self._fparams.contentSizeFlag = write_content_size

        self._fparams.noDictIDFlag = not write_dict_id

        cctx = lib.ZSTD_createCCtx()

        if cctx == ffi.NULL:

            raise MemoryError()

        self._cctx = ffi.gc(cctx, lib.ZSTD_freeCCtx)

    def compress(self, data, allow_empty=False):

        if len(data) == 0 and self._fparams.contentSizeFlag and not allow_empty:

            raise ValueError('cannot write empty inputs when writing content sizes')

        # TODO use a CDict for performance.

        dict_data = ffi.NULL

        dict_size = 0

        if self._dict_data:

            dict_data = self._dict_data.as_bytes()

            dict_size = len(self._dict_data)

        params = ffi.new('ZSTD_parameters *')[0]

        if self._cparams:

            params.cParams = self._cparams.as_compression_parameters()

        else:

            params.cParams = lib.ZSTD_getCParams(self._compression_level, len(data),

                                                 dict_size)

        params.fParams = self._fparams

        dest_size = lib.ZSTD_compressBound(len(data))

        out = new_nonzero('char[]', dest_size)

        zresult = lib.ZSTD_compress_advanced(self._cctx,

                                             ffi.addressof(out), dest_size,

                                             data, len(data),

                                             dict_data, dict_size,

                                             params)

        if lib.ZSTD_isError(zresult):

            raise ZstdError('cannot compress: %s' %

                            ffi.string(lib.ZSTD_getErrorName(zresult)))

        return ffi.buffer(out, zresult)[:]

    def compressobj(self, size=0):

        cstream = self._get_cstream(size)