"""
RBSP Unified Vector Store.

Combines HNSW approximate nearest neighbour search with binary quantization
for a multi-mode dense vector store.

Four search modes:
    - ``"exact"``:   Brute-force linear scan.  Perfect recall, O(n).
    - ``"hnsw"``:    Approximate NN via HNSW graph.  O(log n).
    - ``"binary"``:  Hamming distance on binary codes.  Ultra-fast, lower recall.
    - ``"hybrid"``:  Binary pre-filter (top 4K) -> exact rerank.  Fast + high recall.

All modes support optional metadata filtering via ``FilterExpression``.

Thread-safe via ``threading.RLock``.  Zero external dependencies.
"""

from __future__ import annotations

import heapq
import math
import threading
from typing import Any, Callable, Dict, List, Optional, Sequence, Set, Tuple, Union

from rbsp.core.hnsw import HNSWConfig, HNSWIndex, cosine_distance, l2_distance
from rbsp.core.quantize import BinaryQuantizer


# ---------------------------------------------------------------------------
# Filter Expressions
# ---------------------------------------------------------------------------


class FilterExpression:
    """Base class for metadata filter expressions.

    Filter expressions form a composable DSL for matching vector metadata.
    They can be combined with ``&`` (and), ``|`` (or), and ``~`` (not).

    Example::

        f = FieldCondition("category", "eq", "electronics")
        g = FieldCondition("price", "lte", 100.0)
        expr = f & g  # Both must match
    """

    def match(self, metadata: Dict[str, Any]) -> bool:
        """Return ``True`` if *metadata* satisfies this expression.

        Args:
            metadata: The metadata dictionary to test.

        Returns:
            Boolean result.
        """
        raise NotImplementedError

    def __and__(self, other: FilterExpression) -> And:
        return And(self, other)

    def __or__(self, other: FilterExpression) -> Or:
        return Or(self, other)

    def __invert__(self) -> Not:
        return Not(self)


class FieldCondition(FilterExpression):
    """A condition on a single metadata field.

    Supported operators:
        - ``"eq"``:       Equal
        - ``"ne"``:       Not equal
        - ``"gt"``:       Greater than
        - ``"gte"``:      Greater than or equal
        - ``"lt"``:       Less than
        - ``"lte"``:      Less than or equal
        - ``"in"``:       Value is in the given collection
        - ``"contains"``: Field value (a collection) contains the given value

    Args:
        field:    Metadata field name.
        operator: Comparison operator string.
        value:    Value to compare against.
    """

    __slots__ = ("field", "operator", "value")

    _VALID_OPS = ("eq", "ne", "gt", "gte", "lt", "lte", "in", "contains")

    def __init__(self, field: str, operator: str, value: Any) -> None:
        if operator not in self._VALID_OPS:
            raise ValueError(
                f"operator must be one of {self._VALID_OPS}, got {operator!r}"
            )
        self.field = field
        self.operator = operator
        self.value = value

    def match(self, metadata: Dict[str, Any]) -> bool:
        """Return ``True`` if the field condition is satisfied."""
        val = metadata.get(self.field)
        if val is None:
            return False
        op = self.operator
        if op == "eq":
            return val == self.value
        if op == "ne":
            return val != self.value
        if op == "gt":
            return val > self.value
        if op == "gte":
            return val >= self.value
        if op == "lt":
            return val < self.value
        if op == "lte":
            return val <= self.value
        if op == "in":
            return val in self.value
        if op == "contains":
            try:
                return self.value in val
            except TypeError:
                return False
        return False  # pragma: no cover

    def __repr__(self) -> str:
        return f"FieldCondition({self.field!r}, {self.operator!r}, {self.value!r})"


class And(FilterExpression):
    """Logical AND of two filter expressions.

    Args:
        left:  First expression.
        right: Second expression.
    """

    __slots__ = ("left", "right")

    def __init__(self, left: FilterExpression, right: FilterExpression) -> None:
        self.left = left
        self.right = right

    def match(self, metadata: Dict[str, Any]) -> bool:
        return self.left.match(metadata) and self.right.match(metadata)

    def __repr__(self) -> str:
        return f"And({self.left!r}, {self.right!r})"


class Or(FilterExpression):
    """Logical OR of two filter expressions.

    Args:
        left:  First expression.
        right: Second expression.
    """

    __slots__ = ("left", "right")

    def __init__(self, left: FilterExpression, right: FilterExpression) -> None:
        self.left = left
        self.right = right

    def match(self, metadata: Dict[str, Any]) -> bool:
        return self.left.match(metadata) or self.right.match(metadata)

    def __repr__(self) -> str:
        return f"Or({self.left!r}, {self.right!r})"


