"""ResearchPath Streamlit app.

Two tabs:
  1. Reading Path  — prerequisite-chain planner (the unique agentic feature)
  2. Ask           — grounded Q&A over the indexed RL corpus
"""
from __future__ import annotations

import sys
from pathlib import Path

sys.path.insert(0, str(Path(__file__).resolve().parent))

import streamlit as st

from researchpath.corpus import CANONICAL_RL_PAPERS
from researchpath.planning import plan_reading_path

ROOT = Path(__file__).resolve().parent
INDEX_PATH = ROOT / "data" / "index.faiss"

_TAG_TO_ID = {p.tag.lower(): p.arxiv_id for p in CANONICAL_RL_PAPERS}
_ID_TO_PAPER = {p.arxiv_id: p for p in CANONICAL_RL_PAPERS}
_ERA_ORDER = ["value-based", "policy-gradient", "actor-critic", "model-based", "rlhf"]

st.set_page_config(
    page_title="ResearchPath",
    page_icon="🗺️",
    layout="wide",
    initial_sidebar_state="collapsed",
)

# ── minimal style ──────────────────────────────────────────────────────────
st.markdown(
    """
    <style>
    .step-card {
        background: #1e2330;
        border-left: 4px solid #4f8ef7;
        border-radius: 6px;
        padding: 14px 18px;
        margin-bottom: 12px;
    }
    .step-card h4 { margin: 0 0 4px 0; color: #e0e6f0; }
    .step-card .concepts { color: #a0aec0; font-size: 0.85rem; }
    .step-card .why { color: #cbd5e0; font-size: 0.9rem; margin-top: 6px; }
    .tag-chip {
        display: inline-block;
        background: #2d3748;
        color: #90cdf4;
        border-radius: 12px;
        padding: 2px 10px;
        font-size: 0.78rem;
        margin-right: 4px;
    }
    </style>
    """,
    unsafe_allow_html=True,
)

# ── header ─────────────────────────────────────────────────────────────────
st.title("🗺️ ResearchPath")
st.caption(
    "An agentic research-onboarding companion for Reinforcement Learning. "
    "Give it a target paper and your background; it builds a personalized, "
    "dependency-ordered reading plan."
)

tab_plan, tab_ask = st.tabs(["📚 Reading Path", "💬 Ask a Question"])


# ── helpers ────────────────────────────────────────────────────────────────

def _resolve(token: str) -> str | None:
    token = token.strip()
    if token in _ID_TO_PAPER:
        return token
    if token.lower() in _TAG_TO_ID:
        return _TAG_TO_ID[token.lower()]
    return None


def _paper_label(arxiv_id: str) -> str:
    p = _ID_TO_PAPER[arxiv_id]
    return f"{p.tag} ({p.year})"


@st.cache_resource(show_spinner="Loading index and embedder…")
def _load_retrieval():
    """Load FAISS index + embedder once per session."""
    from researchpath.embeddings import Embedder
    from researchpath.index import load_index
    from researchpath.retrieval import HybridRetriever

    if not INDEX_PATH.exists():
        return None, None, None
    index, chunks = load_index(INDEX_PATH)
    embedder = Embedder()
    retriever = HybridRetriever(index, chunks, embedder)
    return index, chunks, retriever


# ═══════════════════════════════════════════════════════════════════════════
# TAB 1 — Reading Path
# ═══════════════════════════════════════════════════════════════════════════
with tab_plan:
    st.subheader("Build your personalized reading path")
    st.markdown(
        "Select a **target paper** you want to understand, then tell us what you **already know**. "
        "ResearchPath traces the prerequisite chain and gives you a topologically-sorted reading list "
        "with *why each paper is next*."
    )

    # Group papers by era for a nicer selectbox
    era_groups: dict[str, list] = {era: [] for era in _ERA_ORDER}
    for p in CANONICAL_RL_PAPERS:
        era_groups[p.era].append(p)

    paper_options = []
    for era in _ERA_ORDER:
        for p in era_groups[era]:
            paper_options.append(p.arxiv_id)

    def _format_option(arxiv_id: str) -> str:
        p = _ID_TO_PAPER[arxiv_id]
        return f"{p.tag}  —  {p.title[:60]}{'…' if len(p.title) > 60 else ''}  ({p.year})"

    col_target, col_known = st.columns([1, 1])

    with col_target:
        target_id = st.selectbox(
            "Target paper",
            options=paper_options,
            format_func=_format_option,
            index=paper_options.index("1710.02298"),  # default: Rainbow
        )

    with col_known:
        known_options = [aid for aid in paper_options if aid != target_id]
        known_ids_raw = st.multiselect(
            "Papers you already know (optional)",
            options=known_options,
            format_func=_format_option,
            default=[],
        )

    if st.button("Generate reading path", type="primary", use_container_width=True):
        plan = plan_reading_path(target_id, known_ids=set(known_ids_raw))
        target_paper = _ID_TO_PAPER[target_id]

        if not plan.steps:
            st.success(f"You're already ready to read **{target_paper.tag}** directly!")
        else:
            st.markdown(
                f"**{len(plan.steps)} paper(s)** to read before (and including) "
                f"**{target_paper.tag}**, in order:"
            )
            for i, step in enumerate(plan.steps, 1):
                p = step.paper
                is_target = p.arxiv_id == target_id
                border_color = "#f6c90e" if is_target else "#4f8ef7"
                label = "🎯 TARGET" if is_target else f"Step {i}"
                concepts_html = "".join(
                    f'<span class="tag-chip">{c}</span>' for c in step.concepts
                )
                bridges_text = (
                    "→ " + ", ".join(_paper_label(b) for b in step.bridges_to)
                    if step.bridges_to
                    else ""
                )
                st.markdown(
                    f"""
                    <div class="step-card" style="border-left-color:{border_color}">
                        <h4>{label} &nbsp; [{p.arxiv_id}] &nbsp; <b>{p.tag}</b> &nbsp; <span style="color:#718096;font-weight:normal">({p.year})</span></h4>
                        <div style="color:#a0aec0;font-size:0.9rem">{p.title}</div>
                        {f'<div class="concepts" style="margin-top:8px">{concepts_html}</div>' if concepts_html else ""}
                        {f'<div class="why">{step.why}</div>' if step.why else ""}
                        {f'<div style="color:#718096;font-size:0.82rem;margin-top:6px">{bridges_text}</div>' if bridges_text else ""}
                    </div>
                    """,
                    unsafe_allow_html=True,
                )

    with st.expander("Browse all canonical RL papers"):
        for era in _ERA_ORDER:
            st.markdown(f"**{era.replace('-', ' ').title()}**")
            for p in era_groups[era]:
                st.markdown(
                    f"&nbsp;&nbsp;&nbsp;`{p.arxiv_id}` &nbsp; **{p.tag}** ({p.year}) — {p.title}"
                )


# ═══════════════════════════════════════════════════════════════════════════
# TAB 2 — Ask a Question
# ═══════════════════════════════════════════════════════════════════════════
with tab_ask:
    st.subheader("Ask anything about the RL corpus")
    st.markdown(
        "Every answer is grounded in the indexed corpus with **`[source_id, p<page>]` citations**. "
        "Corpus: 17 RL papers + Sutton & Barto textbook + RLHF Book + CS224R notes + 5 tutorials "
        "(5,531 chunks total). No hallucination — if the corpus doesn't contain an answer, the model says so."
    )

    index_ready = INDEX_PATH.exists()
    if not index_ready:
        st.warning(
            "No FAISS index found at `data/index.faiss`. "
            "Run `uv run python scripts/build_index.py` first."
        )

    with st.form("ask_form"):
        question = st.text_area(
            "Your question",
            placeholder="What is the main idea behind Proximal Policy Optimization?",
            height=80,
        )
        col_k, col_mode, _ = st.columns([1, 2, 3])
        with col_k:
            k = st.number_input("Top-k chunks", min_value=1, max_value=20, value=5)
        with col_mode:
            retrieval_mode = st.selectbox(
                "Retrieval mode",
                ["Hybrid (BM25 + FAISS)", "Dense (FAISS only)"],
            )
        submitted = st.form_submit_button("Ask", type="primary", use_container_width=True)

    import os
    _gemini_key = os.environ.get("GEMINI_API_KEY", "")
    if not _gemini_key:
        st.info(
            "**GEMINI_API_KEY not set** — the Ask tab needs a Gemini API key to generate answers. "
            "Set it in your `.env` file (local) or as a Space secret (HF Spaces). "
            "The Reading Path tab above works entirely offline with no API key."
        )

    if submitted and question.strip() and index_ready:
        if not _gemini_key:
            st.error("Cannot generate an answer: GEMINI_API_KEY is not set.")
        else:
            _, _, retriever = _load_retrieval()
            if retriever is None:
                st.error("Failed to load index.")
            else:
                with st.spinner("Retrieving and generating…"):
                    from researchpath.rag import answer as rag_answer
                    from researchpath.index import load_index, search
                    from researchpath.embeddings import Embedder

                    if retrieval_mode.startswith("Hybrid"):
                        hits = retriever.search(question.strip(), k=int(k))
                    else:
                        index, chunks = load_index(INDEX_PATH)
                        embedder = Embedder()
                        hits = search(index, chunks, embedder, question.strip(), k=int(k))

                    result = rag_answer(question.strip(), hits)

                st.markdown("### Answer")
                st.markdown(result.answer)

                st.markdown("---")
                col_meta1, col_meta2, col_meta3 = st.columns(3)
                col_meta1.metric("Chunks retrieved", len(hits))
                col_meta2.metric("Tokens in", result.llm.input_tokens or "—")
                col_meta3.metric("Tokens out", result.llm.output_tokens or "—")

                with st.expander(f"📄 Retrieved chunks ({len(hits)})"):
                    for i, h in enumerate(hits, 1):
                        _type_badge = {
                            "paper": "📄", "textbook": "📚",
                            "course": "🎓", "tutorial": "🌐",
                        }.get(h.source_type, "📄")
                        st.markdown(
                            f"**#{i}** &nbsp; `[{h.arxiv_id}, p{h.page}]` &nbsp; "
                            f"{_type_badge} *{h.source_type}* &nbsp; "
                            f"score={h.score:.3f}"
                        )
                        st.text(h.text[:400] + ("…" if len(h.text) > 400 else ""))
                        st.divider()

    # Example questions
    with st.expander("Example questions to try"):
        examples = [
            "What are the six components Rainbow combines?",
            "How does PPO avoid the large policy update problem in TRPO?",
            "What is the key insight behind Generalized Advantage Estimation?",
            "How does DDPG handle continuous action spaces?",
            "What is the DPO loss function and why does it not need a reward model?",
            "How does A3C's asynchronous training replace experience replay?",
        ]
        for ex in examples:
            st.markdown(f"- *{ex}*")