"""ChromaDB-backed document retriever for the Alzheimer's knowledge base."""

from __future__ import annotations

import logging
import time
from typing import Callable, Optional

from .types import RetrievalResult

logger = logging.getLogger(__name__)


class Retriever:
    """Semantic retrieval over a ChromaDB collection.

    Parameters
    ----------
    db_path:
        Filesystem path to the persistent ChromaDB database.
    collection_name:
        Name of the collection to query (default: ``"publications"``).
    embedding_function:
        ChromaDB-compatible embedding function.  Defaults to
        ``DefaultEmbeddingFunction`` (``all-MiniLM-L6-v2``).
        Pass the same function that was used during indexing.

    Examples
    --------
    >>> from neuromind.alzheimers.retriever import Retriever
    >>> ret = Retriever(db_path="./neuromind_db")
    >>> result = ret.retrieve("gut microbiome Alzheimer's", top_k=5)
    >>> print(result.documents[0][:200])
    """

    def __init__(
        self,
        db_path: str,
        collection_name: str = "publications",
        embedding_function: Optional[Callable] = None,
    ) -> None:
        try:
            import chromadb
            from chromadb.utils.embedding_functions import DefaultEmbeddingFunction
        except ImportError as exc:
            raise ImportError(
                "chromadb is required.\n"
                "Install with:  pip install chromadb"
            ) from exc

        ef = embedding_function if embedding_function is not None else DefaultEmbeddingFunction()

        client = chromadb.PersistentClient(path=db_path)
        self._collection = client.get_or_create_collection(
            name=collection_name,
            embedding_function=ef,
        )
        logger.info(
            "Retriever ready — %d chunks in collection '%s'",
            self._collection.count(),
            collection_name,
        )

    # ------------------------------------------------------------------
    # Properties
    # ------------------------------------------------------------------

    @property
    def chunk_count(self) -> int:
        """Number of indexed document chunks in the collection."""
        return self._collection.count()

    # ------------------------------------------------------------------
    # Retrieval
    # ------------------------------------------------------------------

    def retrieve(self, query: str, top_k: int = 5) -> RetrievalResult:
        """Return the *top_k* chunks most semantically similar to *query*.

        Parameters
        ----------
        query:
            Natural-language query string.
        top_k:
            Number of chunks to retrieve.

        Returns
        -------
        :class:`~neuromind.alzheimers.types.RetrievalResult`
        """
        if self._collection.count() == 0:
            logger.warning(
                "Collection is empty — run `alzheimers index` or `AlzheimersRAG.index()` first."
            )
            return RetrievalResult(
                documents=[], metadatas=[], similarities=[], retrieve_time=0.0
            )

        t0 = time.time()
        results = self._collection.query(
            query_texts=[query],
            n_results=min(top_k, self._collection.count()),
            include=["documents", "metadatas", "distances"],
        )
        elapsed = time.time() - t0

        docs   = results["documents"][0]
        metas  = results["metadatas"][0]
        # ChromaDB returns squared-L2 or cosine distances; convert to [0, 1] similarity
        similarities = [max(0.0, 1.0 - d) for d in results["distances"][0]]

        return RetrievalResult(
            documents=docs,
            metadatas=metas,
            similarities=similarities,
            retrieve_time=elapsed,
        )
