# Copyright (c) 2016-present, Gregory Szorc

# All rights reserved.

#

# This software may be modified and distributed under the terms

# of the BSD license. See the LICENSE file for details.

"""Python interface to the Zstandard (zstd) compression library."""

from __future__ import absolute_import, unicode_literals

# This should match what the C extension exports.

__all__ = [

    #'BufferSegment',

    #'BufferSegments',

    #'BufferWithSegments',

    #'BufferWithSegmentsCollection',

    'CompressionParameters',

    'ZstdCompressionDict',

    'ZstdCompressionParameters',

    'ZstdCompressor',

    'ZstdError',

    'ZstdDecompressor',

    'FrameParameters',

    'estimate_decompression_context_size',

    'frame_content_size',

    'frame_header_size',

    'get_frame_parameters',

    'train_dictionary',

    # Constants.

    'COMPRESSOBJ_FLUSH_FINISH',

    'COMPRESSOBJ_FLUSH_BLOCK',

    'ZSTD_VERSION',

    'FRAME_HEADER',

    'CONTENTSIZE_UNKNOWN',

    'CONTENTSIZE_ERROR',

    'MAX_COMPRESSION_LEVEL',

    'COMPRESSION_RECOMMENDED_INPUT_SIZE',

    'COMPRESSION_RECOMMENDED_OUTPUT_SIZE',

    'DECOMPRESSION_RECOMMENDED_INPUT_SIZE',

    'DECOMPRESSION_RECOMMENDED_OUTPUT_SIZE',

    'MAGIC_NUMBER',

    'WINDOWLOG_MIN',

    'WINDOWLOG_MAX',

    'CHAINLOG_MIN',

    'CHAINLOG_MAX',

    'HASHLOG_MIN',

    'HASHLOG_MAX',

    'HASHLOG3_MAX',

    'SEARCHLOG_MIN',

    'SEARCHLOG_MAX',

    'SEARCHLENGTH_MIN',

    'SEARCHLENGTH_MAX',

    'TARGETLENGTH_MIN',

    'LDM_MINMATCH_MIN',

    'LDM_MINMATCH_MAX',

    'LDM_BUCKETSIZELOG_MAX',

    'STRATEGY_FAST',

    'STRATEGY_DFAST',

    'STRATEGY_GREEDY',

    'STRATEGY_LAZY',

    'STRATEGY_LAZY2',

    'STRATEGY_BTLAZY2',

    'STRATEGY_BTOPT',

    'STRATEGY_BTULTRA',

    'DICT_TYPE_AUTO',

    'DICT_TYPE_RAWCONTENT',

    'DICT_TYPE_FULLDICT',

    'FORMAT_ZSTD1',

    'FORMAT_ZSTD1_MAGICLESS',

]

import io

import os

import sys

from _zstd_cffi import (

    ffi,

    lib,

)

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'

CONTENTSIZE_UNKNOWN = lib.ZSTD_CONTENTSIZE_UNKNOWN

CONTENTSIZE_ERROR = lib.ZSTD_CONTENTSIZE_ERROR

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

LDM_MINMATCH_MIN = lib.ZSTD_LDM_MINMATCH_MIN

LDM_MINMATCH_MAX = lib.ZSTD_LDM_MINMATCH_MAX

LDM_BUCKETSIZELOG_MAX = lib.ZSTD_LDM_BUCKETSIZELOG_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

STRATEGY_BTULTRA = lib.ZSTD_btultra

DICT_TYPE_AUTO = lib.ZSTD_dct_auto

DICT_TYPE_RAWCONTENT = lib.ZSTD_dct_rawContent

DICT_TYPE_FULLDICT = lib.ZSTD_dct_fullDict

FORMAT_ZSTD1 = lib.ZSTD_f_zstd1

FORMAT_ZSTD1_MAGICLESS = lib.ZSTD_f_zstd1_magicless

COMPRESSOBJ_FLUSH_FINISH = 0

