"""
A pure python version of the hdr_histogram code

Ported from
https://github.com/HdrHistogram/HdrHistogram (Java)
https://github.com/HdrHistogram/HdrHistogram_c (C)

Coding / Decoding of histograms using the
HdrHistogram V1 compressed transportable format

Written by Alec Hothan

Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at

    http://www.apache.org/licenses/LICENSE-2.0

Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
"""
from __future__ import print_function

import math

from builtins import str
from builtins import range
from builtins import object

import base64
import ctypes
from ctypes import BigEndianStructure
from ctypes import addressof
from ctypes import c_byte
from ctypes import c_uint
from ctypes import c_longlong
from ctypes import c_ulonglong

from pyhdrh import add_array
from pyhdrh import sub_array
from pyhdrh import decode
from pyhdrh import encode
import zlib

XTXV1_ENCODING_COOKIE_BASE = 0x1c84930a
XTXV1_COMPRESSION_COOKIE_BASE = 0x1c84930b

# LEB128 + ZigZag require up to 9 bytes per word
XTXV1_MAX_WORD_SIZE_IN_BYTES = 9

# allow at most 4 MB counter array size
# that represents 500,000 8 byte counters
MAX_COUNTS_SIZE = 4 * 1024 * 1024


def get_cookie_base(cookie):
    return cookie & ~0xf0


def get_encoding_cookie():
    # LSBit of wordsize byte indicates TLZE Encoding
    return XTXV1_ENCODING_COOKIE_BASE | 0x10


def get_compression_cookie():
    # LSBit of wordsize byte indicates TLZE Encoding
    return XTXV1_COMPRESSION_COOKIE_BASE | 0x10


class HdrCookieException(Exception):
    pass


class HdrLengthException(Exception):
    pass


class HdrHistogramSettingsException(Exception):
    pass


# External encoding header
class ExternalHeader(BigEndianStructure):
    _pack_ = 1
    _fields_ = [
        ("cookie", c_uint),
        ("length", c_uint)]
ext_header_size = ctypes.sizeof(ExternalHeader)


# Header for the zlib compressed part
class PayloadHeader(BigEndianStructure):
    _pack_ = 1
    _fields_ = [
        ("cookie", c_uint),
        ("payload_len", c_uint),
        ("normalizing_index_offset", c_uint),
        ("significant_figures", c_uint),
        ("lowest_trackable_value", c_ulonglong),
        ("highest_trackable_value", c_ulonglong),
        ("conversion_ratio_bits", c_ulonglong)]
payload_header_size = ctypes.sizeof(PayloadHeader)

class HdrPayload(object):
    """A class that wraps the ctypes big endian struct that will hold the
    histogram wire format content (including the counters).
    """
    def __init__(self, counts_len=0, compressed_payload=None):
        """Two ways to use this class:
        - for an empty histogram (pass counts_len>0 and compressed_payload=None)
        - for a decoded histogram (counts_len=0 and compressed_payload!=None)

        Params:
            counts_len number of counters to allocate
                ignored if a compressed payload is provided (not None)

            compressed_payload (string) a payload in zlib compressed form,
                decompress and decode the payload header.
                Caller must then invoke init_counts to pass in counts_len so that the
                counts array can be updated from the decoded varint buffer
                None if no compressed payload is to be associated to this instance
        """
        self.counts_len = counts_len
        self._data = None
        try:
            # ctype counter type
            self.counter_ctype = c_longlong
        except IndexError:
            raise ValueError('Invalid word size')
        if not self.counter_ctype:
            raise ValueError('Invalid word size')
        if compressed_payload:
            self._decompress(compressed_payload)
        elif counts_len:
            self.payload = PayloadHeader()
            self.payload.cookie = get_encoding_cookie()
            self._init_counts()
        else:
            raise RuntimeError('counts_len cannot be zero')

    def _init_counts(self):
        self.counts = (self.counter_ctype * self.counts_len)()

    def init_counts(self, counts_len):
        """Called after instantiating with a compressed payload
        Params:
            counts_len counts size to use based on decoded settings in the header
        """
        assert(self._data and counts_len and self.counts_len == 0)
        self.counts_len = counts_len
        self._init_counts()

        results = decode(self._data, payload_header_size, addressof(self.counts), counts_len)
        # no longer needed
        self._data = None
        return results

    def get_counts(self):
        return self.counts

    def _decompress(self, compressed_payload):
        """Decompress a compressed payload into this payload wrapper.
        Note that the decompressed buffer is saved in self._data and the
        counts array is not yet allocated.

        Args:
            compressed_payload (string) a payload in zlib compressed form
        Exception:
            HdrCookieException:
                the compressed payload has an invalid cookie
            HdrLengthException:
                the decompressed size is too small for the HdrPayload structure
                or is not aligned or is too large for the passed payload class
            HdrHistogramSettingsException:
                mismatch in the significant figures, lowest and highest
                         trackable value
        """
        # make sure this instance is pristine
        if self._data:
            raise RuntimeError('Cannot decompress to an instance with payload')
        # Here it is important to keep a reference to the decompressed
        # string so that it does not get garbage collected
        self._data = zlib.decompress(compressed_payload)
        len_data = len(self._data)

        counts_size = len_data - payload_header_size
        if payload_header_size > counts_size > MAX_COUNTS_SIZE:
            raise HdrLengthException('Invalid size:' + str(len_data))

        # copy the first bytes for the header
        self.payload = PayloadHeader.from_buffer_copy(self._data)

        cookie = self.payload.cookie
        if get_cookie_base(cookie) != XTXV1_ENCODING_COOKIE_BASE:
            raise HdrCookieException('Invalid cookie: %x' % cookie)

    def compress(self, counts_limit):
        """Compress this payload instance
        Args:
            counts_limit how many counters should be encoded
                           starting from index 0 (can be 0),
        Return:
            the compressed payload (python string)
        """
        if self.payload:
            # worst case varint encoded length is when each counter is at the maximum value
            # in this case 1 more byte per counter is needed due to the more bits
            varint_len = counts_limit * XTXV1_MAX_WORD_SIZE_IN_BYTES
            # Add on space for the non zero value bits
            varint_len += int(math.ceil(self.counts_len/8))
            # allocate enough space to fit the header and the varint string
            encode_buf = (c_byte * (payload_header_size + varint_len))()

            # encode past the payload header
            varint_len = encode(addressof(self.counts), counts_limit, self.counts_len,
                                addressof(encode_buf) + payload_header_size,
                                varint_len)

            # copy the header after updating the varint stream length
            self.payload.payload_len = varint_len
            ctypes.memmove(addressof(encode_buf), addressof(self.payload), payload_header_size)

            cdata = zlib.compress(ctypes.string_at(encode_buf, payload_header_size + varint_len))
            return cdata
        # can't compress if no payload
        raise RuntimeError('No payload to compress')

    def dump(self, label=None):
        if label: