"""
RBSP Skip List — Fast sorted posting list with O(log n) advance.

Provides a probabilistic skip list for efficient posting list traversal and
intersection.  Skip lists maintain sorted order with multiple levels of
forward pointers, enabling O(log n) search and advance-to operations.

Used for posting list intersection during query evaluation — the
``advance_to`` method allows skipping past irrelevant document IDs, giving
O(n log m) intersection of two posting lists of sizes n and m.

    - ``SkipList``   Probabilistic skip list with O(log n) lookup / advance.
"""

from __future__ import annotations

import random
from typing import Iterator, Optional

_MAX_LEVEL = 16
_P = 0.5


def _random_level() -> int:
    """Return a random level in [0, _MAX_LEVEL) with geometric distribution."""
    level = 0
    while level < _MAX_LEVEL - 1 and random.random() < _P:
        level += 1
    return level


class _Node:
    """Internal skip list node with forward pointers at multiple levels."""

    __slots__ = ("value", "forward")

    def __init__(self, value: Optional[int], level: int) -> None:
        self.value = value
        self.forward: list[Optional[_Node]] = [None] * (level + 1)


class SkipList:
    """Probabilistic skip list for sorted integer sets.

    Supports O(log n) lookup, O(log n) advance-to, and O(n) ordered
    iteration.  Designed for posting list operations in the RBSP query
    evaluator.

    Example::

        sl = SkipList()
        for doc_id in [10, 20, 30, 40, 50]:
            sl.add(doc_id)

        assert 30 in sl
        assert sl.advance_to(25) == 30
    """

    def __init__(self) -> None:
        self._head = _Node(None, _MAX_LEVEL - 1)
        self._level = 0
        self._size = 0

    # -- Insertion --------------------------------------------------------

    def add(self, value: int) -> None:
        """Insert *value* into the skip list.

        If *value* already exists, this is a no-op (set semantics).
        """
        update: list[Optional[_Node]] = [None] * _MAX_LEVEL
        current = self._head

        for i in range(self._level, -1, -1):
            while current.forward[i] is not None and current.forward[i].value < value:  # type: ignore[union-attr]
                current = current.forward[i]  # type: ignore[assignment]
            update[i] = current

        # Check if value already exists
        candidate = current.forward[0]
        if candidate is not None and candidate.value == value:
            return

        new_level = _random_level()
        if new_level > self._level:
            for i in range(self._level + 1, new_level + 1):
                update[i] = self._head
            self._level = new_level

        node = _Node(value, new_level)
        for i in range(new_level + 1):
            node.forward[i] = update[i].forward[i]  # type: ignore[union-attr]
            update[i].forward[i] = node  # type: ignore[union-attr]

        self._size += 1

    # -- Lookup -----------------------------------------------------------

    def __contains__(self, value: int) -> bool:  # type: ignore[override]
        """O(log n) membership test."""
        current = self._head
        for i in range(self._level, -1, -1):
            while current.forward[i] is not None and current.forward[i].value < value:  # type: ignore[union-attr]
                current = current.forward[i]  # type: ignore[assignment]
        candidate = current.forward[0]
        return candidate is not None and candidate.value == value

    # -- Advance ----------------------------------------------------------

    def advance_to(self, target: int) -> Optional[int]:
        """Advance to the first value >= *target*.

        Returns the value found, or ``None`` if all values are less than
        *target*.  This is the key operation for posting list intersection.
        """
        current = self._head
        for i in range(self._level, -1, -1):
            while current.forward[i] is not None and current.forward[i].value < target:  # type: ignore[union-attr]
                current = current.forward[i]  # type: ignore[assignment]
        candidate = current.forward[0]
        if candidate is not None:
            return candidate.value
        return None

    # -- Iteration --------------------------------------------------------

    def __iter__(self) -> Iterator[int]:
        """Iterate over all values in sorted order."""
        node = self._head.forward[0]
        while node is not None:
            yield node.value  # type: ignore[misc]
            node = node.forward[0]

    # -- Size -------------------------------------------------------------

    def __len__(self) -> int:
        """Return the number of elements in the skip list."""
        return self._size

    # -- Set operations (static) ------------------------------------------

    @staticmethod
    def intersection(a: SkipList, b: SkipList) -> SkipList:
        """Compute the intersection of two skip lists.

        Uses skip-accelerated merge: walks the shorter list and uses
        ``advance_to`` on the longer list to skip past non-matching values.
        Runs in O(min(n, m) * log(max(n, m))) time.
        """
        result = SkipList()

        # Walk the smaller list, probe the larger
        if len(a) > len(b):
            a, b = b, a

        b_cursor = b._head
        for val in a:
            # Advance b_cursor using skip pointers
            found_val = _advance_cursor(b_cursor, b._level, val)
            if found_val == val:
                result.add(val)
                # Advance b_cursor past val for the next iteration
                b_cursor_new = b._head
                for i in range(b._level, -1, -1):
                    while (
                        b_cursor_new.forward[i] is not None
                        and b_cursor_new.forward[i].value <= val  # type: ignore[union-attr]
                    ):
                        b_cursor_new.forward[i]  # type: ignore[union-attr]
                        b_cursor_new = b_cursor_new.forward[i]  # type: ignore[assignment]

        return result

    @staticmethod
    def union(a: SkipList, b: SkipList) -> SkipList:
        """Compute the union of two skip lists.

        Merges both lists in sorted order.  Runs in O(n + m) time.
        """
        result = SkipList()
        node_a = a._head.forward[0]
        node_b = b._head.forward[0]

        while node_a is not None and node_b is not None:
            va = node_a.value  # type: ignore[union-attr]
            vb = node_b.value  # type: ignore[union-attr]
            if va < vb:
                result.add(va)
                node_a = node_a.forward[0]
            elif va > vb:
                result.add(vb)
                node_b = node_b.forward[0]
            else:
                result.add(va)
                node_a = node_a.forward[0]
                node_b = node_b.forward[0]

        while node_a is not None:
            result.add(node_a.value)  # type: ignore[arg-type]
            node_a = node_a.forward[0]

        while node_b is not None:
            result.add(node_b.value)  # type: ignore[arg-type]
            node_b = node_b.forward[0]

        return result

    # -- Repr -------------------------------------------------------------

    def __repr__(self) -> str:
        items = list(self)
        if len(items) > 8:
            preview = ", ".join(str(x) for x in items[:8]) + ", ..."
        else:
            preview = ", ".join(str(x) for x in items)
        return f"SkipList([{preview}], len={self._size})"


def _advance_cursor(head: _Node, level: int, target: int) -> Optional[int]:
    """Advance from *head* using skip pointers to find first value >= *target*."""
    current = head
    for i in range(level, -1, -1):
        while current.forward[i] is not None and current.forward[i].value < target:  # type: ignore[union-attr]
            current = current.forward[i]  # type: ignore[assignment]
    candidate = current.forward[0]
    if candidate is not None:
        return candidate.value
    return None