class Not(FilterExpression):
    """Logical NOT of a filter expression.

    Args:
        expr: Expression to negate.
    """

    __slots__ = ("expr",)

    def __init__(self, expr: FilterExpression) -> None:
        self.expr = expr

    def match(self, metadata: Dict[str, Any]) -> bool:
        return not self.expr.match(metadata)

    def __repr__(self) -> str:
        return f"Not({self.expr!r})"


# ---------------------------------------------------------------------------
# Vector Store
# ---------------------------------------------------------------------------


class VectorStore:
    """Unified dense vector store with multiple search modes.

    Maintains:
        - A raw vector dictionary (id -> vector) for exact search.
        - An HNSW index for approximate NN search.
        - A binary quantizer for ultra-fast Hamming distance search.
        - A metadata dictionary (id -> dict) for filtered search.

    Args:
        dim:       Vector dimensionality (inferred on first insert if None).
        distance:  Distance metric -- ``"cosine"`` or ``"l2"`` (default ``"cosine"``).
        hnsw_config: Optional HNSW configuration override.
    """

    _VALID_MODES = ("exact", "hnsw", "binary", "hybrid")
    _FILTER_EF_MULTIPLIER = 4

    def __init__(
        self,
        dim: Optional[int] = None,
        *,
        distance: str = "cosine",
        hnsw_config: Optional[HNSWConfig] = None,
    ) -> None:
        self._dim = dim
        self._distance = distance
        self._lock = threading.RLock()

        # Raw vectors
        self._vectors: Dict[str, List[float]] = {}

        # Metadata store: id -> arbitrary dict
        self._metadata: Dict[str, Dict[str, Any]] = {}

        # HNSW index
        if hnsw_config is None:
            hnsw_config = HNSWConfig(distance=distance)
        self._hnsw = HNSWIndex(hnsw_config)

        # Binary quantizer
        self._quantizer = BinaryQuantizer(dim)
        self._quantizer_fitted = False

        # Distance function for exact search
        if distance == "cosine":
            self._dist_fn = cosine_distance
        else:
            self._dist_fn = l2_distance

    # -- Properties -----------------------------------------------------------

    @property
    def dim(self) -> Optional[int]:
        """Vector dimensionality (``None`` until first insert)."""
        return self._dim

    @property
    def size(self) -> int:
        """Number of stored vectors."""
        with self._lock:
            return len(self._vectors)

    # -- Insert ---------------------------------------------------------------

    def insert(
        self,
        id: str,
        vector: Sequence[float],
        metadata: Optional[Dict[str, Any]] = None,
    ) -> None:
        """Insert or update a vector with optional metadata.

        Args:
            id:       Unique string identifier.
            vector:   Float vector.
            metadata: Optional dict of arbitrary key-value metadata.
        """
        vec = list(vector)
        with self._lock:
            if self._dim is None:
                self._dim = len(vec)
            self._vectors[id] = vec
            if metadata is not None:
                self._metadata[id] = dict(metadata)
            elif id not in self._metadata:
                self._metadata[id] = {}
            self._hnsw.insert(id, vec)
            self._quantizer.insert(id, vec)

    def batch_insert(
        self,
        ids: List[str],
        vectors: List[Sequence[float]],
        metadata_list: Optional[List[Optional[Dict[str, Any]]]] = None,
    ) -> None:
        """Bulk-insert multiple vectors with optional metadata.

        Args:
            ids:           List of unique string identifiers.
            vectors:       List of float vectors (same length as *ids*).
            metadata_list: Optional list of metadata dicts (same length as *ids*).
                           Individual entries may be ``None``.

        Raises:
            ValueError: If *ids* and *vectors* differ in length, or if
                        *metadata_list* is given and differs in length.
        """
        if len(ids) != len(vectors):
            raise ValueError(
                f"ids ({len(ids)}) and vectors ({len(vectors)}) must have "
                f"the same length"
            )
        if metadata_list is not None and len(metadata_list) != len(ids):
            raise ValueError(
                f"metadata_list ({len(metadata_list)}) must match ids ({len(ids)})"
            )
        with self._lock:
            for i, (id_, vec) in enumerate(zip(ids, vectors)):
                v = list(vec)
                if self._dim is None:
                    self._dim = len(v)
                self._vectors[id_] = v
                if metadata_list is not None and metadata_list[i] is not None:
                    self._metadata[id_] = dict(metadata_list[i])
                elif id_ not in self._metadata:
                    self._metadata[id_] = {}
                self._hnsw.insert(id_, v)
                self._quantizer.insert(id_, v)

    # -- Metadata access ------------------------------------------------------

    def set_metadata(self, id: str, metadata: Dict[str, Any]) -> None:
        """Replace the metadata for *id*.

        Args:
            id:       Vector identifier (must already exist).
            metadata: New metadata dict.

        Raises:
            KeyError: If *id* is not in the store.
        """
        with self._lock:
            if id not in self._vectors:
                raise KeyError(f"Vector {id!r} not found")
            self._metadata[id] = dict(metadata)

    def update_metadata(self, id: str, updates: Dict[str, Any]) -> None:
        """Merge *updates* into the existing metadata for *id*.

        Args:
            id:      Vector identifier (must already exist).
            updates: Key-value pairs to merge.

        Raises:
            KeyError: If *id* is not in the store.
        """
        with self._lock:
            if id not in self._vectors:
                raise KeyError(f"Vector {id!r} not found")
            meta = self._metadata.get(id)
            if meta is None:
                self._metadata[id] = dict(updates)
            else:
                meta.update(updates)

    def get_metadata(self, id: str) -> Optional[Dict[str, Any]]:
        """Return the metadata dict for *id*, or ``None`` if absent.

        Args:
            id: Vector identifier.

        Returns:
            A copy of the metadata dict, or ``None``.
        """
        with self._lock:
            meta = self._metadata.get(id)
            return dict(meta) if meta is not None else None

    # -- Remove ---------------------------------------------------------------

    def remove(self, id: str) -> bool:
        """Remove a vector by ID.

        Returns:
            ``True`` if found and removed, ``False`` otherwise.
        """
        with self._lock:
            if id not in self._vectors:
                return False
            del self._vectors[id]
            self._metadata.pop(id, None)
            self._hnsw.remove(id)
            # BinaryQuantizer doesn't have a remove method; rebuild codes dict
            self._quantizer._codes.pop(id, None)
            return True

    # -- Search ---------------------------------------------------------------

    def _build_filter_fn(
        self, filter_expr: Optional[FilterExpression]
    ) -> Optional[Callable[[str], bool]]:
        """Convert a ``FilterExpression`` into a node-ID predicate."""
        if filter_expr is None:
            return None

        def _filter(node_id: str) -> bool:
            meta = self._metadata.get(node_id)
            if meta is None:
                return False
            return filter_expr.match(meta)

        return _filter

    def search(
        self,
        query: Sequence[float],
        k: int = 10,
        *,
        mode: str = "hnsw",
        filter_expr: Optional[FilterExpression] = None,
    ) -> List[Tuple[str, float]]:
        """Search for the *k* nearest vectors, optionally filtered.

        Args:
            query:       Query float vector.
            k:           Number of results.
            mode:        Search mode -- ``"exact"``, ``"hnsw"``, ``"binary"``,
                         or ``"hybrid"`` (default ``"hnsw"``).
            filter_expr: Optional ``FilterExpression`` to restrict results to
                         vectors whose metadata satisfies the expression.

        Returns:
            List of ``(id, distance)`` tuples sorted by ascending distance.

        Raises:
            ValueError: If *mode* is not recognised.
        """
        if mode not in self._VALID_MODES:
            raise ValueError(
                f"mode must be one of {self._VALID_MODES}, got {mode!r}"
            )

        q = list(query)
        with self._lock:
            filter_fn = self._build_filter_fn(filter_expr)
            if mode == "exact":
                return self._search_exact(q, k, filter_fn)
            if mode == "hnsw":
                return self._search_hnsw(q, k, filter_fn)
            if mode == "binary":
                return self._search_binary(q, k, filter_fn)
            # hybrid
            return self._search_hybrid(q, k, filter_fn)

    def _search_exact(
        self,
        query: List[float],
        k: int,
        filter_fn: Optional[Callable[[str], bool]] = None,
    ) -> List[Tuple[str, float]]:
        """Brute-force exact nearest neighbour search."""
        if not self._vectors:
            return []
        scored = []
        for id_, vec in self._vectors.items():
            if filter_fn is not None and not filter_fn(id_):
                continue
            scored.append((id_, self._dist_fn(query, vec)))
        return heapq.nsmallest(k, scored, key=lambda x: x[1])

    def _search_hnsw(
        self,
        query: List[float],
        k: int,
        filter_fn: Optional[Callable[[str], bool]] = None,
    ) -> List[Tuple[str, float]]:
        """HNSW approximate nearest neighbour search.

        When a filter is active and the initial search returns fewer than
        *k* results, an adaptive 3-tier fallback is used:

        1. Normal HNSW search (with in-graph filtering).
        2. Retry with ``ef_search`` widened by ``_FILTER_EF_MULTIPLIER``.
        3. Fall back to brute-force exact filtered search.
        """
        results = self._hnsw.search(query, k, filter_fn=filter_fn)

        if filter_fn is not None and len(results) < k:
            # Tier 2: widen ef_search and retry
            original_ef = self._hnsw._config.ef_search
            widened_ef = original_ef * self._FILTER_EF_MULTIPLIER
            self._hnsw._config.ef_search = widened_ef
            try:
                results = self._hnsw.search(query, k, filter_fn=filter_fn)
            finally:
                self._hnsw._config.ef_search = original_ef

            # Tier 3: ultimate fallback to exact filtered search
            if len(results) < k:
                return self._search_exact(query, k, filter_fn)

        return results

    def _search_binary(
        self,
        query: List[float],
        k: int,
        filter_fn: Optional[Callable[[str], bool]] = None,
    ) -> List[Tuple[str, float]]:
        """Binary quantization Hamming distance search.

        Returns Hamming distance as the distance value (int cast to float).
        When filtering, retrieves extra candidates and filters post-hoc.
        """
        if filter_fn is None:
            ham_results = self._quantizer.search(query, k)
            return [(id_, float(d)) for id_, d in ham_results]

        # Over-fetch then filter
        fetch_k = min(k * 10, len(self._quantizer._codes))
        ham_results = self._quantizer.search(query, fetch_k)
        filtered: List[Tuple[str, float]] = []
        for id_, d in ham_results:
            if filter_fn(id_):
                filtered.append((id_, float(d)))
                if len(filtered) >= k:
                    break
        return filtered

    def _search_hybrid(
        self,
        query: List[float],
        k: int,
        filter_fn: Optional[Callable[[str], bool]] = None,
    ) -> List[Tuple[str, float]]:
        """Hybrid search: binary pre-filter then exact rerank.

        1. Binary quantization retrieves top 4K candidates by Hamming distance.
        2. Exact distance reranks those candidates.
        """
        if not self._vectors:
            return []

        # Pre-filter: retrieve up to 4K candidates via binary codes
        pre_k = min(4096, len(self._vectors))
        ham_results = self._quantizer.search(query, pre_k)

        # Rerank with exact distance
        reranked: List[Tuple[str, float]] = []
        for id_, _ham_dist in ham_results:
            if filter_fn is not None and not filter_fn(id_):
                continue
            vec = self._vectors.get(id_)
            if vec is not None:
                d = self._dist_fn(query, vec)
                reranked.append((id_, d))

        return heapq.nsmallest(k, reranked, key=lambda x: x[1])

    # -- Quantizer management -------------------------------------------------

    def rebuild_quantizer(self, sample_size: Optional[int] = None) -> None:
        """Refit the binary quantizer thresholds from current vectors.

        Args:
            sample_size: Max vectors to use for fitting.  If ``None``, uses
                         all vectors.
        """
        with self._lock:
            if not self._vectors:
                return

            vecs = list(self._vectors.values())
            if sample_size is not None and sample_size < len(vecs):
                import random
                vecs = random.sample(vecs, sample_size)

            self._quantizer.fit(vecs)
            self._quantizer_fitted = True

            # Re-quantize all stored vectors with new thresholds
            for id_, vec in self._vectors.items():
                self._quantizer.insert(id_, vec)

    # -- Utility --------------------------------------------------------------

    def __len__(self) -> int:
        with self._lock:
            return len(self._vectors)

    def __contains__(self, id: str) -> bool:  # type: ignore[override]
        with self._lock:
            return id in self._vectors

    def __repr__(self) -> str:
        return (
            f"VectorStore(size={len(self._vectors)}, "
            f"dim={self._dim}, distance={self._distance!r})"
        )

    def get_vector(self, id: str) -> Optional[List[float]]:
        """Retrieve the raw vector for *id*, or ``None``."""
        with self._lock:
            vec = self._vectors.get(id)
            return list(vec) if vec is not None else None
