"""Cosine-similarity vector index.

FAISS `IndexFlatIP` wrapped in `IndexIDMap2` when faiss is installed, plain
numpy otherwise — same interface, same results. Choices that fix the old
codebase's failure modes:

- Exact inner-product search (vectors are unit-norm, so IP == cosine).
  No IVF: nothing to train, no wrong-metric trap, and exact search is
  microseconds up to hundreds of thousands of vectors.
- `-1` ids from an empty index are filtered, never used as list indices.
- The index is rebuilt from SQLite at startup (SQLite is the source of
  truth), so index and metadata can never drift apart.
"""
from __future__ import annotations

import logging

import numpy as np

log = logging.getLogger(__name__)

try:
    import faiss  # type: ignore
    _HAVE_FAISS = True
except ImportError:  # pragma: no cover - environment dependent
    faiss = None
    _HAVE_FAISS = False


class VectorIndex:
    def __init__(self, dim: int):
        self.dim = dim
        if _HAVE_FAISS:
            self._index = faiss.IndexIDMap2(faiss.IndexFlatIP(dim))
            self._ids = None
            self._vecs = None
        else:
            log.warning("faiss not installed - using exact numpy search "
                        "(identical results, slower at large scale)")
            self._index = None
            self._ids = np.empty((0,), dtype=np.int64)
            self._vecs = np.empty((0, dim), dtype=np.float32)

    def __len__(self) -> int:
        if self._index is not None:
            return self._index.ntotal
        return len(self._ids)

    def add(self, ids: "list[int]", vectors: np.ndarray) -> None:
        if len(ids) == 0:
            return
        vectors = np.ascontiguousarray(vectors, dtype=np.float32).reshape(len(ids), self.dim)
        id_arr = np.asarray(ids, dtype=np.int64)
        if self._index is not None:
            self._index.add_with_ids(vectors, id_arr)
        else:
            self._ids = np.concatenate([self._ids, id_arr])
            self._vecs = np.vstack([self._vecs, vectors])

    def remove(self, ids: "list[int]") -> None:
        if len(ids) == 0:
            return
        id_arr = np.asarray(ids, dtype=np.int64)
        if self._index is not None:
            self._index.remove_ids(id_arr)
        else:
            keep = ~np.isin(self._ids, id_arr)
            self._ids = self._ids[keep]
            self._vecs = self._vecs[keep]

    def search(self, vector: np.ndarray, k: int = 1) -> "list[tuple[int, float]]":
        """Top-k (embedding_id, cosine_similarity), best first."""
        if len(self) == 0:
            return []
        q = np.ascontiguousarray(vector, dtype=np.float32).reshape(1, self.dim)
        k = min(k, len(self))
        if self._index is not None:
            scores, ids = self._index.search(q, k)
            return [(int(i), float(s))
                    for i, s in zip(ids[0], scores[0]) if i != -1]
        sims = self._vecs @ q[0]
        order = np.argsort(-sims)[:k]
        return [(int(self._ids[i]), float(sims[i])) for i in order]