"""Columnar Document Values for RBSP.

Stores per-document field values in a column-oriented format, enabling
O(1) field access for sorting, faceting, and scripting without loading
the full document.  Inspired by Lucene's doc values.

Supported types:

- **NUMERIC**:     int/float values per document (sorting, range queries)
- **BINARY**:      arbitrary bytes per document
- **SORTED**:      ordinal-encoded string values (sorting, grouping)
- **SORTED_SET**:  multi-valued ordinal-encoded strings (faceting)

Thread-safe for concurrent reads and writes.  Zero external dependencies.
"""

from __future__ import annotations

import bisect
import threading
from collections import defaultdict
from enum import Enum, auto
from typing import Dict, FrozenSet, List, Optional, Set, Union


# ---------------------------------------------------------------------------
# Type enumeration
# ---------------------------------------------------------------------------

class DocValueType(Enum):
    """Supported doc-value storage types."""

    NUMERIC = auto()
    BINARY = auto()
    SORTED = auto()
    SORTED_SET = auto()


# ---------------------------------------------------------------------------
# Numeric doc values
# ---------------------------------------------------------------------------

class NumericDocValues:
    """Column of numeric (int | float) values keyed by document ID.

    Provides O(1) get/set and efficient sorting and range queries.

    Usage::

        nv = NumericDocValues()
        nv.set("doc1", 42)
        nv.set("doc2", 17)
        nv.sort_docs(["doc1", "doc2"])  # ["doc2", "doc1"]
    """

    def __init__(self) -> None:
        self._lock = threading.RLock()
        self._data: Dict[str, Union[int, float]] = {}

    @property
    def type(self) -> DocValueType:
        return DocValueType.NUMERIC

    def set(self, doc_id: str, value: Union[int, float]) -> None:
        """Store a numeric value for *doc_id*."""
        with self._lock:
            self._data[doc_id] = value

    def get(self, doc_id: str) -> Optional[Union[int, float]]:
        """Return the value for *doc_id*, or ``None`` if absent."""
        with self._lock:
            return self._data.get(doc_id)

    def delete(self, doc_id: str) -> bool:
        """Remove a document's value.  Returns True if it existed."""
        with self._lock:
            return self._data.pop(doc_id, None) is not None

    def sort_docs(
        self, doc_ids: List[str], reverse: bool = False
    ) -> List[str]:
        """Sort *doc_ids* by their stored numeric value.

        Documents without a stored value are placed at the end.

        Args:
            doc_ids:  Document IDs to sort.
            reverse:  If True, sort descending.

        Returns:
            A new list of document IDs in sorted order.
        """
        sentinel = float("inf") if not reverse else float("-inf")
        with self._lock:
            return sorted(
                doc_ids,
                key=lambda d: self._data.get(d, sentinel),
                reverse=reverse,
            )

    def range_query(
        self,
        min_val: Union[int, float, None] = None,
        max_val: Union[int, float, None] = None,
    ) -> Set[str]:
        """Return doc IDs whose value falls within [min_val, max_val].

        Either bound may be ``None`` to indicate unbounded.

        Args:
            min_val: Inclusive lower bound (or None for no lower bound).
            max_val: Inclusive upper bound (or None for no upper bound).

        Returns:
            Set of matching document IDs.
        """
        with self._lock:
            result: Set[str] = set()
            for doc_id, val in self._data.items():
                if min_val is not None and val < min_val:
                    continue
                if max_val is not None and val > max_val:
                    continue
                result.add(doc_id)
            return result

    def doc_count(self) -> int:
        """Number of documents with a stored value."""
        with self._lock:
            return len(self._data)

    def clear(self) -> None:
        """Remove all stored values."""
        with self._lock:
            self._data.clear()


# ---------------------------------------------------------------------------
# Binary doc values
# ---------------------------------------------------------------------------

class BinaryDocValues:
    """Column of arbitrary ``bytes`` values keyed by document ID.

    Usage::

        bv = BinaryDocValues()
        bv.set("doc1", b"\\x00\\x01\\x02")
        bv.get("doc1")  # b"\\x00\\x01\\x02"
    """

    def __init__(self) -> None:
        self._lock = threading.RLock()
        self._data: Dict[str, bytes] = {}

    @property
    def type(self) -> DocValueType:
        return DocValueType.BINARY

    def set(self, doc_id: str, value: bytes) -> None:
        """Store a binary value for *doc_id*."""
        with self._lock:
            self._data[doc_id] = value

    def get(self, doc_id: str) -> Optional[bytes]:
        """Return the value for *doc_id*, or ``None`` if absent."""
        with self._lock:
            return self._data.get(doc_id)

    def delete(self, doc_id: str) -> bool:
        """Remove a document's value.  Returns True if it existed."""
        with self._lock:
            return self._data.pop(doc_id, None) is not None

    def doc_count(self) -> int:
        """Number of documents with a stored value."""
        with self._lock:
            return len(self._data)

    def clear(self) -> None:
        """Remove all stored values."""
        with self._lock:
            self._data.clear()


