"""
RBSP Roaring Bitmap Layer.

Provides compressed bitmap data structures for ultra-fast set operations on
document IDs.  Roaring bitmaps partition the 32-bit integer space into chunks
of 65536 (high 16 bits = chunk key) and use the most compact representation
per chunk:

    - **Array container**: Sorted ``array.array('H')`` for sparse chunks
      (fewer than 4096 elements).
    - **Bitset container**: 8 KB bit vector for dense chunks (4096+ elements).

Container types are promoted/demoted automatically on insertion and removal,
giving both excellent compression on sparse data and O(1) bit operations on
dense data.

Typical use in RBSP is to represent posting lists (sets of doc IDs) so that
Boolean query operations (AND/OR/NOT) run in compressed space without
materializing full sets.

Zero external dependencies.  Pure Python.
"""

from __future__ import annotations

import array
import struct
from typing import Iterator, Union

# Threshold at which an array container is promoted to a bitset container.
# At 4096 entries the sorted array (4096 * 2 bytes = 8 KB) costs the same as
# a full 8 KB bitset, so bitset becomes more efficient.
_ARRAY_MAX = 4096

# Number of 16-bit values in one chunk.
_CHUNK = 1 << 16  # 65536

# Size of the bitset container in bytes (65536 bits = 8192 bytes).
_BITSET_BYTES = _CHUNK // 8


# ---------------------------------------------------------------------------
# Container types
# ---------------------------------------------------------------------------

class _ArrayContainer:
    """Sorted array of 16-bit low values.  Used when cardinality < 4096."""

    __slots__ = ("data",)

    def __init__(self) -> None:
        self.data: array.array[int] = array.array("H")

    def add(self, low: int) -> bool:
        """Insert *low*.  Return True if the element was new."""
        # Binary search
        lo, hi = 0, len(self.data)
        while lo < hi:
            mid = (lo + hi) >> 1
            if self.data[mid] < low:
                lo = mid + 1
            else:
                hi = mid
        if lo < len(self.data) and self.data[lo] == low:
            return False
        self.data.insert(lo, low)
        return True

    def remove(self, low: int) -> bool:
        """Remove *low*.  Return True if the element was present."""
        lo, hi = 0, len(self.data)
        while lo < hi:
            mid = (lo + hi) >> 1
            if self.data[mid] < low:
                lo = mid + 1
            else:
                hi = mid
        if lo < len(self.data) and self.data[lo] == low:
            del self.data[lo]
            return True
        return False

    def __contains__(self, low: int) -> bool:
        lo, hi = 0, len(self.data)
        while lo < hi:
            mid = (lo + hi) >> 1
            if self.data[mid] < low:
                lo = mid + 1
            else:
                hi = mid
        return lo < len(self.data) and self.data[lo] == low

    def __len__(self) -> int:
        return len(self.data)

    def __iter__(self) -> Iterator[int]:
        return iter(self.data)

    # -- Set operations (return new containers) ----------------------------

    def and_array(self, other: _ArrayContainer) -> _ArrayContainer:
        """Intersection of two array containers."""
        result = _ArrayContainer()
        a, b = self.data, other.data
        i, j = 0, 0
        la, lb = len(a), len(b)
        while i < la and j < lb:
            if a[i] < b[j]:
                i += 1
            elif a[i] > b[j]:
                j += 1
            else:
                result.data.append(a[i])
                i += 1
                j += 1
        return result

    def or_array(self, other: _ArrayContainer) -> Union[_ArrayContainer, _BitsetContainer]:
        """Union of two array containers.  May promote to bitset."""
        result = _ArrayContainer()
        a, b = self.data, other.data
        i, j = 0, 0
        la, lb = len(a), len(b)
        while i < la and j < lb:
            if a[i] < b[j]:
                result.data.append(a[i])
                i += 1
            elif a[i] > b[j]:
                result.data.append(b[j])
                j += 1
            else:
                result.data.append(a[i])
                i += 1
                j += 1
        while i < la:
            result.data.append(a[i])
            i += 1
        while j < lb:
            result.data.append(b[j])
            j += 1
        if len(result) >= _ARRAY_MAX:
            return result.to_bitset()
        return result

    def sub_array(self, other: _ArrayContainer) -> _ArrayContainer:
        """Difference: self - other."""
        result = _ArrayContainer()
        a, b = self.data, other.data
        i, j = 0, 0
        la, lb = len(a), len(b)
        while i < la and j < lb:
            if a[i] < b[j]:
                result.data.append(a[i])
                i += 1
            elif a[i] > b[j]:
                j += 1
            else:
                i += 1
                j += 1
        while i < la:
            result.data.append(a[i])
            i += 1
        return result

    def to_bitset(self) -> _BitsetContainer:
        """Promote to a bitset container."""
        bc = _BitsetContainer()
        for v in self.data:
            bc._set_bit(v)
        bc._card = len(self.data)
        return bc

    def to_bytes(self) -> bytes:
        """Serialize: 1-byte type tag (0x01) + raw uint16 array."""
        return b"\x01" + self.data.tobytes()

    @classmethod
    def from_bytes(cls, data: bytes) -> _ArrayContainer:
        c = cls()
        c.data = array.array("H")
        c.data.frombytes(data)
        return c


