"""
# flake8: noqa
faiss_index.py
---------------
Lightweight vector index for fast similarity search.
"""

from __future__ import annotations

import json
import logging
from pathlib import Path
from typing import Any, Iterable, List, Tuple

import numpy as np

try:  # pragma: no cover - optional dependency
    import faiss

    _HAS_FAISS = True
except Exception:  # pragma: no cover - fallback when faiss missing
    faiss = None
    _HAS_FAISS = False


class FaissIndex:
    """A convenience wrapper around FAISS (or NumPy fallback)."""

    def __init__(
        self,
        dim: int,
        metric: str = "ip",
        normalize: bool = True,
        debug: bool = False,
        *,
        index_type: str = "ivfflat",
        nlist: int = 1024,
        nprobe: int = 10,
        hnsw_m: int = 32,
        hnsw_ef_search: int = 64,
        hnsw_ef_construction: int = 200,
        index: Any | None = None,
    ) -> None:
        self.dim = int(dim)
        self.metric = metric.lower()
        if self.metric not in {"ip", "l2"}:
            raise ValueError("metric must be 'ip' or 'l2'")
        self.normalize = normalize
        self.debug = debug
        self.index_type = index_type if index_type in {
            "ivfflat", "hnsw"} else "ivfflat"
        self.nlist = max(1, int(nlist))
        self.nprobe = max(1, int(nprobe))
        self.hnsw_m = max(4, int(hnsw_m))
        self.hnsw_ef_search = max(4, int(hnsw_ef_search))
        self.hnsw_ef_construction = max(16, int(hnsw_ef_construction))

        self._ids: List[str] = []
        self._vecs: List[np.ndarray] = []

        self._logger = logging.getLogger(self.__class__.__name__)
        if self.debug and not self._logger.handlers:
            handler = logging.StreamHandler()
            handler.setFormatter(logging.Formatter(
                "[%(levelname)s] %(name)s: %(message)s"))
            self._logger.addHandler(handler)
            self._logger.setLevel(logging.INFO)

        self._index: Any | None = None
        if _HAS_FAISS:
            # Build the underlying FAISS index if not provided
            self._index = index if index is not None else self._build_index()

    # ------------------------------------------------------------------ #
    # Public API
    # ------------------------------------------------------------------ #
    def add(self, episode_id: str, vec: np.ndarray) -> None:
        self.add_embeddings(
    [episode_id], np.asarray(
        vec, dtype=np.float32)[
            np.newaxis, :])

    def add_many(self, ids: Iterable[str], matrix: np.ndarray) -> None:
        self.add_embeddings(list(ids), matrix)

    def add_embeddings(self, ids: Iterable[str], matrix: np.ndarray) -> None:
        ids_list = list(ids)
        if not ids_list:
            return
        prepared = self._prep_mat(matrix)
        if prepared.shape[0] != len(ids_list):
            raise ValueError("Number of ids must match number of vectors")
        self._ids.extend(ids_list)
        self._vecs.extend(prepared)
        if _HAS_FAISS and prepared.size:
            self._ensure_trained(prepared)
            self._add_to_faiss(prepared)
        self._log(f"Added {len(ids_list)} vectors")

    def remove(self, episode_id: str) -> None:
        if episode_id in self._ids:
            idx = self._ids.index(episode_id)
            del self._ids[idx]
            del self._vecs[idx]
            if _HAS_FAISS:
                self._rebuild_faiss()
            self._log(f"Removed id={episode_id}")

    def clear(self) -> None:
        self._ids.clear()
        self._vecs.clear()
        if _HAS_FAISS:
            self._index = self._build_index()
        self._log("Cleared index")

    @property
    def ids(self) -> List[str]:
        """Return list of indexed IDs."""
        return self._ids.copy()

    def search(self, query: np.ndarray, k: int = 5) -> List[Tuple[str, float]]:
        if not self._ids:
            return []

        q = self._prep_vec(query)
        limit = min(max(1, int(k)), len(self._ids))

        if _HAS_FAISS and self._index is not None:
            self._configure_search_params()
            Q = q[np.newaxis, :]
            distances, indices = self._index.search(Q, limit)
            idxs = indices[0]
            scores = distances[0]
            out: List[Tuple[str, float]] = []
            for j, s in zip(idxs, scores):
                if j == -1:
                    continue
                score = float(s if self.metric == "ip" else -s)
                out.append((self._ids[j], score))
            return out

        # NumPy fallback
        M = np.vstack(self._vecs)  # (N, d)
        if self.metric == "ip":
            sims = M @ q
            idxs = np.argsort(-sims)[:limit]
            return [(self._ids[i], float(sims[i])) for i in idxs]
        dists = np.sum((M - q) ** 2, axis=1)
        idxs = np.argsort(dists)[:limit]
        return [(self._ids[i], float(-dists[i])) for i in idxs]

    # ------------------------------------------------------------------ #
    # Persistence helpers
    # ------------------------------------------------------------------ #
    def save(self, path: str | Path) -> None:
        path = Path(path)
        path.mkdir(parents=True, exist_ok=True)
        meta = {
            "dim": self.dim,
            "metric": self.metric,
            "normalize": self.normalize,
            "index_type": self.index_type,
            "nlist": self.nlist,
            "nprobe": self.nprobe,
            "hnsw_m": self.hnsw_m,
            "hnsw_ef_search": self.hnsw_ef_search,
            "hnsw_ef_construction": self.hnsw_ef_construction,
        }
        (path / "meta.json").write_text(json.dumps(meta), encoding="utf-8")
        (path / "mapping.json").write_text(json.dumps(self._ids), encoding="utf-8")
        V = np.vstack(
    self._vecs) if self._vecs else np.zeros(
        (0, self.dim), dtype=np.float32)
        np.save(path / "vectors.npy", V)
        if _HAS_FAISS and self._index is not None:
            faiss.write_index(self._index, str(path / "faiss.index"))
        self._log(f"Saved index to {path}")

    @classmethod
    def load(
        cls,
        path: str | Path,
        mmap: bool = False,
        **overrides: Any,
    ) -> "FaissIndex":
        """Load a persisted FaissIndex from disk.

        Parameters
        ----------
        path: str | Path
            Directory containing the saved index files.
        mmap: bool
            Whether to memory‑map the index file for faster reads.
        overrides: Any
            Optional overrides for index parameters (e.g., dim, metric).
        """
        path = Path(path)
        # Create a temporary instance to hold metadata
        instance = cls(
            dim=overrides.get("dim", 0),
            metric=overrides.get("metric", "ip"),
            normalize=overrides.get("normalize", True),
            debug=overrides.get("debug", False),
            index_type=overrides.get("index_type", "ivfflat"),
            nlist=overrides.get("nlist", 1024),
            nprobe=overrides.get("nprobe", 10),
            hnsw_m=overrides.get("hnsw_m", 32),
            hnsw_ef_search=overrides.get("hnsw_ef_search", 64),
            hnsw_ef_construction=overrides.get("hnsw_ef_construction", 200),
        )
        # Load vectors and metadata (optional, not needed for index operation)
        # Load the FAISS index if available
        if _HAS_FAISS:
            read_flags = faiss.IO_FLAG_MMAP if mmap else 0
            instance._index = faiss.read_index(str(path / "faiss.index"), read_flags)
        return instance

    # ------------------------------------------------------------------ #
    # Internal helpers
    # ------------------------------------------------------------------ #
    def _build_index(self) -> Any:
        if not _HAS_FAISS:
            return None

        metric_type = (
            faiss.METRIC_INNER_PRODUCT
            if self.metric == "ip"
            else faiss.METRIC_L2
        )

        if self.index_type == "ivfflat":
            quantizer = faiss.IndexFlatIP if self.metric == "ip" else faiss.IndexFlatL2
            index = faiss.IndexIVFFlat(
                quantizer(self.dim), self.dim, self.nlist, metric_type
            )
            index.nprobe = self.nprobe
            self._log(
                f"Created Faiss IndexIVFFlat(dim={self.dim}, nlist={self.nlist}, metric={self.metric})"
            )
        elif self.index_type == "hnsw":
            index = faiss.IndexHNSWFlat(self.dim, self.hnsw_m, metric_type)
            index.hnsw.efSearch = self.hnsw_ef_search
            index.hnsw.efConstruction = self.hnsw_ef_construction
            self._log(
                f"Created Faiss IndexHNSWFlat(dim={self.dim}, M={self.hnsw_m}, metric={self.metric})"
            )
        else:
            raise ValueError(f"Unknown index type: {self.index_type}")

        return index

    def _add_to_faiss(self, matrix: np.ndarray) -> None:
        if self._index is None:
            return
        self._index.add(matrix)

    def _rebuild_faiss(self) -> None:
        if not _HAS_FAISS:
            return
        self._index = self._build_index()
        if self._vecs:
            matrix = np.vstack(self._vecs)
            self._ensure_trained(matrix)
            self._add_to_faiss(matrix)

    def _ensure_trained(self, matrix: np.ndarray) -> None:
        if self._index is None:
            return
        if not self._index.is_trained and self.index_type == "ivfflat":
            if matrix.shape[0] < self.nlist:
                self._log(
                    f"Warning: Not enough vectors ({matrix.shape[0]}) to train IVFFlat index with nlist={self.nlist}. Skipping training."
                )
                return
            self._index.train(matrix)
            self._log("Faiss index trained")

    def _prep_vec(self, vec: np.ndarray) -> np.ndarray:
        v = np.asarray(vec, dtype=np.float32)
        if v.ndim == 1:
            v = v.reshape(1, -1)
        if v.shape[1] != self.dim:
            raise ValueError(
                f"Vector dimension mismatch: expected {self.dim}, got {v.shape[1]}"
            )
        if self.normalize and self.metric == "ip":
            faiss.normalize_L2(v)
        return v[0]

    def _prep_mat(self, matrix: np.ndarray) -> np.ndarray:
        m = np.asarray(matrix, dtype=np.float32)
        if m.ndim == 1:
            m = m.reshape(1, -1)
        if m.shape[1] != self.dim:
            raise ValueError(
                f"Matrix dimension mismatch: expected {self.dim}, got {m.shape[1]}"
            )
        if self.normalize and self.metric == "ip":
            faiss.normalize_L2(m)
        return m

    def _configure_search_params(self) -> None:
        if self._index is None:
            return
        if self.index_type == "ivfflat":
            self._index.nprobe = self.nprobe
        elif self.index_type == "hnsw":
            self._index.hnsw.efSearch = self.hnsw_ef_search

    def _log(self, message: str) -> None:
        if self.debug:
            self._logger.info(message)