# 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.

    'FLUSH_BLOCK',

    'FLUSH_FRAME',

    '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',

    'BLOCKSIZELOG_MAX',

    'BLOCKSIZE_MAX',

    'WINDOWLOG_MIN',

    'WINDOWLOG_MAX',

    'CHAINLOG_MIN',

    'CHAINLOG_MAX',

    'HASHLOG_MIN',

    'HASHLOG_MAX',

    'HASHLOG3_MAX',

    'MINMATCH_MIN',

    'MINMATCH_MAX',

    'SEARCHLOG_MIN',

    'SEARCHLOG_MAX',

    'SEARCHLENGTH_MIN',

    'SEARCHLENGTH_MAX',

    'TARGETLENGTH_MIN',

    'TARGETLENGTH_MAX',

    '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',

    'STRATEGY_BTULTRA2',

    '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)

BLOCKSIZELOG_MAX = lib.ZSTD_BLOCKSIZELOG_MAX

BLOCKSIZE_MAX = lib.ZSTD_BLOCKSIZE_MAX

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

MINMATCH_MIN = lib.ZSTD_MINMATCH_MIN

MINMATCH_MAX = lib.ZSTD_MINMATCH_MAX

SEARCHLOG_MIN = lib.ZSTD_SEARCHLOG_MIN

SEARCHLOG_MAX = lib.ZSTD_SEARCHLOG_MAX

SEARCHLENGTH_MIN = lib.ZSTD_MINMATCH_MIN

SEARCHLENGTH_MAX = lib.ZSTD_MINMATCH_MAX

TARGETLENGTH_MIN = lib.ZSTD_TARGETLENGTH_MIN

TARGETLENGTH_MAX = lib.ZSTD_TARGETLENGTH_MAX

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

STRATEGY_BTULTRA2 = lib.ZSTD_btultra2

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

FLUSH_BLOCK = 0