class _BitsetContainer:
    """8 KB bitset for dense chunks.  Used when cardinality >= 4096."""

    __slots__ = ("bits", "_card")

    def __init__(self) -> None:
        self.bits = bytearray(_BITSET_BYTES)
        self._card = 0

    def _set_bit(self, low: int) -> None:
        self.bits[low >> 3] |= 1 << (low & 7)

    def _clear_bit(self, low: int) -> None:
        self.bits[low >> 3] &= ~(1 << (low & 7))

    def _get_bit(self, low: int) -> bool:
        return bool(self.bits[low >> 3] & (1 << (low & 7)))

    def add(self, low: int) -> bool:
        if self._get_bit(low):
            return False
        self._set_bit(low)
        self._card += 1
        return True

    def remove(self, low: int) -> bool:
        if not self._get_bit(low):
            return False
        self._clear_bit(low)
        self._card -= 1
        return True

    def __contains__(self, low: int) -> bool:
        return self._get_bit(low)

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

    def __iter__(self) -> Iterator[int]:
        bits = self.bits
        for byte_idx in range(_BITSET_BYTES):
            b = bits[byte_idx]
            if b == 0:
                continue
            base = byte_idx << 3
            while b:
                # Isolate lowest set bit
                lsb = b & -b
                yield base + lsb.bit_length() - 1
                b ^= lsb

    # -- Set operations ----------------------------------------------------

    def and_bitset(self, other: _BitsetContainer) -> Union[_ArrayContainer, _BitsetContainer]:
        """Intersection of two bitset containers.  May demote to array."""
        result = _BitsetContainer()
        sa, sb = self.bits, other.bits
        card = 0
        rb = result.bits
        for i in range(_BITSET_BYTES):
            v = sa[i] & sb[i]
            rb[i] = v
            # popcount via bin()
            card += bin(v).count("1")
        result._card = card
        if card < _ARRAY_MAX:
            return result.to_array()
        return result

    def or_bitset(self, other: _BitsetContainer) -> _BitsetContainer:
        """Union of two bitset containers."""
        result = _BitsetContainer()
        sa, sb = self.bits, other.bits
        card = 0
        rb = result.bits
        for i in range(_BITSET_BYTES):
            v = sa[i] | sb[i]
            rb[i] = v
            card += bin(v).count("1")
        result._card = card
        return result

    def sub_bitset(self, other: _BitsetContainer) -> Union[_ArrayContainer, _BitsetContainer]:
        """Difference: self - other.  May demote to array."""
        result = _BitsetContainer()
        sa, sb = self.bits, other.bits
        card = 0
        rb = result.bits
        for i in range(_BITSET_BYTES):
            v = sa[i] & ~sb[i]
            rb[i] = v
            card += bin(v).count("1")
        result._card = card
        if card < _ARRAY_MAX:
            return result.to_array()
        return result

    def and_array(self, arr: _ArrayContainer) -> _ArrayContainer:
        """Intersection of bitset with array.  Always yields array."""
        result = _ArrayContainer()
        for v in arr.data:
            if self._get_bit(v):
                result.data.append(v)
        return result

    def or_array(self, arr: _ArrayContainer) -> _BitsetContainer:
        """Union of bitset with array.  Always stays bitset."""
        result = _BitsetContainer()
        result.bits = bytearray(self.bits)
        result._card = self._card
        for v in arr.data:
            if result.add(v):
                pass  # add already increments _card
        return result

    def sub_array(self, arr: _ArrayContainer) -> Union[_ArrayContainer, _BitsetContainer]:
        """Difference: bitset - array.  May demote to array."""
        result = _BitsetContainer()
        result.bits = bytearray(self.bits)
        result._card = self._card
        for v in arr.data:
            result.remove(v)
        if result._card < _ARRAY_MAX:
            return result.to_array()
        return result

    def to_array(self) -> _ArrayContainer:
        """Demote to an array container."""
        ac = _ArrayContainer()
        ac.data = array.array("H", iter(self))
        return ac

    def to_bytes(self) -> bytes:
        """Serialize: 1-byte type tag (0x02) + raw bitset bytes."""
        return b"\x02" + bytes(self.bits)

    @classmethod
    def from_bytes(cls, data: bytes) -> _BitsetContainer:
        c = cls()
        c.bits = bytearray(data)
        c._card = sum(bin(b).count("1") for b in c.bits)
        return c


