"""Compressed Postings Lists for RBSP.

Lucene-style postings lists with delta encoding, variable-byte integer
compression, and skip pointers for fast intersection.

Each postings list stores (doc_id, term_freq, positions) tuples for a
single term, sorted by doc_id.  Serialization uses:

- **Delta encoding** on doc_ids (store gaps instead of absolute values)
- **Variable-byte (varint) encoding** for all integers (7 data bits per
  byte, MSB=1 means more bytes follow)
- **Skip pointers** every SKIP_INTERVAL entries for O(sqrt(n)) advance

Thread-safe for concurrent reads; not designed for concurrent mutation.
"""

from __future__ import annotations

from dataclasses import dataclass, field
from typing import List, Optional


# ---------------------------------------------------------------------------
# Constants
# ---------------------------------------------------------------------------

SKIP_INTERVAL: int = 128


# ---------------------------------------------------------------------------
# Varint helpers (self-contained, no external deps)
# ---------------------------------------------------------------------------

def _encode_varint(value: int) -> bytes:
    """Encode a non-negative integer as a variable-byte sequence.

    Uses MSB continuation scheme: 7 data bits per byte, bit 7 (MSB) is
    set to 1 when more bytes follow, 0 on the final byte.
    """
    if value < 0:
        raise ValueError("varint value must be non-negative")
    buf = bytearray()
    while True:
        low7 = value & 0x7F
        value >>= 7
        if value:
            buf.append(low7 | 0x80)
        else:
            buf.append(low7)
            break
    return bytes(buf)


def _decode_varint(data: bytes, offset: int) -> tuple[int, int]:
    """Decode a varint starting at *offset* in *data*.

    Returns:
        (decoded_value, new_offset) where *new_offset* points to the
        byte immediately after the last consumed byte.
    """
    result = 0
    shift = 0
    while True:
        if offset >= len(data):
            raise ValueError("unexpected end of data while decoding varint")
        b = data[offset]
        offset += 1
        result |= (b & 0x7F) << shift
        if not (b & 0x80):
            break
        shift += 7
    return result, offset


# ---------------------------------------------------------------------------
# PostingEntry
# ---------------------------------------------------------------------------

@dataclass(slots=True)
class PostingEntry:
    """A single posting: one (document, term_frequency) pair with optional
    position offsets."""

    doc_id: int
    term_freq: int
    positions: Optional[List[int]] = None


# ---------------------------------------------------------------------------
# PostingList
# ---------------------------------------------------------------------------

