"""
RBSP Bloom Filter Layer.

Provides O(1) probabilistic set membership with configurable false-positive
rates.  Two variants are included:

    - ``BloomFilter``         Standard space-efficient bloom filter.
    - ``CountingBloomFilter`` Supports deletion by maintaining per-slot counters.

Hash functions are derived from the RBSP hash core using double-hashing
(Kirschner & Mitzenmacher technique) rather than relying on SHA/MD5.
"""

from __future__ import annotations

import array
import math
import mmap
import os
import struct
import tempfile
from pathlib import Path
from typing import Optional, Union

_MASK64 = 0xFFFFFFFFFFFFFFFF

# RBSP-specific mixing constants (same family as signature.py)
_SEED_A = 0xA5B9C3D7E1F20463
_SEED_B = 0x3C6EF372FE94F82B


def _rbsp_hash_pair(data: bytes) -> tuple[int, int]:
    """Compute two independent 64-bit hashes of *data* using the RBSP core.

    Returns (h1, h2) which are then combined via double-hashing to produce
    an arbitrary number of hash functions:  ``h_i(x) = (h1 + i*h2) % m``.
    """
    h1 = _SEED_A
    h2 = _SEED_B

    for b in data:
        h1 ^= b
        h1 = (h1 * 0x9E3779B97F4A7C15) & _MASK64
        h1 = ((h1 << 31) | (h1 >> 33)) & _MASK64

        h2 ^= b
        h2 = (h2 * 0xBF58476D1CE4E5B9) & _MASK64
        h2 = ((h2 << 27) | (h2 >> 37)) & _MASK64

    # Finalise
    h1 ^= h1 >> 30
    h1 = (h1 * 0xBF58476D1CE4E5B9) & _MASK64
    h1 ^= h1 >> 27

    h2 ^= h2 >> 30
    h2 = (h2 * 0x94D049BB133111EB) & _MASK64
    h2 ^= h2 >> 27

    return h1, h2


def _optimal_params(
    expected_items: int, fp_rate: float
) -> tuple[int, int]:
    """Compute optimal (num_bits, num_hashes) for a standard Bloom filter.

    Args:
        expected_items: Number of elements expected to be inserted.
        fp_rate:        Desired false-positive probability (e.g. 0.01 for 1 %).

    Returns:
        (num_bits, num_hashes) tuple.
    """
    if expected_items <= 0:
        expected_items = 1
    if fp_rate <= 0 or fp_rate >= 1:
        raise ValueError("fp_rate must be in (0, 1)")

    num_bits = int(-expected_items * math.log(fp_rate) / (math.log(2) ** 2))
    num_bits = max(num_bits, 64)  # floor at 8 bytes
    num_hashes = max(1, int((num_bits / expected_items) * math.log(2)))
    return num_bits, num_hashes


def _to_bytes(item: Union[str, bytes, int]) -> bytes:
    """Normalise an item to bytes for hashing."""
    if isinstance(item, bytes):
        return item
    if isinstance(item, str):
        return item.encode("utf-8")
    if isinstance(item, int):
        return item.to_bytes(8, "big", signed=True)
    raise TypeError(f"Unsupported type: {type(item)}")


# ---------------------------------------------------------------------------
# Standard Bloom Filter
# ---------------------------------------------------------------------------

class BloomFilter:
    """Space-efficient probabilistic set with O(1) membership test.

    Uses a bit array and *k* hash functions derived from the RBSP hash via
    double-hashing.

    Args:
        expected_items: Anticipated number of insertions.
        fp_rate:        Target false-positive rate (default 0.01).
    """

    def __init__(
        self,
        expected_items: int = 10_000,
        fp_rate: float = 0.01,
    ) -> None:
        self._expected = expected_items
        self._fp_rate = fp_rate
        self._num_bits, self._num_hashes = _optimal_params(expected_items, fp_rate)
        self._bits = bytearray(math.ceil(self._num_bits / 8))
        self._count = 0

    # -- Properties -------------------------------------------------------

    @property
    def num_bits(self) -> int:
        """Total bits in the filter."""
        return self._num_bits

    @property
    def num_hashes(self) -> int:
        """Number of hash functions."""
        return self._num_hashes

    @property
    def count(self) -> int:
        """Number of items inserted (not accounting for duplicates)."""
        return self._count

    @property
    def estimated_fp_rate(self) -> float:
        """Current estimated false-positive rate given the insertion count."""
        if self._count == 0:
            return 0.0
        exponent = -self._num_hashes * self._count / self._num_bits
        return (1 - math.exp(exponent)) ** self._num_hashes

    # -- Bit operations ---------------------------------------------------

    def _set_bit(self, pos: int) -> None:
        byte_idx, bit_idx = divmod(pos, 8)
        self._bits[byte_idx] |= 1 << bit_idx

    def _get_bit(self, pos: int) -> bool:
        byte_idx, bit_idx = divmod(pos, 8)
        return bool(self._bits[byte_idx] & (1 << bit_idx))

    def _hash_positions(self, data: bytes) -> list[int]:
        """Return the *k* bit positions for *data* using double-hashing."""
        h1, h2 = _rbsp_hash_pair(data)
        return [((h1 + i * h2) & _MASK64) % self._num_bits for i in range(self._num_hashes)]

    # -- Public API -------------------------------------------------------

    def add(self, item: Union[str, bytes, int]) -> None:
        """Add *item* to the filter."""
        data = _to_bytes(item)
        for pos in self._hash_positions(data):
            self._set_bit(pos)
        self._count += 1

    def __contains__(self, item: Union[str, bytes, int]) -> bool:
        """Test whether *item* might be in the filter (O(1)).

        Returns ``True`` if the item *may* be present (with false-positive
        probability <= *fp_rate*), or ``False`` if it is **definitely** not
        present.
        """
        data = _to_bytes(item)
        return all(self._get_bit(pos) for pos in self._hash_positions(data))

    def __len__(self) -> int:
        return self._count

    def clear(self) -> None:
        """Reset the filter to empty."""
        self._bits = bytearray(math.ceil(self._num_bits / 8))
        self._count = 0

    def union(self, other: "BloomFilter") -> "BloomFilter":
        """Return a new BloomFilter that is the union of *self* and *other*.

        Both filters must have the same parameters.

        Raises:
            ValueError: If the filters have incompatible parameters.
        """
        if self._num_bits != other._num_bits or self._num_hashes != other._num_hashes:
            raise ValueError("Cannot union bloom filters with different parameters")
        result = BloomFilter.__new__(BloomFilter)
        result._expected = self._expected
        result._fp_rate = self._fp_rate
        result._num_bits = self._num_bits
        result._num_hashes = self._num_hashes
        result._bits = bytearray(a | b for a, b in zip(self._bits, other._bits))
        result._count = self._count + other._count  # upper bound
        return result

    # -- Serialization ----------------------------------------------------

    def to_bytes(self) -> bytes:
        """Serialize the entire filter to bytes (for persistence)."""
        header = struct.pack(">IIQ", self._num_bits, self._num_hashes, self._count)
        return header + bytes(self._bits)

    @classmethod
    def from_bytes(cls, data: bytes) -> "BloomFilter":
        """Reconstruct a BloomFilter from serialized bytes."""
        if len(data) < 16:
            raise ValueError("Data too short for bloom filter header")
        num_bits, num_hashes, count = struct.unpack(">IIQ", data[:16])
        bf = cls.__new__(cls)
        bf._num_bits = num_bits
        bf._num_hashes = num_hashes
        bf._count = count
        bf._expected = 0
        bf._fp_rate = 0.0
        bf._bits = bytearray(data[16:])
        expected_bytes = math.ceil(num_bits / 8)
        if len(bf._bits) < expected_bytes:
            raise ValueError("Data too short for bit array")
        return bf


