"""Vector store wrapper around Chroma for relevance scoring of papers.

The embedding model is loaded once per process (singleton) and computed
embeddings are cached on disk by content hash, so repeated papers across runs
skip re-encoding. Each run still uses its own Chroma collection so concurrent
jobs never contaminate one another's candidate set.
"""

from __future__ import annotations

import hashlib
import os
import uuid

from .tools import Paper

EMBED_MODEL = "all-MiniLM-L6-v2"
CACHE_DIR = os.getenv("EMBED_CACHE_DIR", "papers/.embcache")

_model = None  # lazily-loaded SentenceTransformer singleton


def _get_model():
    global _model
    if _model is None:
        from sentence_transformers import SentenceTransformer

        _model = SentenceTransformer(EMBED_MODEL)
    return _model


def _embed(texts: list[str]) -> list[list[float]]:
    """Embed texts, using an on-disk per-text cache keyed by content hash."""
    import numpy as np

    os.makedirs(CACHE_DIR, exist_ok=True)
    out: list[list[float] | None] = [None] * len(texts)
    to_compute: list[int] = []
    paths: list[str] = []

    for i, t in enumerate(texts):
        key = hashlib.sha1(t.encode("utf-8")).hexdigest()
        path = os.path.join(CACHE_DIR, f"{key}.npy")
        paths.append(path)
        if os.path.exists(path):
            try:
                out[i] = np.load(path).tolist()
                continue
            except Exception:
                pass
        to_compute.append(i)

    if to_compute:
        vecs = _get_model().encode([texts[i] for i in to_compute])
        for j, i in enumerate(to_compute):
            vec = np.asarray(vecs[j], dtype="float32")
            out[i] = vec.tolist()
            try:
                np.save(paths[i], vec)
            except Exception:
                pass
    return [v for v in out]  # type: ignore[return-value]


class PaperVectorDB:
    """In-memory Chroma collection (per run) over cached MiniLM embeddings."""

    def __init__(self):
        import chromadb

        self._client = chromadb.Client()
        self._name = f"papers_{uuid.uuid4().hex[:8]}"
        self._collection = self._client.get_or_create_collection(name=self._name)

    @staticmethod
    def _document(paper: Paper) -> str:
        first_claim = paper.claims[0] if paper.claims else ""
        return f"{paper.title}\n{paper.abstract}\n{first_claim}".strip()

    def index(self, papers: list[Paper]) -> None:
        if not papers:
            return
        docs = [self._document(p) for p in papers]
        self._collection.upsert(
            ids=[p.id for p in papers],
            documents=docs,
            embeddings=_embed(docs),
            metadatas=[{"title": p.title, "year": p.year} for p in papers],
        )

    def search(self, query: str, n: int = 10) -> list[str]:
        """Return paper IDs ranked by similarity to ``query`` (best first)."""
        count = self._collection.count()
        if count == 0:
            return []
        res = self._collection.query(
            query_embeddings=_embed([query]), n_results=min(n, count)
        )
        ids = res.get("ids") or [[]]
        return ids[0]

    def close(self) -> None:
        """Drop this run's collection to free memory."""
        try:
            self._client.delete_collection(self._name)
        except Exception:
            pass