FLUSH_FRAME = 1

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_c_format, params.format),

        (lib.ZSTD_c_compressionLevel, params.compression_level),

        (lib.ZSTD_c_windowLog, params.window_log),

        (lib.ZSTD_c_hashLog, params.hash_log),

        (lib.ZSTD_c_chainLog, params.chain_log),

        (lib.ZSTD_c_searchLog, params.search_log),

        (lib.ZSTD_c_minMatch, params.min_match),

        (lib.ZSTD_c_targetLength, params.target_length),

        (lib.ZSTD_c_strategy, params.compression_strategy),

        (lib.ZSTD_c_contentSizeFlag, params.write_content_size),

        (lib.ZSTD_c_checksumFlag, params.write_checksum),

        (lib.ZSTD_c_dictIDFlag, params.write_dict_id),

        (lib.ZSTD_c_nbWorkers, params.threads),

        (lib.ZSTD_c_jobSize, params.job_size),

        (lib.ZSTD_c_overlapLog, params.overlap_log),

        (lib.ZSTD_c_forceMaxWindow, params.force_max_window),

        (lib.ZSTD_c_enableLongDistanceMatching, params.enable_ldm),

        (lib.ZSTD_c_ldmHashLog, params.ldm_hash_log),

        (lib.ZSTD_c_ldmMinMatch, params.ldm_min_match),

        (lib.ZSTD_c_ldmBucketSizeLog, params.ldm_bucket_size_log),

        (lib.ZSTD_c_ldmHashRateLog, params.ldm_hash_rate_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': 'minMatch',

            'target_length': 'targetLength',

            'compression_strategy': 'strategy',

        }

        for arg, attr in args.items():

            if arg not in kwargs:

                kwargs[arg] = getattr(params, attr)

        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,

                 strategy=-1, compression_strategy=-1,

                 write_content_size=1, write_checksum=0,

                 write_dict_id=0, job_size=0, overlap_log=-1,

                 overlap_size_log=-1, force_max_window=0, enable_ldm=0,

                 ldm_hash_log=0, ldm_min_match=0, ldm_bucket_size_log=0,

                 ldm_hash_rate_log=-1, ldm_hash_every_log=-1, threads=0):

        params = lib.ZSTD_createCCtxParams()

        if params == ffi.NULL:

            raise MemoryError()

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

        self._params = params

        if threads < 0:

            threads = _cpu_count()

        # We need to set ZSTD_c_nbWorkers before ZSTD_c_jobSize and ZSTD_c_overlapLog

        # because setting ZSTD_c_nbWorkers resets the other parameters.

        _set_compression_parameter(params, lib.ZSTD_c_nbWorkers, threads)

        _set_compression_parameter(params, lib.ZSTD_c_format, format)

        _set_compression_parameter(params, lib.ZSTD_c_compressionLevel, compression_level)

        _set_compression_parameter(params, lib.ZSTD_c_windowLog, window_log)

        _set_compression_parameter(params, lib.ZSTD_c_hashLog, hash_log)

        _set_compression_parameter(params, lib.ZSTD_c_chainLog, chain_log)

        _set_compression_parameter(params, lib.ZSTD_c_searchLog, search_log)

        _set_compression_parameter(params, lib.ZSTD_c_minMatch, min_match)

        _set_compression_parameter(params, lib.ZSTD_c_targetLength, target_length)

        if strategy != -1 and compression_strategy != -1:

            raise ValueError('cannot specify both compression_strategy and strategy')

        if compression_strategy != -1:

            strategy = compression_strategy

        elif strategy == -1:

            strategy = 0

        _set_compression_parameter(params, lib.ZSTD_c_strategy, strategy)

        _set_compression_parameter(params, lib.ZSTD_c_contentSizeFlag, write_content_size)

        _set_compression_parameter(params, lib.ZSTD_c_checksumFlag, write_checksum)

        _set_compression_parameter(params, lib.ZSTD_c_dictIDFlag, write_dict_id)

        _set_compression_parameter(params, lib.ZSTD_c_jobSize, job_size)

        if overlap_log != -1 and overlap_size_log != -1:

            raise ValueError('cannot specify both overlap_log and overlap_size_log')

        if overlap_size_log != -1:

            overlap_log = overlap_size_log

        elif overlap_log == -1:

            overlap_log = 0

        _set_compression_parameter(params, lib.ZSTD_c_overlapLog, overlap_log)

        _set_compression_parameter(params, lib.ZSTD_c_forceMaxWindow, force_max_window)

        _set_compression_parameter(params, lib.ZSTD_c_enableLongDistanceMatching, enable_ldm)

        _set_compression_parameter(params, lib.ZSTD_c_ldmHashLog, ldm_hash_log)

        _set_compression_parameter(params, lib.ZSTD_c_ldmMinMatch, ldm_min_match)

        _set_compression_parameter(params, lib.ZSTD_c_ldmBucketSizeLog, ldm_bucket_size_log)

        if ldm_hash_rate_log != -1 and ldm_hash_every_log != -1:

            raise ValueError('cannot specify both ldm_hash_rate_log and ldm_hash_every_log')

        if ldm_hash_every_log != -1:

            ldm_hash_rate_log = ldm_hash_every_log

        elif ldm_hash_rate_log == -1:

            ldm_hash_rate_log = 0

        _set_compression_parameter(params, lib.ZSTD_c_ldmHashRateLog, ldm_hash_rate_log)

    @property

    def format(self):

        return _get_compression_parameter(self._params, lib.ZSTD_c_format)

    @property

    def compression_level(self):

        return _get_compression_parameter(self._params, lib.ZSTD_c_compressionLevel)

    @property

    def window_log(self):

        return _get_compression_parameter(self._params, lib.ZSTD_c_windowLog)

    @property

    def hash_log(self):

        return _get_compression_parameter(self._params, lib.ZSTD_c_hashLog)

    @property

    def chain_log(self):

        return _get_compression_parameter(self._params, lib.ZSTD_c_chainLog)

    @property

    def search_log(self):

        return _get_compression_parameter(self._params, lib.ZSTD_c_searchLog)

    @property

    def min_match(self):

        return _get_compression_parameter(self._params, lib.ZSTD_c_minMatch)

    @property

    def target_length(self):

        return _get_compression_parameter(self._params, lib.ZSTD_c_targetLength)

    @property

    def compression_strategy(self):

        return _get_compression_parameter(self._params, lib.ZSTD_c_strategy)

    @property

    def write_content_size(self):

        return _get_compression_parameter(self._params, lib.ZSTD_c_contentSizeFlag)

    @property

    def write_checksum(self):

        return _get_compression_parameter(self._params, lib.ZSTD_c_checksumFlag)

    @property

    def write_dict_id(self):

        return _get_compression_parameter(self._params, lib.ZSTD_c_dictIDFlag)

    @property

    def job_size(self):

        return _get_compression_parameter(self._params, lib.ZSTD_c_jobSize)

    @property

    def overlap_log(self):

        return _get_compression_parameter(self._params, lib.ZSTD_c_overlapLog)

    @property

    def overlap_size_log(self):

        return self.overlap_log

    @property

    def force_max_window(self):

        return _get_compression_parameter(self._params, lib.ZSTD_c_forceMaxWindow)

    @property

    def enable_ldm(self):

        return _get_compression_parameter(self._params, lib.ZSTD_c_enableLongDistanceMatching)

    @property

    def ldm_hash_log(self):

        return _get_compression_parameter(self._params, lib.ZSTD_c_ldmHashLog)

    @property

    def ldm_min_match(self):

        return _get_compression_parameter(self._params, lib.ZSTD_c_ldmMinMatch)

    @property

    def ldm_bucket_size_log(self):

        return _get_compression_parameter(self._params, lib.ZSTD_c_ldmBucketSizeLog)

    @property

    def ldm_hash_rate_log(self):

        return _get_compression_parameter(self._params, lib.ZSTD_c_ldmHashRateLog)

    @property

    def ldm_hash_every_log(self):

        return self.ldm_hash_rate_log

    @property

    def threads(self):

        return _get_compression_parameter(self._params, lib.ZSTD_c_nbWorkers)

    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_CCtxParams_setParameter(params, param, value)

    if lib.ZSTD_isError(zresult):

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

                        _zstd_error(zresult))

