"""Nearest-neighbor annotation lookup against a FAISS index + parquet table
mounted from the HF bucket (via hfmount) -- reference_min.index /
reference_min.parquet, row i of the parquet matching vector i of the index.
Table columns: id, source (swissprot / acr / phrog), annotation, category."""
import os

from utils.vector_store import VectorStore

REFERENCE_DIR = os.environ.get("REFERENCE_DATA_DIR", "/data/reference")
INDEX_FILENAME = os.environ.get("REFERENCE_INDEX_FILE", "reference_min.index")
PARQUET_FILENAME = os.environ.get("REFERENCE_PARQUET_FILE", "reference_min.parquet")


class AnnotationRetriever:
    def __init__(self, data_dir=REFERENCE_DIR, index_filename=INDEX_FILENAME,
                 parquet_filename=PARQUET_FILENAME, top_k=16):
        self.store = VectorStore(
            os.path.join(data_dir, index_filename),
            os.path.join(data_dir, parquet_filename),
            id_column=None, top_k=top_k,
        )

    def query(self, esm_embeddings):
        """esm_embeddings: [N, D] array-like. Returns a list of N dicts, one
        per query: {"top_hit", "other_source_hit"}."""
        scores, indices = self.store.search(esm_embeddings)
        rows = self.store.fetch_rows(indices.ravel())
        return [self._summarize(s, i, rows) for s, i in zip(scores, indices)]

    def _as_dict(self, score, row):
        return {
            "id": row["id"], "source": row["source"], "annotation": row["annotation"],
            "category": row["category"], "score": float(score),
        }

    def _summarize(self, scores, indices, rows):
        valid = [(float(s), int(i)) for s, i in zip(scores, indices) if i >= 0]
        if not valid:
            return {"top_hit": None, "other_source_hit": None}

        top_score, top_idx = valid[0]
        top_hit = self._as_dict(top_score, rows[top_idx])
        top_source = rows[top_idx]["source"]

        other_source_hit = None
        for score, idx in valid[1:]:
            if rows[idx]["source"] != top_source:
                other_source_hit = self._as_dict(score, rows[idx])
                break

        return {"top_hit": top_hit, "other_source_hit": other_source_hit}