# ---------------------------------------------------------------------------
# Cross-container operations
# ---------------------------------------------------------------------------

_Container = Union[_ArrayContainer, _BitsetContainer]


def _intersect(a: _Container, b: _Container) -> _Container:
    if isinstance(a, _ArrayContainer) and isinstance(b, _ArrayContainer):
        return a.and_array(b)
    if isinstance(a, _BitsetContainer) and isinstance(b, _BitsetContainer):
        return a.and_bitset(b)
    # Mixed: use bitset.and_array
    if isinstance(a, _BitsetContainer):
        return a.and_array(b)  # type: ignore[arg-type]
    return b.and_array(a)  # type: ignore[union-attr, arg-type]


def _union(a: _Container, b: _Container) -> _Container:
    if isinstance(a, _ArrayContainer) and isinstance(b, _ArrayContainer):
        return a.or_array(b)
    if isinstance(a, _BitsetContainer) and isinstance(b, _BitsetContainer):
        return a.or_bitset(b)
    # Mixed: use bitset.or_array
    if isinstance(a, _BitsetContainer):
        return a.or_array(b)  # type: ignore[arg-type]
    return b.or_array(a)  # type: ignore[union-attr, arg-type]


def _difference(a: _Container, b: _Container) -> _Container:
    if isinstance(a, _ArrayContainer) and isinstance(b, _ArrayContainer):
        return a.sub_array(b)
    if isinstance(a, _BitsetContainer) and isinstance(b, _BitsetContainer):
        return a.sub_bitset(b)
    if isinstance(a, _BitsetContainer) and isinstance(b, _ArrayContainer):
        return a.sub_array(b)
    # a is array, b is bitset: iterate a, keep only items NOT in b
    result = _ArrayContainer()
    for v in a:  # type: ignore[union-attr]
        if v not in b:
            result.data.append(v)
    return result


# ---------------------------------------------------------------------------
# RoaringBitmap
# ---------------------------------------------------------------------------

