"""RAG pipeline for Alzheimer's disease research.

Implements six evidence-depth conditions that mirror the
``RAG_Comparison_Demo`` notebook:

    1. NO_RAG          — LLM answers from training knowledge alone
    2. RETRIEVAL_ONLY  — Raw excerpts, no generation
    3. BASIC           — 5 papers  (~20 K context tokens)
    4. STANDARD        — 10 papers (~40 K context tokens)
    5. DETAILED        — 15 papers (~70 K context tokens)  ← recommended default
    6. COMPREHENSIVE   — 25 papers (~100 K context tokens)
"""

from __future__ import annotations

import logging
from typing import TYPE_CHECKING, Dict, List

from .types import RAGMode, RAGResult, RAG_MODE_MAX_TOKENS, RAG_MODE_TOP_K, RetrievalResult

if TYPE_CHECKING:
    from .llm import OllamaLLM
    from .retriever import Retriever

logger = logging.getLogger(__name__)


# ---------------------------------------------------------------------------
# Internal helpers
# ---------------------------------------------------------------------------


def _estimate_tokens(text: str) -> int:
    """Rough token count estimate (1 token ≈ 4 characters)."""
    return len(text) // 4


def _build_context(
    documents: List[str],
    metadatas: List[Dict],
    max_tokens: int,
) -> tuple[str, int]:
    """Pack retrieved documents into a single context string.

    Stops adding documents once *max_tokens* would be exceeded.

    Returns
    -------
    ``(context_string, tokens_used)``
    """
    parts: List[str] = []
    tokens_used = 0

    for i, (doc, meta) in enumerate(zip(documents, metadatas), 1):
        pmcid = meta.get("pmcid") or f"PMC{meta.get('pmc_id', 'N/A')}"
        title = meta.get("title", "")
        header = f"[Source {i}: {pmcid}]"
        if title:
            header += f" — {title[:100]}"
        entry = f"{header}\n{doc}\n"
        entry_tokens = _estimate_tokens(entry)
        if tokens_used + entry_tokens > max_tokens:
            break
        parts.append(entry)
        tokens_used += entry_tokens

    return "\n".join(parts), tokens_used


def _build_prompt(
    query: str,
    mode: RAGMode,
    context: str = "",
    top_k: int = 0,
) -> str:
    """Return the full prompt for a given RAG mode."""

    expert_base = "You are a medical research expert on Alzheimer's disease."

    if mode == RAGMode.NO_RAG:
        return (
            f"{expert_base}\n\n"
            "Answer the following question based on your training knowledge.\n\n"
            f"QUESTION: {query}\n\n"
            "Provide a structured answer with specific evidence and mechanisms."
        )

    if mode == RAGMode.BASIC:
        return (
            f"{expert_base}\n\n"
            f"Based on the following {top_k} scientific literature excerpts, "
            "answer the question.\n"
            "Cite sources as [Source N].\n\n"
            f"SCIENTIFIC LITERATURE:\n{context}\n\n"
            f"QUESTION: {query}\n\n"
            "ANSWER (cite sources, be specific):"
        )

    if mode == RAGMode.STANDARD:
        return (
            f"{expert_base}\n\n"
            f"Based on the following {top_k} scientific literature excerpts, "
            "provide a thorough analysis.\n"
            "Cite sources as [Source N]. Discuss mechanisms, biomarkers, "
            "and therapeutic strategies.\n\n"
            f"SCIENTIFIC LITERATURE:\n{context}\n\n"
            f"QUESTION: {query}\n\n"
            "Provide your analysis in these sections:\n"
            "## 1. Key Evidence from the Literature\n"
            "## 2. Mechanisms and Pathophysiology\n"
            "## 3. Clinical Implications and Therapeutic Strategies\n\n"
            "ANALYSIS:"
        )

    if mode == RAGMode.DETAILED:
        return (
            "You are a leading medical research expert specializing in Alzheimer's disease.\n\n"
            f"Based on the following {top_k} scientific publications, provide a comprehensive, "
            "evidence-based analysis.\n"
            "Cite sources as [Source N]. Include specific biomarkers, "
            "mechanisms, and clinical evidence.\n\n"
            f"SCIENTIFIC LITERATURE ({top_k} sources):\n{context}\n\n"
            f"QUESTION: {query}\n\n"
            "Provide your analysis in these sections:\n"
            "## 1. Supporting Evidence from the Literature\n"
            "## 2. Pathophysiological Mechanisms\n"
            "## 3. Biomarker Connections\n"
            "## 4. Therapeutic Strategies & Clinical Evidence\n\n"
            "COMPREHENSIVE ANALYSIS:"
        )

    # COMPREHENSIVE (default for unrecognised modes)
    return (
        "You are a leading medical research expert specializing in Alzheimer's disease, "
        "with deep expertise in neuroinflammation, biomarkers, and multi-modal analysis.\n\n"
        f"Based on the following {top_k} scientific publications, provide an exhaustive, "
        "research-grade analysis.\n\n"
        "Cite sources as [Source N]. Include specific biomarker values, effect sizes, "
        "molecular pathways, and clinical trial outcomes where available.\n\n"
        f"SCIENTIFIC LITERATURE ({top_k} sources):\n{context}\n\n"
        f"QUESTION: {query}\n\n"
        "Provide your analysis in these sections:\n"
        "## 1. Epidemiological & Clinical Evidence\n"
        "## 2. Pathophysiological Mechanisms\n"
        "## 3. Biomarker Connections\n"
        "## 4. Therapeutic Strategies & Clinical Evidence\n"
        "## 5. Future Directions & Research Gaps\n\n"
        "COMPREHENSIVE ANALYSIS:"
    )


