"""neuromind.alzheimers — Alzheimer's disease RAG research library.

Quick start
-----------
>>> from neuromind.alzheimers import AlzheimersRAG
>>>
>>> rag = AlzheimersRAG()  # defaults: gpt-oss:20b, db from NEUROMIND_DB_PATH env var
>>>
>>> # 1. Download papers from PubMed Central (requires pubmed-stream)
>>> rag.download("Alzheimer's disease gut microbiome", max_results=50)
>>>
>>> # 2. Index into ChromaDB
>>> rag.index()
>>>
>>> # 3. Query with RAG
>>> result = rag.query(
...     "What is the evidence for gut microbiome dysbiosis in Alzheimer's disease?",
...     mode="detailed",
... )
>>> print(result.text)
>>> print(result.summary())
"""

from __future__ import annotations

import os
from typing import Any, Callable, Dict, List, Optional

from .downloader import AD_SEARCH_TERMS, download_ad_corpus, download_papers
from .indexer import index_directory
from .rag import RAGPipeline
from .retriever import Retriever
from .types import IndexStats, RAGMode, RAGResult, RetrievalResult

__version__ = "0.1.1"


def _make_medembed_fn() -> Optional[Callable]:
    """Return a MedEmbed ChromaDB embedding function, or None (→ DefaultEmbeddingFunction)."""
    try:
        from .medembed_embedder import LangchainMedEmbedEmbeddings  # type: ignore
        return LangchainMedEmbedEmbeddings(device="cpu")
    except Exception:  # sentence-transformers not installed or model not cached
        import logging
        logging.getLogger(__name__).warning(
            "MedEmbed not available — falling back to all-MiniLM-L6-v2. "
            "Install sentence-transformers to use biomedical embeddings."
        )
        return None


__all__ = [
    # High-level façade
    "AlzheimersRAG",
    # Types
    "RAGMode",
    "RAGResult",
    "RetrievalResult",
    "IndexStats",
    # Building blocks
    "RAGPipeline",
    "Retriever",
    # Functions
    "download_papers",
    "download_ad_corpus",
    "index_directory",
    # Data
    "AD_SEARCH_TERMS",
]