# ---------------------------------------------------------------------------
# Sorted doc values
# ---------------------------------------------------------------------------

class SortedDocValues:
    """Column of string values with ordinal encoding.

    Each unique string is assigned an ordinal based on its position in
    sorted order.  This allows O(1) ordinal lookup and efficient sorting.

    Usage::

        sv = SortedDocValues()
        sv.set("doc1", "banana")
        sv.set("doc2", "apple")
        sv.ordinal("doc2")  # 0  (apple < banana)
    """

    def __init__(self) -> None:
        self._lock = threading.RLock()
        self._data: Dict[str, str] = {}
        # Sorted list of unique values for ordinal computation.
        self._sorted_vals: List[str] = []
        self._val_set: Set[str] = set()
        self._ordinals_dirty = False

    @property
    def type(self) -> DocValueType:
        return DocValueType.SORTED

    def _rebuild_ordinals(self) -> None:
        """Rebuild the sorted value list when the value set has changed."""
        if self._ordinals_dirty:
            self._sorted_vals = sorted(self._val_set)
            self._ordinals_dirty = False

    def set(self, doc_id: str, value: str) -> None:
        """Store a string value for *doc_id*."""
        with self._lock:
            old = self._data.get(doc_id)
            self._data[doc_id] = value
            if value not in self._val_set:
                self._val_set.add(value)
                self._ordinals_dirty = True
            # If old value is no longer referenced, remove it.
            if old is not None and old != value:
                if old not in self._data.values():
                    self._val_set.discard(old)
                    self._ordinals_dirty = True

    def get(self, doc_id: str) -> Optional[str]:
        """Return the value for *doc_id*, or ``None`` if absent."""
        with self._lock:
            return self._data.get(doc_id)

    def delete(self, doc_id: str) -> bool:
        """Remove a document's value.  Returns True if it existed."""
        with self._lock:
            old = self._data.pop(doc_id, None)
            if old is None:
                return False
            if old not in self._data.values():
                self._val_set.discard(old)
                self._ordinals_dirty = True
            return True

    def ordinal(self, doc_id: str) -> int:
        """Return the ordinal of the value for *doc_id*.

        The ordinal is the zero-based index of the value in the
        globally sorted list of all unique values.

        Args:
            doc_id: The document identifier.

        Returns:
            Ordinal index (>= 0).

        Raises:
            KeyError: If *doc_id* has no stored value.
        """
        with self._lock:
            val = self._data.get(doc_id)
            if val is None:
                raise KeyError(f"No value for doc {doc_id!r}")
            self._rebuild_ordinals()
            return bisect.bisect_left(self._sorted_vals, val)

    def sort_docs(
        self, doc_ids: List[str], reverse: bool = False
    ) -> List[str]:
        """Sort *doc_ids* by their stored string value.

        Documents without a stored value are placed at the end.

        Args:
            doc_ids:  Document IDs to sort.
            reverse:  If True, sort descending.

        Returns:
            A new list of document IDs in sorted order.
        """
        with self._lock:
            self._rebuild_ordinals()
            max_ord = len(self._sorted_vals)

            def _key(d: str) -> int:
                val = self._data.get(d)
                if val is None:
                    return max_ord  # push to end
                return bisect.bisect_left(self._sorted_vals, val)

            return sorted(doc_ids, key=_key, reverse=reverse)

    def doc_count(self) -> int:
        """Number of documents with a stored value."""
        with self._lock:
            return len(self._data)

    def value_count(self) -> int:
        """Number of unique string values."""
        with self._lock:
            return len(self._val_set)

    def clear(self) -> None:
        """Remove all stored values."""
        with self._lock:
            self._data.clear()
            self._sorted_vals.clear()
            self._val_set.clear()
            self._ordinals_dirty = False


# ---------------------------------------------------------------------------
# Sorted-set doc values
# ---------------------------------------------------------------------------