def _get_compression_parameter(params, param):

    result = ffi.new('int *')

    zresult = lib.ZSTD_CCtxParams_getParameter(params, param, result)

    if lib.ZSTD_isError(zresult):

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

                        _zstd_error(zresult))

    return result[0]

class ZstdCompressionWriter(object):

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

                 write_return_read):

        self._compressor = compressor

        self._writer = writer

        self._write_size = write_size

        self._write_return_read = bool(write_return_read)

        self._entered = False

        self._closed = False

        self._bytes_compressed = 0

        self._dst_buffer = ffi.new('char[]', write_size)

        self._out_buffer = ffi.new('ZSTD_outBuffer *')

        self._out_buffer.dst = self._dst_buffer

        self._out_buffer.size = len(self._dst_buffer)

        self._out_buffer.pos = 0

        zresult = lib.ZSTD_CCtx_setPledgedSrcSize(compressor._cctx,

                                                  source_size)

        if lib.ZSTD_isError(zresult):

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

                            _zstd_error(zresult))

    def __enter__(self):

        if self._closed:

            raise ValueError('stream is closed')

        if self._entered:

            raise ZstdError('cannot __enter__ multiple times')

        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:

            self.close()

        self._compressor = None

        return False

    def memory_size(self):

        return lib.ZSTD_sizeof_CCtx(self._compressor._cctx)

    def fileno(self):

        f = getattr(self._writer, 'fileno', None)

        if f:

            return f()

        else:

            raise OSError('fileno not available on underlying writer')

    def close(self):

        if self._closed:

            return