COMPRESSOBJ_FLUSH_BLOCK = 1

def _cpu_count():

    # os.cpu_count() was introducd in Python 3.4.

    try:

        return os.cpu_count() or 0

    except AttributeError:

        pass

    # Linux.

    try:

        if sys.version_info[0] == 2:

            return os.sysconf(b'SC_NPROCESSORS_ONLN')

        else:

            return os.sysconf(u'SC_NPROCESSORS_ONLN')

    except (AttributeError, ValueError):

        pass

    # TODO implement on other platforms.

    return 0

class ZstdError(Exception):

    pass

def _zstd_error(zresult):

    # Resolves to bytes on Python 2 and 3. We use the string for formatting

    # into error messages, which will be literal unicode. So convert it to

    # unicode.

    return ffi.string(lib.ZSTD_getErrorName(zresult)).decode('utf-8')

def _make_cctx_params(params):

    res = lib.ZSTD_createCCtxParams()

    if res == ffi.NULL:

        raise MemoryError()

    res = ffi.gc(res, lib.ZSTD_freeCCtxParams)

    attrs = [

        (lib.ZSTD_p_format, params.format),

        (lib.ZSTD_p_compressionLevel, params.compression_level),

        (lib.ZSTD_p_windowLog, params.window_log),

        (lib.ZSTD_p_hashLog, params.hash_log),

        (lib.ZSTD_p_chainLog, params.chain_log),

        (lib.ZSTD_p_searchLog, params.search_log),

        (lib.ZSTD_p_minMatch, params.min_match),

        (lib.ZSTD_p_targetLength, params.target_length),

        (lib.ZSTD_p_compressionStrategy, params.compression_strategy),

        (lib.ZSTD_p_contentSizeFlag, params.write_content_size),

        (lib.ZSTD_p_checksumFlag, params.write_checksum),

        (lib.ZSTD_p_dictIDFlag, params.write_dict_id),

        (lib.ZSTD_p_nbWorkers, params.threads),

        (lib.ZSTD_p_jobSize, params.job_size),

        (lib.ZSTD_p_overlapSizeLog, params.overlap_size_log),

        (lib.ZSTD_p_compressLiterals, params.compress_literals),

        (lib.ZSTD_p_forceMaxWindow, params.force_max_window),

        (lib.ZSTD_p_enableLongDistanceMatching, params.enable_ldm),

        (lib.ZSTD_p_ldmHashLog, params.ldm_hash_log),

        (lib.ZSTD_p_ldmMinMatch, params.ldm_min_match),

        (lib.ZSTD_p_ldmBucketSizeLog, params.ldm_bucket_size_log),

        (lib.ZSTD_p_ldmHashEveryLog, params.ldm_hash_every_log),

    ]

    for param, value in attrs:

        _set_compression_parameter(res, param, value)

    return res

