"""
FAISS-backed ANN index over paper embeddings.

The index persists three things to disk so a query process never has to
re-embed the full corpus:

  * ``<name>.index``      -- the FAISS flat L2 index (float32 vectors)
  * ``<name>.meta.pkl``   -- parallel arrays: paper_id, year, category, url,
                             title, abstract (aligned to the FAISS rows)
  * ``<name>.meta.json``  -- provenance (embedder name, dim, source, counts)

The index is *incremental*: ``build_index`` can be pointed at a fresh scan of
the metadata and will add any papers not already present (keyed by paper_id),
rather than rebuilding from scratch every time.
"""

from __future__ import annotations

import json
import os
import pickle
from dataclasses import dataclass

import numpy as np

try:
    import faiss
except ImportError:  # pragma: no cover - faiss is a hard dependency
    faiss = None  # type: ignore[assignment]

from data import DataStore, Paper


def _require_faiss() -> None:
    if faiss is None:
        raise RuntimeError("faiss-cpu is required for ANN indexing. Install it (see requirements.txt).")


@dataclass
class IndexMeta:
    embedder: str
    dim: int
    source: str
    count: int
    built_at: str = ""


class AnnIndex:
    """
    Read-only view over a persisted FAISS index, exposing the parallel metadata
    arrays needed to attach titles / years / categories / URLs to ANN hits.
    """

    def __init__(self, builder: IndexBuilder, meta: IndexMeta):
        self.index = builder.index
        self.meta = meta
        self._ids = builder._ids
        self._years = builder._years
        self._cats = builder._cats
        self._urls = builder._urls
        self._titles = builder._titles
        self._abstracts = builder._abstracts

    def __len__(self) -> int:
        return len(self._ids)


class IndexBuilder:
    """
    Build (and incrementally extend) a FAISS index from a DataStore scan.

    Example
    -------
    >>> builder = IndexBuilder("arxiv", embedder=HashingEmbedder())
    >>> store = DataStore()
    >>> builder.build(store.scan_filtered(), store)   # doctest: +SKIP
    >>> builder.save("arxiv")
    """

    def __init__(self, name: str, embedder, chunk_size: int = 4096):
        _require_faiss()
        self.name = name
        self.embedder = embedder
        self.dim = int(embedder.DIM)
        self.chunk_size = chunk_size
        self._existing: set[str] = set()
        self._ids: list[str] = []
        self._years: list[int] = []
        self._cats: list[str] = []
        self._urls: list[str] = []
        self._titles: list[str] = []
        self._abstracts: list[str] = []
        self._loaded = False

    # -- loading an existing index ------------------------------------------
    def load(self, path: str | None = None) -> None:
        path = path or self.default_path()
        if not os.path.exists(path):
            return
        self.index = faiss.read_index(path)
        with open(path.replace(".index", ".meta.pkl"), "rb") as fh:
            blob = pickle.load(fh)
        (
            self._ids,
            self._years,
            self._cats,
            self._urls,
            self._titles,
            self._abstracts,
        ) = blob
        self._existing = set(self._ids)
        self._loaded = True

    def default_path(self, root: str | None = None) -> str:
        root = root or os.path.join(os.path.dirname(__file__), "index")
        return os.path.join(root, f"{self.name}.index")

    # -- building -----------------------------------------------------------
    def build(self, scan, store: DataStore) -> int:
        """
        Embed every row of ``scan`` (a polars LazyFrame) and add new papers.

        ``scan`` must project the DISPLAY_COLUMNS; ``store`` is used only for its
        provenance metadata.  Returns the number of papers newly added.
        """
        import datetime as _dt

        if not self._loaded:
            self.index = faiss.IndexFlatL2(self.dim)
        added = 0
        batch: list[str] = []
        batch_ids: list[str] = []
        batch_years: list[int] = []
        batch_cats: list[str] = []
        batch_urls: list[str] = []
        batch_titles: list[str] = []
        batch_abstracts: list[str] = []

        def flush() -> None:
            nonlocal added
            if not batch:
                return
            # Only embed/add papers not already in the index.
            keep = [i for i, pid in enumerate(batch_ids) if pid not in self._existing]
            if not keep:
                batch.clear()
                batch_ids.clear()
                batch_years.clear()
                batch_cats.clear()
                batch_urls.clear()
                batch_titles.clear()
                batch_abstracts.clear()
                return
            keep_ids = [batch_ids[i] for i in keep]
            keep_texts = [batch[i] for i in keep]
            vecs = self.embedder.embed(keep_texts)
            vecs = np.ascontiguousarray(vecs, dtype=np.float32)
            self.index.add(vecs)
            for i in keep:
                self._ids.append(batch_ids[i])
                self._years.append(batch_years[i])
                self._cats.append(batch_cats[i])
                self._urls.append(batch_urls[i])
                self._titles.append(batch_titles[i])
                self._abstracts.append(batch_abstracts[i])
                self._existing.add(batch_ids[i])
            added += len(keep_ids)
            batch.clear()
            batch_ids.clear()
            batch_years.clear()
            batch_cats.clear()
            batch_urls.clear()
            batch_titles.clear()
            batch_abstracts.clear()

        # Collect all matching rows, then embed in chunks.
        rows = scan.select(
            ["paper_id", "title", "primary_category", "first_version_date", "arxiv_abs_url", "abstract"]
        ).collect().to_dicts()

        for r in rows:
            rec = dict(r)
            fvd = rec["first_version_date"]
            year = fvd.year if isinstance(fvd, _dt.datetime) else (int(fvd) if fvd is not None else 0)
            pid = str(rec["paper_id"])
            if pid in self._existing:
                continue
            text = f"{rec['title']} {rec['abstract']}"
            batch.append(text)
            batch_ids.append(pid)
            batch_years.append(year)
            batch_cats.append(str(rec["primary_category"] or ""))
            batch_urls.append(str(rec["arxiv_abs_url"] or ""))
            batch_titles.append(str(rec["title"] or ""))
            batch_abstracts.append(str(rec["abstract"] or ""))
            if len(batch) >= self.chunk_size:
                flush()
        flush()
        return added

    # -- persistence --------------------------------------------------------
    def save(self, path: str | None = None) -> IndexMeta:
        path = path or self.default_path()
        os.makedirs(os.path.dirname(path), exist_ok=True)
        faiss.write_index(self.index, path)
        meta_pkl = path.replace(".index", ".meta.pkl")
        with open(meta_pkl, "wb") as fh:
            pickle.dump(
                (self._ids, self._years, self._cats, self._urls, self._titles, self._abstracts),
                fh,
            )
        meta = IndexMeta(
            embedder=type(self.embedder).__name__,
            dim=self.dim,
            source=os.environ.get("ARXIV_METADATA_PARQUET", "arxiv-complete/metadata"),
            count=len(self._ids),
        )
        with open(path.replace(".index", ".meta.json"), "w") as fh:
            json.dump(meta.__dict__, fh, indent=2)
        return meta


def load_index(name: str, embedder, root: str | None = None) -> tuple["AnnIndex", IndexMeta]:
    """Load a previously-built index plus its embedder for querying."""
    _require_faiss()
    root = root or os.path.join(os.path.dirname(__file__), "index")
    path = os.path.join(root, f"{name}.index")
    builder = IndexBuilder(name, embedder)
    builder.load(path)
    meta_path = path.replace(".index", ".meta.json")
    with open(meta_path) as fh:
        meta = IndexMeta(**json.load(fh))
    return AnnIndex(builder, meta), meta