# 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

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'

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

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

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

        zresult = lib.ZSTD_checkCParams(self.as_compression_parameters())

        if lib.ZSTD_isError(zresult):

            raise ValueError('invalid compression parameters: %s',

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

    def estimated_compression_context_size(self):

        return lib.ZSTD_estimateCCtxSize(self.as_compression_parameters())

    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._writer = writer

        self._source_size = source_size

        self._write_size = write_size

        self._entered = False

        self._mtcctx = compressor._cctx if compressor._multithreaded else None

    def __enter__(self):

        if self._entered:

            raise ZstdError('cannot __enter__ multiple times')

        if self._mtcctx:

            self._compressor._init_mtcstream(self._source_size)

        else:

            self._compressor._ensure_cstream(self._source_size)

        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:

            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 True:

                if self._mtcctx:

                    zresult = lib.ZSTDMT_endStream(self._mtcctx, out_buffer)

                else:

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

                if lib.ZSTD_isError(zresult):

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

                                    ffi.string(lib.ZSTD_getErrorName(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_CStream(self._compressor._cstream)

    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:

            if self._mtcctx:

                zresult = lib.ZSTDMT_compressStream(self._mtcctx, out_buffer,

                                                    in_buffer)

            else:

                zresult = lib.ZSTD_compressStream(self._compressor._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

        out_buffer.pos = 0

        while True:

            if self._mtcctx:

                zresult = lib.ZSTDMT_flushStream(self._mtcctx, out_buffer)

            else:

                zresult = lib.ZSTD_flushStream(self._compressor._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):

            if self._mtcctx:

                zresult = lib.ZSTDMT_compressStream(self._mtcctx,

                                                    self._out, source)

            else:

                zresult = lib.ZSTD_compressStream(self._compressor._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:

            if self._mtcctx:

                zresult = lib.ZSTDMT_flushStream(self._mtcctx, self._out)

            else:

                zresult = lib.ZSTD_flushStream(self._compressor._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:

            if self._mtcctx:

                zresult = lib.ZSTDMT_endStream(self._mtcctx, self._out)

            else:

                zresult = lib.ZSTD_endStream(self._compressor._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

        return b''.join(chunks)

class ZstdCompressor(object):

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

                 write_checksum=False, write_content_size=False,

                 write_dict_id=True, threads=0):

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