class ZstdCompressionParameters(object):

    @staticmethod

    def from_level(level, source_size=0, dict_size=0, **kwargs):

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

        args = {

            'window_log': 'windowLog',

            'chain_log': 'chainLog',

            'hash_log': 'hashLog',

            'search_log': 'searchLog',

            'min_match': 'searchLength',

            'target_length': 'targetLength',

            'compression_strategy': 'strategy',

        }

        for arg, attr in args.items():

            if arg not in kwargs:

                kwargs[arg] = getattr(params, attr)

        if 'compress_literals' not in kwargs:

            kwargs['compress_literals'] = 1 if level >= 0 else 0

        return ZstdCompressionParameters(**kwargs)

    def __init__(self, format=0, compression_level=0, window_log=0, hash_log=0,

                 chain_log=0, search_log=0, min_match=0, target_length=0,

                 compression_strategy=0, write_content_size=1, write_checksum=0,

                 write_dict_id=0, job_size=0, overlap_size_log=0,

                 force_max_window=0, enable_ldm=0, ldm_hash_log=0,

                 ldm_min_match=0, ldm_bucket_size_log=0, ldm_hash_every_log=0,

                 threads=0, compress_literals=None):

        if threads < 0:

            threads = _cpu_count()

        if compress_literals is None:

            compress_literals = compression_level >= 0

        self.format = format

        self.compression_level = compression_level

        self.window_log = window_log

        self.hash_log = hash_log

        self.chain_log = chain_log

        self.search_log = search_log

        self.min_match = min_match

        self.target_length = target_length

        self.compression_strategy = compression_strategy

        self.write_content_size = write_content_size

        self.write_checksum = write_checksum

        self.write_dict_id = write_dict_id

        self.job_size = job_size

        self.overlap_size_log = overlap_size_log

        self.compress_literals = compress_literals

        self.force_max_window = force_max_window

        self.enable_ldm = enable_ldm

        self.ldm_hash_log = ldm_hash_log

        self.ldm_min_match = ldm_min_match

        self.ldm_bucket_size_log = ldm_bucket_size_log

        self.ldm_hash_every_log = ldm_hash_every_log

        self.threads = threads

        self.params = _make_cctx_params(self)

    def estimated_compression_context_size(self):

        return lib.ZSTD_estimateCCtxSize_usingCCtxParams(self.params)

CompressionParameters = ZstdCompressionParameters

def estimate_decompression_context_size():

    return lib.ZSTD_estimateDCtxSize()

def _set_compression_parameter(params, param, value):

    zresult = lib.ZSTD_CCtxParam_setParameter(params, param,

                                              ffi.cast('unsigned', value))

    if lib.ZSTD_isError(zresult):

        raise ZstdError('unable to set compression context parameter: %s' %

                        _zstd_error(zresult))

class ZstdCompressionWriter(object):

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

        self._compressor = compressor

        self._writer = writer

        self._source_size = source_size

        self._write_size = write_size

        self._entered = False

        self._bytes_compressed = 0

    def __enter__(self):

        if self._entered:

            raise ZstdError('cannot __enter__ multiple times')

        zresult = lib.ZSTD_CCtx_setPledgedSrcSize(self._compressor._cctx,

                                                  self._source_size)

        if lib.ZSTD_isError(zresult):

            raise ZstdError('error setting source size: %s' %

                            _zstd_error(zresult))

        self._entered = True

        return self

    def __exit__(self, exc_type, exc_value, exc_tb):

        self._entered = False

        if not exc_type and not exc_value and not exc_tb:

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

            out_buffer = ffi.new('ZSTD_outBuffer *')

            in_buffer = ffi.new('ZSTD_inBuffer *')

            out_buffer.dst = dst_buffer

            out_buffer.size = len(dst_buffer)

            out_buffer.pos = 0

            in_buffer.src = ffi.NULL

            in_buffer.size = 0

            in_buffer.pos = 0

            while True:

                zresult = lib.ZSTD_compress_generic(self._compressor._cctx,

                                                    out_buffer, in_buffer,

                                                    lib.ZSTD_e_end)

                if lib.ZSTD_isError(zresult):

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

                                    _zstd_error(zresult))

                if out_buffer.pos:

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

                    out_buffer.pos = 0

                if zresult == 0:

                    break

        self._compressor = None

        return False

    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_CCtx(self._compressor._cctx)

    def write(self, data):

        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

        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

        out_buffer.pos = 0

        while in_buffer.pos < in_buffer.size:

            zresult = lib.ZSTD_compress_generic(self._compressor._cctx,

                                                out_buffer, in_buffer,

                                                lib.ZSTD_e_continue)

            if lib.ZSTD_isError(zresult):

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

                                _zstd_error(zresult))

            if out_buffer.pos:

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

                total_write += out_buffer.pos

                self._bytes_compressed += out_buffer.pos

                out_buffer.pos = 0

        return total_write

    def flush(self):

        if not self._entered: