"""Generic FAISS + DuckDB/Parquet nearest-neighbor store, mounted read-only
(hfmount). FAISS is mmap'd; the parquet table is never fully loaded --
DuckDB queries it directly, pushing id filters down to row-group pruning.

Two id modes, matching how an index was built:
  id_column=None  -- plain index, vector i <-> parquet row i (positional).
                     Table is small enough to load once (e.g. the ~10MB
                     annotation reference); searched ids ARE row positions.
  id_column=<col> -- IndexIDMap2 built with add_with_ids(vecs, <col> values)
                     (e.g. genome_id/gene_id). Searched ids are real key
                     values; fetch_rows does a WHERE <col> IN (...) lookup
                     with no full-table load, for tables too big to load.

One VectorStore is shared (lazy_singleton) across every concurrent request,
so every query goes through self._con.cursor() rather than self._con
directly -- a bare DuckDB connection isn't safe for concurrent queries from
multiple threads (confirmed empirically: concurrent .execute() calls on one
connection returned wrong/None results on ~20% of queries, silently, no
exception); a fresh cursor per call shares the connection but is safe.
"""
import duckdb
import faiss
import numpy as np

from utils.retry import retry


class VectorStore:
    def __init__(self, index_path, parquet_path, id_column=None, top_k=16, preload=False):
        """preload=True reads the whole index into process RAM up front
        (one slower sequential read at construction) instead of mmap'ing it
        (lazy, on-demand network fetch per touched page on every query) --
        worth it for an index small enough to fit comfortably in RAM, to
        make every later query immune to the mount's latency/flakiness."""
        self.parquet_path = parquet_path
        self.id_column = id_column
        self.top_k = top_k
        flags = 0 if preload else faiss.IO_FLAG_MMAP
        self.index = retry(lambda: faiss.read_index(index_path, flags), label=f"read_index {index_path}")
        if hasattr(self.index, "set_direct_map_type"):
            # IVF indexes need this before .reconstruct(id) works at all --
            # the default (Array) direct map assumes dense sequential ids,
            # which ours (genome_id/gene_id) aren't; Hashtable handles
            # arbitrary ids. Flat indexes have no such method/need.
            self.index.set_direct_map_type(faiss.DirectMap.Hashtable)
        self._normalize = (self.index.metric_type == faiss.METRIC_INNER_PRODUCT)
        self._con = duckdb.connect()

        if id_column is None:
            self.table = retry(
                lambda: self._con.cursor().execute("SELECT * FROM read_parquet(?)", [parquet_path]).df(),
                label=f"load {parquet_path}")
            n_rows = len(self.table)
        else:
            self.table = None
            n_rows = retry(
                lambda: self._con.cursor().execute("SELECT count(*) FROM read_parquet(?)", [parquet_path]).fetchone()[0],
                label=f"count {parquet_path}")

        if n_rows != self.index.ntotal:
            raise RuntimeError(
                f"row-count mismatch: {parquet_path} has {n_rows} rows, "
                f"{index_path} has {self.index.ntotal} vectors."
            )

    def search(self, vectors, k=None):
        """vectors: [N, D] array-like. Returns (scores, ids), each [N, k] --
        ids are row positions (id_column=None) or key values (id_column set);
        -1 pads rows with fewer than k hits."""
        vecs = np.asarray(vectors, dtype=np.float32)
        if self._normalize:
            norms = np.linalg.norm(vecs, axis=1, keepdims=True)
            vecs = vecs / np.clip(norms, 1e-12, None)
        return self.index.search(vecs, k or self.top_k)

    def fetch_rows(self, ids):
        """ids: iterable of row positions or key values (matching id_column
        mode), -1 entries ignored. Returns {id: row_dict}."""
        ids = sorted({int(i) for i in ids if i >= 0})
        if not ids:
            return {}
        if self.id_column is None:
            return {i: self.table.iloc[i].to_dict() for i in ids}
        placeholders = ",".join("?" * len(ids))
        query = f"SELECT * FROM read_parquet(?) WHERE {self.id_column} IN ({placeholders})"

        def _run():
            cur = self._con.cursor()
            cur.execute(query, [self.parquet_path] + ids)
            return cur.fetchall(), [d[0] for d in cur.description]

        rows, cols = retry(_run, label=f"fetch_rows {self.parquet_path}")
        return {row[cols.index(self.id_column)]: dict(zip(cols, row)) for row in rows}

    def search_and_fetch(self, vectors, k=None):
        """One search + one batched fetch_rows for the whole query batch.
        Returns a list of per-query hit lists: [{**row, "score": float}, ...]."""
        scores, ids = self.search(vectors, k)
        rows = self.fetch_rows(ids.ravel())
        results = []
        for s_row, i_row in zip(scores, ids):
            results.append([
                {**rows[int(i)], "score": float(s)}
                for s, i in zip(s_row, i_row) if i >= 0 and int(i) in rows
            ])
        return results