"""This module implements a bloom filter probabilistic data structure and
an a Scalable Bloom Filter that grows in size as your add more items to it
without increasing the false positive error_rate.

Bloom filters that allow asynchronous operation.

Requires the bitarray library: http://pypi.python.org/pypi/bitarray/

origin author: https://github.com/joseph-fox
author: ASXE  https://github.com/asxez

"""

from __future__ import absolute_import

import copy
import hashlib
import math
from struct import calcsize, pack, unpack

import aiofiles
import xxhash

from aiobloom_live.utils import is_string_io

try:
    import bitarray
except ImportError:
    raise ImportError('aiobloom requires bitarray >= 0.3.4')


def make_hashfuncs(num_slices, num_bits):
    if num_bits >= (1 << 31):
        fmt_code, chunk_size = 'Q', 8
    elif num_bits >= (1 << 15):
        fmt_code, chunk_size = 'I', 4
    else:
        fmt_code, chunk_size = 'H', 2
    total_hash_bits = 8 * num_slices * chunk_size
    if total_hash_bits > 384:
        hashfn = hashlib.sha512
    elif total_hash_bits > 256:
        hashfn = hashlib.sha384
    elif total_hash_bits > 160:
        hashfn = hashlib.sha256
    elif total_hash_bits > 128:
        hashfn = hashlib.sha1
    else:
        hashfn = xxhash.xxh128

    fmt = fmt_code * (hashfn().digest_size // chunk_size)
    num_salts, extra = divmod(num_slices, len(fmt))
    if extra:
        num_salts += 1
    salts = tuple(hashfn(hashfn(pack('I', i)).digest()) for i in range(num_salts))

    def _hash_maker(key):
        if isinstance(key, str):
            key_bytes = key.encode('utf-8')
        else:
            key_bytes = str(key).encode('utf-8')
        i = 0
        for salt in salts:
            h = salt.copy()
            h.update(key_bytes)
            for uint in unpack(fmt, h.digest()):
                yield uint % num_bits
                i += 1
                if i >= num_slices:
                    return

    return _hash_maker, hashfn


class BloomFilter(object):
    FILE_FMT = b'<dQQQQ'

    def __init__(self, capacity, error_rate=0.001):
        """Implements a space-efficient probabilistic data structure

        capacity
            this BloomFilter must be able to store at least *capacity* elements
            while maintaining no more than *error_rate* chance of false
            positives
        error_rate
            the error_rate of the filter returning false positives. This
            determines the filters capacity. Inserting more than capacity
            elements greatly increases the chance of false positives.
        """
        if not (0 < error_rate < 1):
            raise ValueError("Error_Rate must be between 0 and 1.")
        if not capacity > 0:
            raise ValueError("Capacity must be > 0.")
        # given M = num_bits, k = num_slices, P = error_rate, n = capacity
        #       k = log2(1/P)
        # solving for m = bits_per_slice
        # n ~= M * ((ln(2) ** 2) / abs(ln(P)))
        # n ~= (k * m) * ((ln(2) ** 2) / abs(ln(P)))
        # m ~= n * abs(ln(P)) / (k * (ln(2) ** 2))
        num_slices = int(math.ceil(math.log(1.0 / error_rate, 2)))
        bits_per_slice = int(math.ceil(
            (capacity * abs(math.log(error_rate))) /
            (num_slices * (math.log(2) ** 2))))
        self._setup(error_rate, num_slices, bits_per_slice, capacity, 0)
        self.bitarray = bitarray.bitarray(self.num_bits, endian='little')
        self.bitarray.setall(False)

    def _setup(self, error_rate, num_slices, bits_per_slice, capacity, count):
        self.error_rate = error_rate
        self.num_slices = num_slices
        self.bits_per_slice = bits_per_slice
        self.capacity = capacity
        self.num_bits = num_slices * bits_per_slice
        self.count = count
        self.make_hashes, self.hashfn = make_hashfuncs(self.num_slices, self.bits_per_slice)

    def __contains__(self, key):
        """Tests a key's membership in this bloom filter.
        """
        bits_per_slice = self.bits_per_slice
        bitarray = self.bitarray
        hashes = self.make_hashes(key)
        offset = 0
        for k in hashes:
            if not bitarray[offset + k]:
                return False
            offset += bits_per_slice
        return True

    def __len__(self):
        """Return the number of keys stored by this bloom filter."""
        return self.count

    def add(self, key, skip_check=False):
        """ Adds a key to this bloom filter. If the key already exists in this
        filter it will return True. Otherwise False.
        """
        bitarray = self.bitarray
        bits_per_slice = self.bits_per_slice
        hashes = self.make_hashes(key)
        found_all_bits = True
        if self.count > self.capacity:
            raise IndexError("BloomFilter is at capacity.")
        offset = 0
        for k in hashes:
            if not skip_check and found_all_bits and not bitarray[offset + k]:
                found_all_bits = False
            self.bitarray[offset + k] = True
            offset += bits_per_slice

        if skip_check:
            self.count += 1
            return False
        elif not found_all_bits:
            self.count += 1
            return False
        else:
            return True

    def copy(self):
        """Return a copy of this bloom filter.
        """
        new_filter = BloomFilter(self.capacity, self.error_rate)
        new_filter.bitarray = self.bitarray.copy()
        return new_filter

    def union(self, other):
        """ Calculates the union of the two underlying bitarrays and returns
        a new bloom filter object."""
        if self.capacity != other.capacity or \
                self.error_rate != other.error_rate:
            raise ValueError(
                "Unioning filters requires both filters to have both the same capacity and error rate.")
        new_bloom = self.copy()
        new_bloom.bitarray = new_bloom.bitarray | other.bitarray
        return new_bloom

    def __or__(self, other):
        return self.union(other)

    def intersection(self, other):
        """ Calculates the intersection of the two underlying bitarrays and returns
        a new bloom filter object."""
        if self.capacity != other.capacity or \
                self.error_rate != other.error_rate:
            raise ValueError(
                "Intersecting filters requires both filters to have equal capacity and error rate.")
        new_bloom = self.copy()
        new_bloom.bitarray = new_bloom.bitarray & other.bitarray
        return new_bloom

    def __and__(self, other):
        return self.intersection(other)

    def tofile(self, f):
        """Write the bloom filter to file object `f'. Underlying bits
        are written as machine values. This is much more space
        efficient than pickling the object."""
        f.write(pack(self.FILE_FMT, self.error_rate, self.num_slices,
                     self.bits_per_slice, self.capacity, self.count))
        (f.write(self.bitarray.tobytes()) if is_string_io(f)
         else self.bitarray.tofile(f))

    @classmethod
    def fromfile(cls, f, n=-1):
        """Read a bloom filter from file-object `f' serialized with
        ``BloomFilter.tofile''. If `n' > 0 read only so many bytes."""
        headerlen = calcsize(cls.FILE_FMT)

        if 0 < n < headerlen:
            raise ValueError('n too small!')

        filter = cls(1)  # Bogus instantiation, we will `_setup'.
        filter._setup(*unpack(cls.FILE_FMT, f.read(headerlen)))
        filter.bitarray = bitarray.bitarray(endian='little')
        if n > 0:
            (filter.bitarray.frombytes(f.read(n - headerlen)) if is_string_io(f)
             else filter.bitarray.fromfile(f, n - headerlen))
        else:
            (filter.bitarray.frombytes(f.read()) if is_string_io(f)
             else filter.bitarray.fromfile(f))
        if filter.num_bits != len(filter.bitarray) and \
                (filter.num_bits + (8 - filter.num_bits % 8) != len(filter.bitarray)):
            raise ValueError('Bit length mismatch!')

        return filter

    # --------------------------- Async ---------------------------
    def to_bytes(self) -> bytes:
        """Serialize this BloomFilter into raw bytes."""
        return pack(self.FILE_FMT, self.error_rate, self.num_slices,
                    self.bits_per_slice, self.capacity, self.count) + self.bitarray.tobytes()

    @classmethod
    def from_bytes(cls, data: bytes) -> "BloomFilter":
        """Deserialize *data* into a BloomFilter instance."""
        header_len = calcsize(cls.FILE_FMT)
        if len(data) < header_len:
            raise ValueError("Data truncated – cannot recover BloomFilter.")
        error_rate, num_slices, bits_per_slice, capacity, count = unpack(cls.FILE_FMT, data[:header_len])
        instance = cls(1)  # temporary values, will be overwritten by _setup
        instance._setup(error_rate, num_slices, bits_per_slice, capacity, count)
        instance.bitarray = bitarray.bitarray(endian='little')
        instance.bitarray.frombytes(data[header_len:])
        return instance

    async def tofile_async(self, file_path: str) -> None:
        """Asynchronously write this filter to *file_path* using asyncio I/O and ``memoryview``"""
        header = pack(self.FILE_FMT, self.error_rate, self.num_slices,
                      self.bits_per_slice, self.capacity, self.count)

        async with aiofiles.open(file_path, "wb") as f:
            await f.write(header)
            await f.write(memoryview(self.bitarray))  # noqa

    @classmethod
    async def fromfile_async(cls, file_path: str) -> "BloomFilter":
        """Asynchronously load a BloomFilter from file path."""
        async with aiofiles.open(file_path, "rb") as f:
            data = await f.read()
        return cls.from_bytes(data)

    def __getstate__(self):
        d = self.__dict__.copy()
        del d['make_hashes']
        return d

    def __setstate__(self, d):
        self.__dict__.update(d)
        self.make_hashes, self.hashfn = make_hashfuncs(self.num_slices, self.bits_per_slice)


class ScalableBloomFilter(object):
    SMALL_SET_GROWTH = 2  # slower, but takes up less memory
    LARGE_SET_GROWTH = 4  # faster, but takes up more memory faster
    FILE_FMT = '<idQd'

    def __init__(self, initial_capacity=100, error_rate=0.001,
                 mode=LARGE_SET_GROWTH):
        """Implements a space-efficient probabilistic data structure that
        grows as more items are added while maintaining a steady false
        positive rate

        initial_capacity
            the initial capacity of the filter
        error_rate
            the error_rate of the filter returning false positives. This
            determines the filters capacity. Going over capacity greatly
            increases the chance of false positives.
        mode
            can be either ScalableBloomFilter.SMALL_SET_GROWTH or
            ScalableBloomFilter.LARGE_SET_GROWTH. SMALL_SET_GROWTH is slower
            but uses less memory. LARGE_SET_GROWTH is faster but consumes
            memory faster.
        """
        if not error_rate or error_rate < 0:
            raise ValueError("Error_Rate must be a decimal less than 0.")
        self._setup(mode, 0.9, initial_capacity, error_rate)
        self.filters = []

    def _setup(self, mode, ratio, initial_capacity, error_rate):
        self.scale = mode
        self.ratio = ratio
        self.initial_capacity = initial_capacity
        self.error_rate = error_rate

    def __contains__(self, key):
        """Tests a key's membership in this bloom filter.
        """
        for f in reversed(self.filters):
            if key in f:
                return True
        return False

    def add(self, key):
        """Adds a key to this bloom filter.
        If the key already exists in this filter it will return True.
        Otherwise False.
        """
        if key in self:
            return True
        if not self.filters:
            filter = BloomFilter(
                capacity=self.initial_capacity,
                error_rate=self.error_rate * self.ratio)
            self.filters.append(filter)
        else:
            filter = self.filters[-1]
            if filter.count >= filter.capacity:
                filter = BloomFilter(
                    capacity=filter.capacity * self.scale,
                    error_rate=filter.error_rate * self.ratio)
                self.filters.append(filter)
        filter.add(key, skip_check=True)
        return False

    def union(self, other):
        """ Calculates the union of the underlying classic bloom filters and returns
        a new scalable bloom filter object."""

        if self.scale != other.scale or \
                self.initial_capacity != other.initial_capacity or \
                self.error_rate != other.error_rate:
            raise ValueError("Unioning two scalable bloom filters requires \
            both filters to have both the same mode, initial capacity and error rate.")
        if len(self.filters) > len(other.filters):
            larger_sbf = copy.deepcopy(self)
            smaller_sbf = other
        else:
            larger_sbf = copy.deepcopy(other)
            smaller_sbf = self
        # Union the underlying classic bloom filters
        new_filters = []
        for i in range(len(smaller_sbf.filters)):
            new_filter = larger_sbf.filters[i] | smaller_sbf.filters[i]
            new_filters.append(new_filter)
        for i in range(len(smaller_sbf.filters), len(larger_sbf.filters)):
            new_filters.append(larger_sbf.filters[i])
        larger_sbf.filters = new_filters
        return larger_sbf

    def __or__(self, other):
        return self.union(other)

    @property
    def capacity(self):
        """Returns the total capacity for all filters in this SBF"""
        return sum(f.capacity for f in self.filters)

    @property
    def count(self):
        return len(self)

    def tofile(self, f):
        """Serialize this ScalableBloomFilter into the file-object
        `f'."""
        f.write(pack(self.FILE_FMT, self.scale, self.ratio,
                     self.initial_capacity, self.error_rate))

        # Write #-of-filters
        f.write(pack(b'<l', len(self.filters)))

        if len(self.filters) > 0:
            # Then each filter directly, with a header describing
            # their lengths.
            headerpos = f.tell()
            headerfmt = b'<' + b'Q' * (len(self.filters))
            f.write(b'.' * calcsize(headerfmt))
            filter_sizes = []
            for filter in self.filters:
                begin = f.tell()
                filter.tofile(f)
                filter_sizes.append(f.tell() - begin)

            f.seek(headerpos)
            f.write(pack(headerfmt, *filter_sizes))

    @classmethod
    def fromfile(cls, f):
        """Deserialize the ScalableBloomFilter in file object `f'."""
        filter = cls()
        filter._setup(*unpack(cls.FILE_FMT, f.read(calcsize(cls.FILE_FMT))))
        nfilters, = unpack(b'<l', f.read(calcsize(b'<l')))
        if nfilters > 0:
            header_fmt = b'<' + b'Q' * nfilters
            bytes = f.read(calcsize(header_fmt))
            filter_lengths = unpack(header_fmt, bytes)
            for fl in filter_lengths:
                filter.filters.append(BloomFilter.fromfile(f, fl))
        else:
            filter.filters = []

        return filter

    # --------------------------- Async ---------------------------
    def to_bytes(self) -> bytes:
        header = pack(self.FILE_FMT, self.scale, self.ratio,
                      self.initial_capacity, self.error_rate)
        num_filters = len(self.filters)
        header += pack(b"<l", num_filters)

        if num_filters:
            filter_bytes_list = [f.to_bytes() for f in self.filters]
            filter_sizes = [len(b) for b in filter_bytes_list]
            header_fmt = b"<" + b"Q" * num_filters
            header += pack(header_fmt, *filter_sizes)
            header += b"".join(filter_bytes_list)
        return header

    @classmethod
    def from_bytes(cls, data: bytes) -> "ScalableBloomFilter":
        """Deserialize *data* into a ScalableBloomFilter instance.

        The implementation operates purely in-memory and performs no blocking
        file I/O, making it suitable for high-concurrency async workloads.
        """
        offset = 0

        header_len = calcsize(cls.FILE_FMT)
        if len(data) < header_len + 4:  # +4 for the nfilters int32
            raise ValueError("Data truncated – cannot recover ScalableBloomFilter.")

        scale, ratio, initial_capacity, error_rate = unpack(
            cls.FILE_FMT, data[offset: offset + header_len]
        )
        offset += header_len

        (nfilters,) = unpack('<l', data[offset: offset + 4])
        offset += 4

        sbf = cls(initial_capacity=int(initial_capacity), error_rate=error_rate, mode=int(scale))
        sbf.ratio = ratio  # preserve original ratio exactly
        sbf.filters = []

        if nfilters:
            header_fmt = '<' + 'Q' * nfilters
            header_size = calcsize(header_fmt)
            if len(data) < offset + header_size:
                raise ValueError("Data truncated – incomplete filter headers.")

            filter_sizes = unpack(header_fmt, data[offset: offset + header_size])
            offset += header_size

            for sz in filter_sizes:
                end = offset + sz
                if len(data) < end:
                    raise ValueError("Data truncated – incomplete nested BloomFilter.")
                bf_bytes = data[offset:end]
                sbf.filters.append(BloomFilter.from_bytes(bf_bytes))
                offset = end
        return sbf

    async def tofile_async(self, file_path: str) -> None:
        """Asynchronously write the SBF to *file_path* using asyncio I/O
        and ``memoryview``"""

        async with aiofiles.open(file_path, "wb") as fp:
            # top-level header
            await fp.write(pack(self.FILE_FMT, self.scale, self.ratio,
                                self.initial_capacity, self.error_rate))

            nfilters = len(self.filters)
            await fp.write(pack('<l', nfilters))

            if nfilters:
                bloom_header_len = calcsize(BloomFilter.FILE_FMT)
                sizes = [bloom_header_len + ((f.num_bits + 7) // 8) for f in self.filters]
                header_fmt = '<' + 'Q' * nfilters
                await fp.write(pack(header_fmt, *sizes))

                for bf in self.filters:
                    header = pack(BloomFilter.FILE_FMT, bf.error_rate, bf.num_slices,
                                  bf.bits_per_slice, bf.capacity, bf.count)
                    await fp.write(header)
                    await fp.write(memoryview(bf.bitarray))

    @classmethod
    async def fromfile_async(cls, file_path: str) -> "ScalableBloomFilter":
        async with aiofiles.open(file_path, "rb") as f:
            data = await f.read()
        return cls.from_bytes(data)

    def __len__(self):
        """Returns the total number of elements stored in this SBF"""
        return sum(f.count for f in self.filters)


if __name__ == "__main__":
    import doctest

    doctest.testmod()
