"""Document Deletion via Tombstone Bitmaps.

Implements Lucene-style soft deletes where documents are marked as deleted in
a bitmap rather than physically removed.  Actual removal happens during segment
merging.

The bitmap is backed by a ``bytearray`` bitset for compact, zero-dependency
storage.  All public methods are thread-safe.
"""

from __future__ import annotations

import struct
import threading
from typing import Dict, Set


# ---------------------------------------------------------------------------
# Bitmap
# ---------------------------------------------------------------------------

class DeletionBitmap:
    """Bitset tracking deleted document IDs within a single segment.

    Internally uses a ``bytearray`` where each bit corresponds to a
    document ID.  The array grows automatically when higher IDs are
    marked.

    Args:
        max_doc_id: Pre-allocate enough bits for IDs in ``[0, max_doc_id)``.
    """

    def __init__(self, max_doc_id: int = 0) -> None:
        self._lock = threading.Lock()
        nbytes = (max_doc_id + 7) // 8 if max_doc_id > 0 else 0
        self._bits = bytearray(nbytes)
        self._count = 0

    # -- mutation ----------------------------------------------------------

    def delete(self, doc_id: int) -> None:
        """Mark *doc_id* as deleted."""
        with self._lock:
            self._ensure_capacity(doc_id)
            byte_idx, bit_idx = divmod(doc_id, 8)
            mask = 1 << bit_idx
            if not (self._bits[byte_idx] & mask):
                self._bits[byte_idx] |= mask
                self._count += 1

    def undelete(self, doc_id: int) -> None:
        """Remove the deletion mark for *doc_id*."""
        with self._lock:
            byte_idx, bit_idx = divmod(doc_id, 8)
            if byte_idx >= len(self._bits):
                return
            mask = 1 << bit_idx
            if self._bits[byte_idx] & mask:
                self._bits[byte_idx] &= ~mask
                self._count -= 1

    # -- queries -----------------------------------------------------------

    def is_deleted(self, doc_id: int) -> bool:
        """Return ``True`` if *doc_id* has been marked as deleted."""
        with self._lock:
            byte_idx, bit_idx = divmod(doc_id, 8)
            if byte_idx >= len(self._bits):
                return False
            return bool(self._bits[byte_idx] & (1 << bit_idx))

    def is_live(self, doc_id: int) -> bool:
        """Return ``True`` if *doc_id* has **not** been deleted."""
        return not self.is_deleted(doc_id)

    def deleted_count(self) -> int:
        """Number of documents currently marked as deleted."""
        with self._lock:
            return self._count

    def live_count(self, total_docs: int) -> int:
        """Return ``total_docs - deleted_count()``."""
        with self._lock:
            return total_docs - self._count

    def deleted_docs(self) -> Set[int]:
        """Return the set of all deleted document IDs."""
        with self._lock:
            result: Set[int] = set()
            for byte_idx, byte_val in enumerate(self._bits):
                if byte_val == 0:
                    continue
                base = byte_idx * 8
                for bit_idx in range(8):
                    if byte_val & (1 << bit_idx):
                        result.add(base + bit_idx)
            return result

    # -- serialization -----------------------------------------------------

    _HEADER = b"RBDL"  # magic
    _VERSION = 1

    def to_bytes(self) -> bytes:
        """Serialize the bitmap to a compact binary format.

        Layout::

            4 bytes  magic   "RBDL"
            1 byte   version
            4 bytes  count   (uint32 LE, number of deleted docs)
            4 bytes  length  (uint32 LE, byte length of bitset)
            N bytes  bitset  (raw bytearray)
        """
        with self._lock:
            parts = [
                self._HEADER,
                struct.pack("<BII", self._VERSION, self._count, len(self._bits)),
                bytes(self._bits),
            ]
            return b"".join(parts)

    @classmethod
    def from_bytes(cls, data: bytes) -> DeletionBitmap:
        """Deserialize a bitmap previously produced by :meth:`to_bytes`."""
        if data[:4] != cls._HEADER:
            raise ValueError("Invalid DeletionBitmap header")
        version, count, length = struct.unpack_from("<BII", data, 4)
        if version != cls._VERSION:
            raise ValueError(f"Unsupported DeletionBitmap version {version}")
        offset = 4 + 1 + 4 + 4  # header + version + count + length
        bm = cls.__new__(cls)
        bm._lock = threading.Lock()
        bm._bits = bytearray(data[offset : offset + length])
        bm._count = count
        return bm

    # -- merge -------------------------------------------------------------

    def merge(self, other: DeletionBitmap) -> DeletionBitmap:
        """Return a new bitmap that is the union of *self* and *other*."""
        with self._lock:
            self_bits = bytearray(self._bits)
        with other._lock:
            other_bits = bytearray(other._bits)

        max_len = max(len(self_bits), len(other_bits))
        merged_bits = bytearray(max_len)
        for i in range(len(self_bits)):
            merged_bits[i] = self_bits[i]
        for i in range(len(other_bits)):
            merged_bits[i] |= other_bits[i]

        bm = DeletionBitmap.__new__(DeletionBitmap)
        bm._lock = threading.Lock()
        bm._bits = merged_bits
        # Recount — cheaper than tracking during OR.
        bm._count = 0
        for byte_val in merged_bits:
            # Brian Kernighan popcount
            v = byte_val
            while v:
                v &= v - 1
                bm._count += 1
        return bm

    # -- internal ----------------------------------------------------------

    def _ensure_capacity(self, doc_id: int) -> None:
        """Grow the bytearray so that *doc_id* fits."""
        needed = (doc_id // 8) + 1
        if needed > len(self._bits):
            self._bits.extend(bytearray(needed - len(self._bits)))

    def __repr__(self) -> str:
        return f"DeletionBitmap(deleted={self._count}, bytes={len(self._bits)})"


# ---------------------------------------------------------------------------
# Tracker (per-segment management)
# ---------------------------------------------------------------------------

class DeletionTracker:
    """Manages :class:`DeletionBitmap` instances for multiple segments.

    Each segment gets its own bitmap, created on first use.  Thread-safe.
    """

    def __init__(self) -> None:
        self._lock = threading.Lock()
        self._bitmaps: Dict[str, DeletionBitmap] = {}

    # -- mutation ----------------------------------------------------------

    def delete_document(self, segment_id: str, doc_id: int) -> None:
        """Soft-delete a single document in a segment."""
        bm = self._get_or_create(segment_id)
        bm.delete(doc_id)

    def delete_by_query(self, segment_id: str, matching_doc_ids: Set[int]) -> None:
        """Bulk soft-delete all *matching_doc_ids* in a segment."""
        bm = self._get_or_create(segment_id)
        for doc_id in matching_doc_ids:
            bm.delete(doc_id)

    # -- queries -----------------------------------------------------------

    def is_deleted(self, segment_id: str, doc_id: int) -> bool:
        """Check if a document is deleted in the given segment."""
        with self._lock:
            bm = self._bitmaps.get(segment_id)
        if bm is None:
            return False
        return bm.is_deleted(doc_id)

    def get_bitmap(self, segment_id: str) -> DeletionBitmap:
        """Return the bitmap for *segment_id*, creating one if needed."""
        return self._get_or_create(segment_id)

    def pending_deletions(self, segment_id: str) -> int:
        """Count of pending (soft) deletions in a segment."""
        with self._lock:
            bm = self._bitmaps.get(segment_id)
        if bm is None:
            return 0
        return bm.deleted_count()

    def should_merge(self, segment_id: str, threshold: float = 0.3) -> bool:
        """Return ``True`` if deleted fraction exceeds *threshold*.

        The deleted fraction is ``deleted_count / (deleted_count + live)``
        where *live* is approximated by the total docs tracked in the
        bitmap (deleted + live IDs seen).  A segment with no documents
        never triggers a merge.

        For accurate results the caller should use::

            bitmap.deleted_count() / total_docs_in_segment > threshold

        This helper provides a reasonable heuristic using only the bitmap
        state (it considers the highest allocated doc-ID as the total).
        """
        with self._lock:
            bm = self._bitmaps.get(segment_id)
        if bm is None:
            return False
        with bm._lock:
            total = len(bm._bits) * 8
        if total == 0:
            return False
        return bm.deleted_count() / total > threshold

    # -- internal ----------------------------------------------------------

    def _get_or_create(self, segment_id: str) -> DeletionBitmap:
        with self._lock:
            bm = self._bitmaps.get(segment_id)
            if bm is None:
                bm = DeletionBitmap()
                self._bitmaps[segment_id] = bm
            return bm