# ---------------------------------------------------------------------------
# RAGPipeline
# ---------------------------------------------------------------------------


class RAGPipeline:
    """Full RAG pipeline combining retrieval and LLM generation.

    Parameters
    ----------
    retriever:
        A :class:`~neuromind.retriever.Retriever` instance.
    llm:
        An :class:`~neuromind.llm.OllamaLLM` instance (or any object
        with a ``generate(prompt, temperature, max_tokens)`` method that
        returns an object with ``.text`` and ``.words`` attributes).

    Examples
    --------
    >>> from neuromind.alzheimers.retriever import Retriever
    >>> from neuromind.alzheimers.llm import OllamaLLM
    >>> from neuromind.alzheimers.rag import RAGPipeline
    >>> from neuromind.alzheimers.types import RAGMode
    >>>
    >>> pipeline = RAGPipeline(
    ...     retriever=Retriever("./neuromind_db"),
    ...     llm=OllamaLLM(model="llama3.2"),
    ... )
    >>> result = pipeline.query(
    ...     "What is the role of gut microbiome in Alzheimer's?",
    ...     mode=RAGMode.DETAILED,
    ... )
    >>> print(result.text)
    """

    def __init__(self, retriever: "Retriever", llm: "OllamaLLM") -> None:
        self._retriever = retriever
        self._llm = llm

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

    def query(
        self,
        question: str,
        mode: RAGMode | str = RAGMode.DETAILED,
        max_response_tokens: int = 3000,
        temperature: float = 0.3,
    ) -> RAGResult:
        """Answer *question* under the specified RAG *mode*.

        Parameters
        ----------
        question:
            Research or clinical question to answer.
        mode:
            RAG mode.  Accepts a :class:`~neuromind.alzheimers.types.RAGMode` value or
            its string equivalent (``"no_rag"``, ``"basic"``,
            ``"standard"``, ``"detailed"``, ``"comprehensive"``).
        max_response_tokens:
            Upper bound on LLM output tokens.
        temperature:
            LLM sampling temperature (0 = deterministic).

        Returns
        -------
        :class:`~neuromind.alzheimers.types.RAGResult`
        """
        if isinstance(mode, str):
            mode = RAGMode(mode)

        top_k = RAG_MODE_TOP_K[mode]
        max_ctx = RAG_MODE_MAX_TOKENS[mode]

        # ── Retrieval ──────────────────────────────────────────────────
        retrieve_time = 0.0
        raw_docs: List[str] = []
        metas: List[Dict] = []
        similarities: List[float] = []

        if top_k > 0:
            ret: RetrievalResult = self._retriever.retrieve(question, top_k=top_k)
            raw_docs = ret.documents
            metas = ret.metadatas
            similarities = ret.similarities
            retrieve_time = ret.retrieve_time
            logger.info(
                "Retrieved %d chunks in %.2fs (mode=%s)",
                len(raw_docs),
                retrieve_time,
                mode.value,
            )

        # ── Retrieval-only: skip generation ────────────────────────────
        if mode == RAGMode.RETRIEVAL_ONLY:
            sources = [
                {**meta, "similarity": sim, "excerpt": doc[:300]}
                for doc, meta, sim in zip(raw_docs, metas, similarities)
            ]
            return RAGResult(
                text=self._format_retrieval_only(raw_docs, metas, similarities),
                mode=mode,
                query=question,
                sources=sources,
                context_tokens=0,
                response_words=0,
                retrieve_time=retrieve_time,
                generate_time=0.0,
            )

        # ── Build context string ───────────────────────────────────────
        context, ctx_tokens = "", 0
        if top_k > 0 and raw_docs:
            context, ctx_tokens = _build_context(raw_docs, metas, max_ctx)

        # ── Generation ────────────────────────────────────────────────
        prompt = _build_prompt(question, mode, context, top_k)
        llm_result = self._llm.generate(
            prompt,
            temperature=temperature,
            max_tokens=max_response_tokens,
        )
        logger.info(
            "Generated %d words in %.1fs (mode=%s)",
            llm_result.words,
            llm_result.generate_time,
            mode.value,
        )

        sources = [
            {**meta, "similarity": sim}
            for meta, sim in zip(metas, similarities)
        ]
        return RAGResult(
            text=llm_result.text,
            mode=mode,
            query=question,
            sources=sources,
            context_tokens=ctx_tokens,
            response_words=llm_result.words,
            retrieve_time=retrieve_time,
            generate_time=llm_result.generate_time,
        )

    # ------------------------------------------------------------------
    # Helpers
    # ------------------------------------------------------------------

    @staticmethod
    def _format_retrieval_only(
        documents: List[str],
        metadatas: List[Dict],
        similarities: List[float],
    ) -> str:
        lines = ["=== RETRIEVAL-ONLY RESULTS (no generation) ===\n"]
        for i, (doc, meta, sim) in enumerate(zip(documents, metadatas, similarities), 1):
            pmcid = meta.get("pmcid", "N/A")
            title = meta.get("title", "")
            lines.append("─" * 70)
            lines.append(f"Source {i}: {pmcid}  (similarity: {sim:.4f})")
            if title:
                lines.append(f"Title: {title}")
            lines.append("")
            lines.append(doc[:500] + ("..." if len(doc) > 500 else ""))
            lines.append("")
        return "\n".join(lines)