class AlzheimersRAG:
    """High-level façade for the Alzheimer's disease RAG pipeline.

    Combines three stages into a single object:

    1. **Download** — pull full-text PMC articles via *pubmed-stream*
    2. **Index** — chunk and embed articles into a persistent ChromaDB store
    3. **Query** — answer biomedical questions under six RAG modes

    Parameters
    ----------
    db_path:
        Filesystem path for the ChromaDB vector database.  Defaults to the
        ``NEUROMIND_DB_PATH`` environment variable, or ``"./neuromind_db"`` if
        that variable is not set.
    collection_name:
        ChromaDB collection name (default: ``"publications"``).
    model:
        Ollama model tag to use for generation (default: ``"gpt-oss:20b"``
        — use any model available in your Ollama instance).
    ollama_host:
        Base URL of the Ollama REST API (default: ``"http://localhost:11434"``).
    num_ctx:
        Context-window size passed to Ollama in tokens (default: 131 072).
    embedding_fn:
        ChromaDB-compatible embedding function.  Defaults to
        ``LangchainMedEmbedEmbeddings`` (``abhinand/MedEmbed-large-v0.1``,
        1024-dim biomedical embeddings) when *sentence-transformers* is
        installed; falls back to ``all-MiniLM-L6-v2`` otherwise.  Must be
        the **same** function used for both indexing and querying.

    Examples
    --------
    >>> rag = AlzheimersRAG()
    >>> rag.download("Alzheimer gut microbiome", max_results=20)
    >>> stats = rag.index()
    >>> result = rag.query("How does gut dysbiosis affect AD?", mode="basic")
    >>> print(result.text)
    """

    def __init__(
        self,
        db_path: Optional[str] = None,
        collection_name: str = "publications",
        model: str = "gpt-oss:20b",
        ollama_host: str = "http://localhost:11434",
        num_ctx: int = 131_072,
        embedding_fn: Optional[Callable] = None,
    ) -> None:
        self._db_path = db_path or os.environ.get("NEUROMIND_DB_PATH", "./neuromind_db")
        self._collection_name = collection_name
        self._model = model
        self._ollama_host = ollama_host
        self._num_ctx = num_ctx
        # Default to MedEmbed (biomedical-grade); falls back to all-MiniLM-L6-v2
        self._embedding_fn = embedding_fn if embedding_fn is not None else _make_medembed_fn()

        # Lazily initialised
        self._retriever: Optional[Retriever] = None
        self._pipeline: Optional[RAGPipeline] = None

    # ------------------------------------------------------------------
    # Download
    # ------------------------------------------------------------------

    def download(
        self,
        keyword: str = "Alzheimer's disease",
        max_results: int = 100,
        output_dir: str = "./publications",
        **kwargs: Any,
    ) -> Any:
        """Download PubMed Central articles matching *keyword*.

        Wraps :func:`~neuromind.alzheimers.download_papers`.  All extra keyword arguments
        are forwarded to that function (e.g. ``api_key``, ``email``,
        ``use_concurrent``, ``fmt``).

        Returns
        -------
        ``DownloadStats`` from *pubmed-stream*.
        """
        return download_papers(
            keyword=keyword,
            max_results=max_results,
            output_dir=output_dir,
            **kwargs,
        )

    def download_corpus(
        self,
        topics: Optional[List[str]] = None,
        max_results_per_topic: int = 100,
        output_dir: str = "./publications",
        **kwargs: Any,
    ) -> List[Any]:
        """Download a comprehensive AD corpus across multiple topics.

        Wraps :func:`~neuromind.alzheimers.download_ad_corpus`.

        Returns
        -------
        List of ``DownloadStats`` objects (one per topic).
        """
        return download_ad_corpus(
            topics=topics,
            max_results_per_topic=max_results_per_topic,
            output_dir=output_dir,
            **kwargs,
        )

    # ------------------------------------------------------------------
    # Index
    # ------------------------------------------------------------------

    def index(
        self,
        papers_dir: str = "./publications",
        chunk_size: int = 2000,
        chunk_overlap: int = 200,
    ) -> IndexStats:
        """Index downloaded articles into the vector database.

        Reads every ``*.json`` file under *papers_dir*, splits text into
        overlapping chunks, and upserts them into ChromaDB.

        After indexing the internal retriever/pipeline are reset so the
        updated collection is used on the next call to :meth:`query`.

        Returns
        -------
        :class:`~neuromind.alzheimers.types.IndexStats`
        """
        stats = index_directory(
            papers_dir=papers_dir,
            db_path=self._db_path,
            collection_name=self._collection_name,
            chunk_size=chunk_size,
            chunk_overlap=chunk_overlap,
            embedding_function=self._embedding_fn,
        )
        # Reset cached components so they pick up the updated collection
        self._retriever = None
        self._pipeline = None
        return stats

    # ------------------------------------------------------------------
    # Query
    # ------------------------------------------------------------------

    def query(
        self,
        question: str,
        mode: RAGMode | str = RAGMode.DETAILED,
        max_response_tokens: int = 3000,
        temperature: float = 0.3,
    ) -> RAGResult:
        """Answer a biomedical research question using RAG.

        Parameters
        ----------
        question:
            Clinical or research question (e.g. *"What is the evidence for
            gut microbiome dysbiosis in Alzheimer's disease?"*).
        mode:
            RAG mode.  Pass a :class:`~neuromind.alzheimers.types.RAGMode` value or one
            of ``"no_rag"``, ``"retrieval_only"``, ``"basic"``,
            ``"standard"``, ``"detailed"``, ``"comprehensive"``.
        max_response_tokens:
            Maximum tokens the LLM may generate.
        temperature:
            LLM sampling temperature (0 = deterministic).

        Returns
        -------
        :class:`~neuromind.alzheimers.types.RAGResult`
        """
        return self._get_pipeline().query(
            question=question,
            mode=mode,
            max_response_tokens=max_response_tokens,
            temperature=temperature,
        )

    def compare_modes(
        self,
        question: str,
        modes: Optional[List[RAGMode]] = None,
        max_response_tokens: int = 2000,
        temperature: float = 0.3,
    ) -> Dict[RAGMode, RAGResult]:
        """Run *question* through multiple RAG modes and return all results.

        Useful for reproducing the ``RAG_Comparison_Demo`` notebook
        programmatically.

        Parameters
        ----------
        question:
            The question to ask under every mode.
        modes:
            Subset of modes to test.  Defaults to all six :class:`RAGMode`
            values.
        max_response_tokens:
            Upper bound on LLM output per mode.
        temperature:
            LLM sampling temperature.

        Returns
        -------
        ``dict`` mapping :class:`~neuromind.alzheimers.types.RAGMode`
        → :class:`~neuromind.alzheimers.types.RAGResult`.
        """
        if modes is None:
            modes = list(RAGMode)

        results: Dict[RAGMode, RAGResult] = {}
        for m in modes:
            results[m] = self.query(
                question=question,
                mode=m,
                max_response_tokens=max_response_tokens,
                temperature=temperature,
            )
        return results

    # ------------------------------------------------------------------
    # Internal
    # ------------------------------------------------------------------

    def _get_retriever(self) -> Retriever:
        if self._retriever is None:
            self._retriever = Retriever(
                db_path=self._db_path,
                collection_name=self._collection_name,
                embedding_function=self._embedding_fn,
            )
        return self._retriever

    def _get_pipeline(self) -> RAGPipeline:
        if self._pipeline is None:
            from .llm import OllamaLLM

            llm = OllamaLLM(
                model=self._model,
                host=self._ollama_host,
                num_ctx=self._num_ctx,
            )
            self._pipeline = RAGPipeline(
                retriever=self._get_retriever(),
                llm=llm,
            )
        return self._pipeline

    def __repr__(self) -> str:
        return (
            f"AlzheimersRAG("
            f"db_path={self._db_path!r}, "
            f"model={self._model!r}, "
            f"collection={self._collection_name!r})"
        )