class RoaringBitmap:
    """Compressed bitmap for efficient integer set operations.

    Splits the 32-bit integer space into 65536-element chunks.  Each chunk
    uses either a sorted array (sparse) or a bitset (dense), switching
    automatically based on cardinality.

    Supports the standard set operators ``&`` (intersection), ``|`` (union),
    ``-`` (difference) as well as membership testing, iteration, and
    serialization.

    Example::

        bm = RoaringBitmap()
        for doc_id in range(10_000):
            bm.add(doc_id)
        assert 42 in bm
        assert len(bm) == 10_000
    """

    __slots__ = ("_containers",)

    def __init__(self) -> None:
        # key = high 16 bits, value = container for the low 16 bits
        self._containers: dict[int, _Container] = {}

    @staticmethod
    def _split(value: int) -> tuple[int, int]:
        """Split a 32-bit integer into (high-16, low-16)."""
        return (value >> 16) & 0xFFFF, value & 0xFFFF

    @staticmethod
    def _join(high: int, low: int) -> int:
        return (high << 16) | low

    # -- Mutation ----------------------------------------------------------

    def add(self, value: int) -> None:
        """Add a 32-bit integer to the bitmap."""
        if value < 0 or value > 0xFFFFFFFF:
            raise ValueError("value must be a 32-bit unsigned integer (0..4294967295)")
        high, low = self._split(value)
        container = self._containers.get(high)
        if container is None:
            container = _ArrayContainer()
            self._containers[high] = container
        container.add(low)
        # Promote array -> bitset if threshold crossed
        if isinstance(container, _ArrayContainer) and len(container) >= _ARRAY_MAX:
            self._containers[high] = container.to_bitset()

    def remove(self, value: int) -> None:
        """Remove a 32-bit integer from the bitmap.

        Raises:
            KeyError: If the value is not present.
        """
        high, low = self._split(value)
        container = self._containers.get(high)
        if container is None or not container.remove(low):
            raise KeyError(value)
        # Clean up empty containers
        if len(container) == 0:
            del self._containers[high]
        # Demote bitset -> array if threshold crossed
        elif isinstance(container, _BitsetContainer) and len(container) < _ARRAY_MAX:
            self._containers[high] = container.to_array()

    def discard(self, value: int) -> None:
        """Remove *value* if present; do nothing otherwise."""
        try:
            self.remove(value)
        except KeyError:
            pass

    # -- Queries -----------------------------------------------------------

    def __contains__(self, value: object) -> bool:
        if not isinstance(value, int):
            return False
        if value < 0 or value > 0xFFFFFFFF:
            return False
        high, low = self._split(value)
        container = self._containers.get(high)
        if container is None:
            return False
        return low in container

    def __len__(self) -> int:
        return sum(len(c) for c in self._containers.values())

    def __bool__(self) -> bool:
        return bool(self._containers)

    def __iter__(self) -> Iterator[int]:
        """Iterate over all values in sorted ascending order."""
        for high in sorted(self._containers):
            base = high << 16
            for low in self._containers[high]:
                yield base | low

    def min(self) -> int:
        """Return the smallest value in the bitmap.

        Raises:
            ValueError: If the bitmap is empty.
        """
        if not self._containers:
            raise ValueError("empty bitmap has no min")
        high = builtins_min(self._containers)
        low = builtins_min(self._containers[high])
        return self._join(high, low)

    def max(self) -> int:
        """Return the largest value in the bitmap.

        Raises:
            ValueError: If the bitmap is empty.
        """
        if not self._containers:
            raise ValueError("empty bitmap has no max")
        high = builtins_max(self._containers)
        low = builtins_max(self._containers[high])
        return self._join(high, low)

    # -- Set operations ----------------------------------------------------

    def __and__(self, other: RoaringBitmap) -> RoaringBitmap:
        """Intersection: return new bitmap with elements in both."""
        if not isinstance(other, RoaringBitmap):
            return NotImplemented
        result = RoaringBitmap()
        for high in self._containers:
            if high in other._containers:
                c = _intersect(self._containers[high], other._containers[high])
                if len(c) > 0:
                    result._containers[high] = c
        return result

    def __or__(self, other: RoaringBitmap) -> RoaringBitmap:
        """Union: return new bitmap with elements in either."""
        if not isinstance(other, RoaringBitmap):
            return NotImplemented
        result = RoaringBitmap()
        all_keys = set(self._containers) | set(other._containers)
        for high in all_keys:
            a = self._containers.get(high)
            b = other._containers.get(high)
            if a is not None and b is not None:
                result._containers[high] = _union(a, b)
            elif a is not None:
                result._containers[high] = a  # shared ref is safe (immutable ops)
            else:
                result._containers[high] = b  # type: ignore[assignment]
        return result

    def __sub__(self, other: RoaringBitmap) -> RoaringBitmap:
        """Difference: return new bitmap with elements in self but not other."""
        if not isinstance(other, RoaringBitmap):
            return NotImplemented
        result = RoaringBitmap()
        for high, container in self._containers.items():
            b = other._containers.get(high)
            if b is None:
                result._containers[high] = container
            else:
                c = _difference(container, b)
                if len(c) > 0:
                    result._containers[high] = c
        return result

    def __eq__(self, other: object) -> bool:
        if not isinstance(other, RoaringBitmap):
            return NotImplemented
        if len(self) != len(other):
            return False
        # Compare sorted values
        return list(self) == list(other)

    def __repr__(self) -> str:
        card = len(self)
        chunks = len(self._containers)
        return f"RoaringBitmap(cardinality={card}, containers={chunks})"

    # -- Serialization -----------------------------------------------------
    #
    # Wire format:
    #   4 bytes: number of containers (uint32 big-endian)
    #   For each container:
    #       2 bytes: chunk key (uint16 big-endian)
    #       4 bytes: container payload length (uint32 big-endian)
    #       N bytes: container payload (type-tagged: 0x01=array, 0x02=bitset)

    _HEADER = struct.Struct(">I")
    _CONTAINER_HEADER = struct.Struct(">HI")

    def to_bytes(self) -> bytes:
        """Serialize the bitmap to a compact binary representation."""
        parts: list[bytes] = [self._HEADER.pack(len(self._containers))]
        for high in sorted(self._containers):
            payload = self._containers[high].to_bytes()
            parts.append(self._CONTAINER_HEADER.pack(high, len(payload)))
            parts.append(payload)
        return b"".join(parts)

    @classmethod
    def from_bytes(cls, data: bytes) -> RoaringBitmap:
        """Deserialize a bitmap from bytes produced by ``to_bytes``."""
        bm = cls()
        if len(data) < 4:
            raise ValueError("Data too short for roaring bitmap header")
        (num_containers,) = cls._HEADER.unpack_from(data, 0)
        offset = 4
        ch = cls._CONTAINER_HEADER
        for _ in range(num_containers):
            if offset + 6 > len(data):
                raise ValueError("Truncated container header")
            high, payload_len = ch.unpack_from(data, offset)
            offset += 6
            if offset + payload_len > len(data):
                raise ValueError("Truncated container payload")
            payload = data[offset : offset + payload_len]
            offset += payload_len
            tag = payload[0:1]
            body = payload[1:]
            if tag == b"\x01":
                bm._containers[high] = _ArrayContainer.from_bytes(body)
            elif tag == b"\x02":
                bm._containers[high] = _BitsetContainer.from_bytes(body)
            else:
                raise ValueError(f"Unknown container tag: {tag!r}")
        return bm

    # -- Convenience constructors ------------------------------------------

    @classmethod
    def from_iterable(cls, values: Iterator[int]) -> RoaringBitmap:
        """Build a RoaringBitmap from an iterable of integers."""
        bm = cls()
        for v in values:
            bm.add(v)
        return bm


# Bring builtins into module scope to avoid name collision with methods.
import builtins as _builtins

builtins_min = _builtins.min
builtins_max = _builtins.max