class SortedSetDocValues:
    """Column of multi-valued string sets keyed by document ID.

    Each document can have zero or more string values, stored as a
    frozenset.  Useful for faceting and multi-valued fields.

    Usage::

        ssv = SortedSetDocValues()
        ssv.add("doc1", "red")
        ssv.add("doc1", "blue")
        ssv.get("doc1")  # frozenset({"red", "blue"})
        ssv.docs_with_value("red")  # {"doc1"}
    """

    def __init__(self) -> None:
        self._lock = threading.RLock()
        self._data: Dict[str, Set[str]] = defaultdict(set)
        # Inverted lookup: value -> set of doc_ids
        self._value_index: Dict[str, Set[str]] = defaultdict(set)

    @property
    def type(self) -> DocValueType:
        return DocValueType.SORTED_SET

    def add(self, doc_id: str, value: str) -> None:
        """Add a value to the document's value set."""
        with self._lock:
            self._data[doc_id].add(value)
            self._value_index[value].add(doc_id)

    def get(self, doc_id: str) -> FrozenSet[str]:
        """Return the set of values for *doc_id* (empty frozenset if absent)."""
        with self._lock:
            vals = self._data.get(doc_id)
            if vals is None:
                return frozenset()
            return frozenset(vals)

    def remove(self, doc_id: str, value: str) -> bool:
        """Remove a single value from a document's set.

        Returns True if the value was present.
        """
        with self._lock:
            vals = self._data.get(doc_id)
            if vals is None or value not in vals:
                return False
            vals.discard(value)
            self._value_index[value].discard(doc_id)
            if not self._value_index[value]:
                del self._value_index[value]
            if not vals:
                del self._data[doc_id]
            return True

    def delete(self, doc_id: str) -> bool:
        """Remove all values for a document.  Returns True if it existed."""
        with self._lock:
            vals = self._data.pop(doc_id, None)
            if vals is None:
                return False
            for v in vals:
                self._value_index[v].discard(doc_id)
                if not self._value_index[v]:
                    del self._value_index[v]
            return True

    def docs_with_value(self, value: str) -> Set[str]:
        """Return all document IDs that contain *value*.

        Returns a copy of the internal set.
        """
        with self._lock:
            return set(self._value_index.get(value, set()))

    def doc_count(self) -> int:
        """Number of documents with at least one value."""
        with self._lock:
            return len(self._data)

    def value_count(self) -> int:
        """Number of unique values across all documents."""
        with self._lock:
            return len(self._value_index)

    def clear(self) -> None:
        """Remove all stored values."""
        with self._lock:
            self._data.clear()
            self._value_index.clear()


# ---------------------------------------------------------------------------
# Doc values store (multi-field manager)
# ---------------------------------------------------------------------------

_DocValuesField = Union[
    NumericDocValues, BinaryDocValues, SortedDocValues, SortedSetDocValues
]


class DocValuesStore:
    """Manages multiple named doc-value fields.

    Each field has a specific type (numeric, binary, sorted, sorted-set)
    and can be retrieved by name.  Thread-safe.

    Usage::

        store = DocValuesStore()
        price = store.add_numeric("price")
        price.set("doc1", 9.99)

        category = store.add_sorted_set("category")
        category.add("doc1", "electronics")

        store.get_field("price").get("doc1")  # 9.99
    """

    def __init__(self) -> None:
        self._lock = threading.RLock()
        self._fields: Dict[str, _DocValuesField] = {}

    def _add_field(self, name: str, field: _DocValuesField) -> _DocValuesField:
        """Register a field, raising if the name is already taken."""
        with self._lock:
            if name in self._fields:
                raise ValueError(f"Field {name!r} already exists")
            self._fields[name] = field
            return field

    def add_numeric(self, field: str) -> NumericDocValues:
        """Create and register a numeric doc-value field.

        Args:
            field: The field name.

        Returns:
            The new :class:`NumericDocValues` instance.

        Raises:
            ValueError: If a field with this name already exists.
        """
        dv = NumericDocValues()
        self._add_field(field, dv)
        return dv

    def add_binary(self, field: str) -> BinaryDocValues:
        """Create and register a binary doc-value field.

        Args:
            field: The field name.

        Returns:
            The new :class:`BinaryDocValues` instance.

        Raises:
            ValueError: If a field with this name already exists.
        """
        dv = BinaryDocValues()
        self._add_field(field, dv)
        return dv

    def add_sorted(self, field: str) -> SortedDocValues:
        """Create and register a sorted doc-value field.

        Args:
            field: The field name.

        Returns:
            The new :class:`SortedDocValues` instance.

        Raises:
            ValueError: If a field with this name already exists.
        """
        dv = SortedDocValues()
        self._add_field(field, dv)
        return dv

    def add_sorted_set(self, field: str) -> SortedSetDocValues:
        """Create and register a sorted-set doc-value field.

        Args:
            field: The field name.

        Returns:
            The new :class:`SortedSetDocValues` instance.

        Raises:
            ValueError: If a field with this name already exists.
        """
        dv = SortedSetDocValues()
        self._add_field(field, dv)
        return dv

    def get_field(self, field: str) -> Optional[_DocValuesField]:
        """Return the doc-value object for *field*, or ``None``."""
        with self._lock:
            return self._fields.get(field)

    def has_field(self, field: str) -> bool:
        """Check whether *field* is registered."""
        with self._lock:
            return field in self._fields

    def field_names(self) -> List[str]:
        """Return a sorted list of all registered field names."""
        with self._lock:
            return sorted(self._fields.keys())

    def remove_field(self, field: str) -> bool:
        """Remove a field from the store.  Returns True if it existed."""
        with self._lock:
            return self._fields.pop(field, None) is not None

    def clear(self) -> None:
        """Remove all fields."""
        with self._lock:
            self._fields.clear()