# ---------------------------------------------------------------------------
# Counting Bloom Filter (supports deletion)
# ---------------------------------------------------------------------------

class CountingBloomFilter:
    """Bloom filter variant that supports deletion via per-slot counters.

    Each slot uses a 4-bit counter (0-15), allowing up to 15 insertions of
    items that hash to the same slot before saturation.  Two slots are packed
    per byte for memory efficiency.

    Args:
        expected_items: Anticipated number of insertions.
        fp_rate:        Target false-positive rate (default 0.01).
    """

    _MAX_COUNT = 15  # 4-bit counter ceiling

    def __init__(
        self,
        expected_items: int = 10_000,
        fp_rate: float = 0.01,
    ) -> None:
        self._expected = expected_items
        self._fp_rate = fp_rate
        self._num_slots, self._num_hashes = _optimal_params(expected_items, fp_rate)
        # Pack two 4-bit counters per byte
        self._counters = bytearray(math.ceil(self._num_slots / 2))
        self._count = 0

    @property
    def num_slots(self) -> int:
        return self._num_slots

    @property
    def num_hashes(self) -> int:
        return self._num_hashes

    @property
    def count(self) -> int:
        return self._count

    def _get_counter(self, pos: int) -> int:
        byte_idx, nibble = divmod(pos, 2)
        if nibble == 0:
            return self._counters[byte_idx] & 0x0F
        return (self._counters[byte_idx] >> 4) & 0x0F

    def _inc_counter(self, pos: int) -> None:
        byte_idx, nibble = divmod(pos, 2)
        current = self._get_counter(pos)
        if current < self._MAX_COUNT:
            if nibble == 0:
                self._counters[byte_idx] = (self._counters[byte_idx] & 0xF0) | (current + 1)
            else:
                self._counters[byte_idx] = (self._counters[byte_idx] & 0x0F) | ((current + 1) << 4)

    def _dec_counter(self, pos: int) -> None:
        byte_idx, nibble = divmod(pos, 2)
        current = self._get_counter(pos)
        if current > 0:
            if nibble == 0:
                self._counters[byte_idx] = (self._counters[byte_idx] & 0xF0) | (current - 1)
            else:
                self._counters[byte_idx] = (self._counters[byte_idx] & 0x0F) | ((current - 1) << 4)

    def _hash_positions(self, data: bytes) -> list[int]:
        h1, h2 = _rbsp_hash_pair(data)
        return [((h1 + i * h2) & _MASK64) % self._num_slots for i in range(self._num_hashes)]

    def add(self, item: Union[str, bytes, int]) -> None:
        """Add *item* to the counting filter."""
        data = _to_bytes(item)
        for pos in self._hash_positions(data):
            self._inc_counter(pos)
        self._count += 1

    def remove(self, item: Union[str, bytes, int]) -> bool:
        """Remove *item* from the counting filter.

        Returns:
            True if the item appeared to be present and was removed.
            False if the item was definitely not present (no counters decremented).
        """
        data = _to_bytes(item)
        positions = self._hash_positions(data)

        # Check membership first
        if not all(self._get_counter(pos) > 0 for pos in positions):
            return False

        for pos in positions:
            self._dec_counter(pos)
        self._count -= 1
        return True

    def __contains__(self, item: Union[str, bytes, int]) -> bool:
        data = _to_bytes(item)
        return all(self._get_counter(pos) > 0 for pos in self._hash_positions(data))

    def __len__(self) -> int:
        return self._count

    def clear(self) -> None:
        """Reset all counters to zero."""
        self._counters = bytearray(math.ceil(self._num_slots / 2))
        self._count = 0