class PostingList:
    """Compressed postings list for a single term.

    Postings are maintained in ascending doc_id order.  Serialization
    compresses doc_ids via delta-encoding and stores all integers as
    varints.  Skip pointers are written every :data:`SKIP_INTERVAL`
    entries so that :meth:`advance_to` can skip large segments of the
    byte stream.

    Usage::

        pl = PostingList("python")
        pl.add(1, 2, [0, 15])
        pl.add(5, 1, [7])
        pl.add(12, 3)

        entry = pl.advance_to(5)   # PostingEntry(doc_id=5, ...)
        raw = pl.to_bytes()
        pl2 = PostingList.from_bytes(raw, "python")
    """

    __slots__ = ("term", "_entries")

    def __init__(self, term: str) -> None:
        self.term: str = term
        self._entries: List[PostingEntry] = []

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

    def add(self, doc_id: int, term_freq: int,
            positions: Optional[List[int]] = None) -> None:
        """Append a posting.  *doc_id* must be strictly greater than the
        last added doc_id (postings must be added in sorted order).

        Args:
            doc_id:    Non-negative document identifier.
            term_freq: Term frequency in this document (>= 1).
            positions: Optional sorted list of positions within the doc.

        Raises:
            ValueError: If doc_id is out of order or negative.
        """
        if doc_id < 0:
            raise ValueError("doc_id must be non-negative")
        if self._entries and doc_id <= self._entries[-1].doc_id:
            raise ValueError(
                f"doc_ids must be added in strictly ascending order "
                f"(got {doc_id} after {self._entries[-1].doc_id})"
            )
        self._entries.append(PostingEntry(doc_id, term_freq, positions))

    # ------------------------------------------------------------------
    # Access
    # ------------------------------------------------------------------

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

    def __iter__(self):
        return iter(self._entries)

    def advance_to(self, target: int) -> Optional[PostingEntry]:
        """Find the first posting with ``doc_id >= target``.

        Uses binary search on the in-memory list (O(log n)).  When
        operating on a deserialized list, the skip pointers embedded in
        the byte representation enable fast seeking during decode; this
        method provides equivalent semantics for the in-memory form.

        Args:
            target: Minimum doc_id to match.

        Returns:
            The matching :class:`PostingEntry`, or ``None`` if every
            posting has ``doc_id < target``.
        """
        entries = self._entries
        lo, hi = 0, len(entries)
        while lo < hi:
            mid = (lo + hi) // 2
            if entries[mid].doc_id < target:
                lo = mid + 1
            else:
                hi = mid
        if lo < len(entries):
            return entries[lo]
        return None

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

    def to_bytes(self) -> bytes:
        """Serialize the postings list to a compact byte representation.

        Format overview::

            [num_entries: varint]
            [num_skips:  varint]
            -- skip table (num_skips entries) --
            [skip_doc_id: varint] [skip_byte_offset: varint]  ...
            -- postings data --
            [delta_doc_id: varint] [term_freq: varint]
            [has_positions: varint (0 or 1)]
            if has_positions:
                [num_positions: varint]
                [delta_pos0: varint] [delta_pos1: varint] ...

        Skip pointers record the *absolute* doc_id and the byte offset
        (relative to the start of the postings-data section) at every
        SKIP_INTERVAL-th posting.
        """
        entries = self._entries
        n = len(entries)
        if n == 0:
            return _encode_varint(0)

        # --- First pass: encode all postings to collect byte offsets ---
        postings_buf = bytearray()
        skip_doc_ids: List[int] = []
        skip_offsets: List[int] = []
        prev_doc_id = 0

        for i, entry in enumerate(entries):
            # Record skip pointer at every SKIP_INTERVAL boundary
            if i > 0 and i % SKIP_INTERVAL == 0:
                skip_doc_ids.append(entry.doc_id)
                skip_offsets.append(len(postings_buf))

            # Delta-encode doc_id
            delta = entry.doc_id - prev_doc_id
            prev_doc_id = entry.doc_id
            postings_buf.extend(_encode_varint(delta))

            # Term frequency
            postings_buf.extend(_encode_varint(entry.term_freq))

            # Positions (optional)
            if entry.positions is not None and len(entry.positions) > 0:
                postings_buf.extend(_encode_varint(1))  # has_positions = 1
                postings_buf.extend(_encode_varint(len(entry.positions)))
                prev_pos = 0
                for pos in entry.positions:
                    postings_buf.extend(_encode_varint(pos - prev_pos))
                    prev_pos = pos
            else:
                postings_buf.extend(_encode_varint(0))  # has_positions = 0

        # --- Build header ---
        header = bytearray()
        header.extend(_encode_varint(n))
        num_skips = len(skip_doc_ids)
        header.extend(_encode_varint(num_skips))
        for j in range(num_skips):
            header.extend(_encode_varint(skip_doc_ids[j]))
            header.extend(_encode_varint(skip_offsets[j]))

        return bytes(header) + bytes(postings_buf)

    @classmethod
    def from_bytes(cls, data: bytes, term: str) -> "PostingList":
        """Deserialize a postings list from bytes produced by :meth:`to_bytes`.

        Args:
            data: The byte representation.
            term: The term this list belongs to.

        Returns:
            A fully-populated :class:`PostingList`.
        """
        pl = cls(term)
        if not data:
            return pl

        offset = 0

        # Number of entries
        n, offset = _decode_varint(data, offset)
        if n == 0:
            return pl

        # Skip table (we read it but don't need it for full decode --
        # skip pointers are used in advance_to_bytes for streaming).
        num_skips, offset = _decode_varint(data, offset)
        for _ in range(num_skips):
            _skip_doc_id, offset = _decode_varint(data, offset)
            _skip_offset, offset = _decode_varint(data, offset)

        # Decode postings
        prev_doc_id = 0
        for _ in range(n):
            delta, offset = _decode_varint(data, offset)
            doc_id = prev_doc_id + delta
            prev_doc_id = doc_id

            term_freq, offset = _decode_varint(data, offset)

            has_positions, offset = _decode_varint(data, offset)
            positions: Optional[List[int]] = None
            if has_positions:
                num_positions, offset = _decode_varint(data, offset)
                positions = []
                prev_pos = 0
                for _ in range(num_positions):
                    pos_delta, offset = _decode_varint(data, offset)
                    prev_pos += pos_delta
                    positions.append(prev_pos)

            pl._entries.append(PostingEntry(doc_id, term_freq, positions))

        return pl

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

    @staticmethod
    def intersect(a: "PostingList", b: "PostingList") -> List[int]:
        """Compute the intersection of two postings lists.

        Uses a two-pointer merge with :meth:`advance_to` for skipping,
        analogous to Lucene's ``ConjunctionScorer``.

        Args:
            a: First postings list (sorted by doc_id).
            b: Second postings list (sorted by doc_id).

        Returns:
            Sorted list of doc_ids present in both *a* and *b*.
        """
        result: List[int] = []
        iter_a = iter(a)
        iter_b = iter(b)

        entry_a = next(iter_a, None)
        entry_b = next(iter_b, None)

        while entry_a is not None and entry_b is not None:
            if entry_a.doc_id == entry_b.doc_id:
                result.append(entry_a.doc_id)
                entry_a = next(iter_a, None)
                entry_b = next(iter_b, None)
            elif entry_a.doc_id < entry_b.doc_id:
                # Try to skip forward in a
                entry_a = a.advance_to(entry_b.doc_id)
            else:
                # Try to skip forward in b
                entry_b = b.advance_to(entry_a.doc_id)

        return result

    @staticmethod
    def union(a: "PostingList", b: "PostingList") -> List[int]:
        """Compute the union of two postings lists.

        Uses a standard two-pointer merge.

        Args:
            a: First postings list (sorted by doc_id).
            b: Second postings list (sorted by doc_id).

        Returns:
            Sorted list of doc_ids present in either *a* or *b*.
        """
        result: List[int] = []
        iter_a = iter(a)
        iter_b = iter(b)

        entry_a = next(iter_a, None)
        entry_b = next(iter_b, None)

        while entry_a is not None and entry_b is not None:
            if entry_a.doc_id == entry_b.doc_id:
                result.append(entry_a.doc_id)
                entry_a = next(iter_a, None)
                entry_b = next(iter_b, None)
            elif entry_a.doc_id < entry_b.doc_id:
                result.append(entry_a.doc_id)
                entry_a = next(iter_a, None)
            else:
                result.append(entry_b.doc_id)
                entry_b = next(iter_b, None)

        # Drain remaining
        while entry_a is not None:
            result.append(entry_a.doc_id)
            entry_a = next(iter_a, None)
        while entry_b is not None:
            result.append(entry_b.doc_id)
            entry_b = next(iter_b, None)

        return result
