server_ui.py

"""
server_ui.py - Multi-Agent Clinical Decision Support Chatbot

Stack: Streamlit, LangGraph, LlamaIndex/ChromaDB, DeepSeek V4, CrossEncoder
"""

# =============================================================================
# 1. Imports & Initializations
# =============================================================================

import os
import time
os.environ["CUDA_VISIBLE_DEVICES"] = ""
import json
import uuid
import logging
from datetime import date
from typing import TypedDict

from rich.logging import RichHandler
from rich.console import Console
from rich.theme import Theme

console = Console(theme=Theme({
    "logging.level.info": "bright_cyan",
    "logging.level.warning": "bright_yellow",
    "logging.level.error": "bright_red",
    "log.time": "bright_black",
    "log.message": "white",
}))

logging.basicConfig(
    level=logging.INFO,
    format="%(message)s",
    datefmt="%Y-%m-%d %H:%M:%S",
    handlers=[RichHandler(
        show_time=True, show_level=True, show_path=False,
        omit_repeated_times=False,
        console=console,
    )],
)
for noisy in ("huggingface_hub", "sentence_transformers", "urllib3", "httpx"):
    logging.getLogger(noisy).setLevel(logging.WARNING)
logger = logging.getLogger(__name__)

import streamlit as st
import openai
from openai import OpenAI

from langgraph.graph import StateGraph, END

import chromadb
from llama_index.core import VectorStoreIndex, StorageContext
from llama_index.vector_stores.chroma import ChromaVectorStore
from llama_index.embeddings.huggingface import HuggingFaceEmbedding
from llama_index.core.vector_stores import (
    MetadataFilter,
    MetadataFilters,
    FilterCondition,
    FilterOperator,
)

from sentence_transformers import CrossEncoder

import transformers
transformers.logging.set_verbosity_error()

CHROMA_DB_PATH = "./chroma_db"
COLLECTION_NAME = "clinical_guidelines"
EMBED_MODEL_NAME = "BAAI/bge-small-en-v1.5"
RERANKER_MODEL_NAME = "BAAI/bge-reranker-base"

DEEPSEEK_API_KEY = os.environ.get("DEEPSEEK_API_KEY", "")
DEEPSEEK_BASE_URL = os.environ.get(
    "DEEPSEEK_BASE_URL", "https://api.deepseek.com/v1"
)
DEEPSEEK_MODEL = os.environ.get("DEEPSEEK_MODEL", "deepseek-chat")

CURRENT_YEAR = date.today().year

AVAILABLE_DOCS = {
    "2026 Guideline for the Early Management of Patients With Acute Ischemic Stroke: A Guideline From the American Heart Association/American Stroke Association": [
        2026
    ],
    "Guidelines for the Early Management of Patients With Acute Ischemic Stroke: 2019 Update to the 2018 Guidelines for the Early Management of Acute Ischemic Stroke: A Guideline for Healthcare Professionals From the American Heart Association/American Stroke Association": [
        2019
    ],
}

DOC_URLS = {
    "2026 Guideline for the Early Management of Patients With Acute Ischemic Stroke: A Guideline From the American Heart Association/American Stroke Association": "https://www.ahajournals.org/doi/10.1161/STR.0000000000000513",
    "Guidelines for the Early Management of Patients With Acute Ischemic Stroke: 2019 Update to the 2018 Guidelines for the Early Management of Acute Ischemic Stroke: A Guideline for Healthcare Professionals From the American Heart Association/American Stroke Association": "https://www.ahajournals.org/doi/full/10.1161/STR.0000000000000211",
}

deepseek_client: OpenAI | None = None


def call_llm(
    system_prompt: str,
    user_prompt: str,
    thinking: bool = False,
    temperature: float = 0.1,
) -> str:
    logger.info("[LLM] → Calling DeepSeek (thinking=%s, temperature=%s)", thinking, temperature)
    t0 = time.time()
    global deepseek_client
    if deepseek_client is None:
        if not DEEPSEEK_API_KEY:
            raise ValueError("API key not configured")
        deepseek_client = OpenAI(
            api_key=DEEPSEEK_API_KEY,
            base_url=DEEPSEEK_BASE_URL,
        )
    messages = [
        {"role": "system", "content": system_prompt},
        {"role": "user", "content": user_prompt},
    ]
    kwargs = {
        "model": DEEPSEEK_MODEL,
        "messages": messages,
        "temperature": temperature,
    }
    if thinking:
        kwargs["extra_body"] = {"thinking": {"type": "enabled"}, "reasoning_effort": "high"}
    response = deepseek_client.chat.completions.create(**kwargs)
    result = response.choices[0].message.content.strip()
    elapsed = time.time() - t0
    logger.info("[LLM] ← Response in %.1fs (%d chars)", elapsed, len(result))
    return result


def format_conversation_history(history: list) -> str:
    if not history:
        return ""
    lines = ["## Previous Conversation"]
    for m in history:
        role_label = "User" if m["role"] == "user" else "Assistant"
        lines.append(f"{role_label}: {m['content']}")
    return "\n\n".join(lines) + "\n\n"


@st.cache_resource
def init_vector_index():
    db = chromadb.PersistentClient(path=CHROMA_DB_PATH)
    chroma_collection = db.get_collection(COLLECTION_NAME)
    vector_store = ChromaVectorStore(chroma_collection=chroma_collection)
    embed_model = HuggingFaceEmbedding(model_name=EMBED_MODEL_NAME)
    return VectorStoreIndex.from_vector_store(
        vector_store, embed_model=embed_model
    )


@st.cache_resource
def load_cross_encoder():
    return CrossEncoder(RERANKER_MODEL_NAME)


@st.cache_data
def get_available_docs() -> list[tuple[str, int]]:
    db = chromadb.PersistentClient(path=CHROMA_DB_PATH)
    chroma_collection = db.get_collection(COLLECTION_NAME)
    result = chroma_collection.get(include=["metadatas"])
    seen = set()
    docs = []
    for meta in result["metadatas"]:
        key = (meta["doc_name"], meta["doc_year"])
        if key not in seen:
            seen.add(key)
            docs.append((meta["doc_name"], meta["doc_year"]))
    return docs


def query_chromadb(
    query_text: str, doc_name: str, doc_year: int, k: int = 5
) -> list:
    logger.info("[DB] → query_chromadb: %.60s… [%s (%d)]", query_text, doc_name, doc_year)
    index = init_vector_index()
    filters = MetadataFilters(
        filters=[
            MetadataFilter(
                key="doc_name",
                value=doc_name,
                operator=FilterOperator.EQ,
            ),
            MetadataFilter(
                key="doc_year",
                value=doc_year,
                operator=FilterOperator.EQ,
            ),
        ],
        condition=FilterCondition.AND,
    )
    retriever = index.as_retriever(similarity_top_k=k, filters=filters)
    results = []
    for node in retriever.retrieve(query_text):
        results.append(
            {
                "text": node.text,
                "doc_name": node.metadata.get("doc_name", ""),
                "doc_year": node.metadata.get("doc_year", 0),
                "score": getattr(node, "score", 0.0),
            }
        )
    logger.info("[DB] ← %d result(s)", len(results))
    return results


# =============================================================================
# 2. LangGraph State & Nodes
# =============================================================================
#
# The graph has 5 nodes executing in this order:
#
#   guard_agent (triages query relevance)
#     -> [if relevant] retrieval_agent (decomposes query, searches ChromaDB)
#                      -> re_ranker (CrossEncoder, keeps top 3)
#                        -> main_agent (synthesizes final answer)
#     -> [if irrelevant] end_inappropriate (returns out-of-scope message)
#


class GraphState(TypedDict):
    """State object passed between LangGraph nodes at each step.

    - user_query:          The original question from the doctor.
    - is_appropriate:      Set by GuardAgent — whether the query is clinically relevant.
    - fetched_chunks:      Accumulated raw chunks from ChromaDB; later filtered by re-ranker.
    - final_response:      The final answer produced by MainAgent (or the "can't answer" message).
    - debug_system_prompt:    System prompt sent to MainAgent (for debug dialog).
    - debug_user_prompt:      User prompt sent to MainAgent (for debug dialog).
    - retrieval_system_prompt: System prompt sent to RetrievalAgent (for debug).
    - retrieval_user_prompt:   User prompt sent to RetrievalAgent (for debug).
    - conversation_history:   Prior turns (list of {"role", "content"}) for context.
    - top_k_retrieval:     Number of chunks to fetch from ChromaDB per sub-query.
    - top_k_rerank:        Number of best chunks to keep after re-ranking.
    - temperature:         LLM temperature for main_agent synthesis.
    """
    user_query: str
    is_appropriate: bool
    fetched_chunks: list
    final_response: str
    debug_system_prompt: str
    debug_user_prompt: str
    retrieval_system_prompt: str
    retrieval_user_prompt: str
    retrieval_raw_response: str
    conversation_history: list
    top_k_retrieval: int
    top_k_rerank: int
    temperature: float


def main_agent(state: GraphState) -> dict:
    """Synthesis node.  Builds final answer from re-ranked chunks.
       Passes through if a final response was already set upstream."""
    logger.info("[AGENT] main_agent → start")
    if state.get("final_response"):
        logger.info("[AGENT] main_agent ← skipped (final_response already set)")
        return {}

    chunks_text = "\n\n---\n\n".join(
        [
            f"[Document: {c['doc_name']} ({c['doc_year']})]\n{c['text']}"
            for c in state["fetched_chunks"]
        ]
    )

    system_prompt = (
        "You are a clinical decision support AI assistant for doctors. "
        "Answer **exclusively** from the retrieved guideline chunks below. "
        "Do not use your own medical knowledge or training data — "
        "only the information provided in the chunks.\n\n"
        "If the chunks do not contain enough information to fully answer "
        "the question, clearly state what is missing rather than making up an answer.\n\n"
        "Each chunk is structured as:\n"
        "  Document: <guideline name>  (use this for citation)\n"
        "  Section: <heading hierarchy>  (use this for section reference)\n"
        "  <content>\n\n"
        "Cite the specific guideline year and relevant section.\n\n"
        "Be concise by default — doctors need key facts quickly. "
        "However, if the user's query explicitly asks for more detail, "
        "elaboration, or a comprehensive explanation, provide a thorough answer."
    )
    history_text = format_conversation_history(state.get("conversation_history", []))
    user_prompt = (
        f"{history_text}"
        f"## User Query\n\n{state['user_query']}\n\n"
        f"## Retrieved Guideline Chunks\n\n{chunks_text}\n\n"
        "Based on the above retrieved information, provide a medical answer for the doctor."
    )

    # Thinking mode ON for careful synthesis
    temp = state.get("temperature", 0.1)
    response = call_llm(system_prompt, user_prompt, thinking=True, temperature=temp)
    logger.info("[AGENT] main_agent ← done (%d chars)", len(response))
    return {
        "final_response": response,
        "debug_system_prompt": system_prompt,
        "debug_user_prompt": user_prompt,
    }


def guard_agent(state: GraphState) -> dict:
    """Guardrail node.  Uses a fast LLM call (thinking OFF) to decide if the
       query belongs to clinical medicine.  Sets `is_appropriate` in state."""
    logger.info("[AGENT] guard_agent → start")
    system_prompt = (
        "You are a medical triage guard for a clinical decision support "
        "system used by doctors. Determine if a user query is relevant to "
        "clinical medicine, patient management, or medical guidelines.\n\n"
        "Ignore years, document names, or specific guideline versions — "
        "that is checked downstream. Only assess medical relevance.\n\n"
        'Respond with ONLY "YES" if clinically relevant, or "NO" if not. '
        "YES: blood pressure, medication, patient care, "
        "symptoms, diagnosis, treatment, rehabilitation, imaging.\n"
        "NO: sports, politics, entertainment, programming, weather, "
        "history, geography, general knowledge."
    )
    history_text = format_conversation_history(state.get("conversation_history", []))
    user_prompt = (
        f"{history_text}"
        f"Query: {state['user_query']}\n\n"
        "Is this query relevant to clinical medicine? Answer YES or NO only."
    )
    # Thinking mode OFF = minimum latency for this simple classification
    response = call_llm(
        system_prompt, user_prompt, thinking=False, temperature=0.0
    )
    is_appropriate = response.strip().upper().startswith("YES")
    logger.info("[AGENT] guard_agent ← done (appropriate=%s)", is_appropriate)
    return {"is_appropriate": is_appropriate}


def end_inappropriate(state: GraphState) -> dict:
    """Terminal node for out-of-scope queries.  Sets a polite refusal message."""
    logger.info("[AGENT] end_inappropriate → start")
    logger.info("[AGENT] end_inappropriate ← done")
    return {
        "final_response": (
        "I can't answer this type of question. Please ask a "
        "clinical/medical question related to the available healthcare "
        "guidelines."
        )
    }


def retrieval_agent(state: GraphState) -> dict:
    """Query-decomposition node.  LLM (thinking ON) analyses the user query
       and outputs a JSON list of sub-queries.  Each sub-query specifies
       keywords, doc_name, and doc_year for a `query_chromadb` call.

       Handles temporal reasoning: explicit years, relative terms like
       "latest", or defaults to the newest version, and validates that
       every sub-query targets an available document.
    """
    logger.info("[AGENT] retrieval_agent → start")
    docs_info = "\n".join(
        [f'- "{name}" (years: {years})' for name, years in AVAILABLE_DOCS.items()]
    )

    system_prompt = (
        "You are a medical retrieval specialist for healthcare guidelines. "
        "Conversation history is provided below for context. "
        "Your ONLY task is to decompose the LATEST user query into vector "
        "database searches. Ignore the assistant responses in history.\n\n"
        f"Current year: {CURRENT_YEAR}\n\n"
        f"Available documents:\n{docs_info}\n\n"
        "Output a JSON object with a 'sub_queries' array. Each item:\n"
        '- "query_text": clinical keywords for semantic search\n'
        '- "doc_name": exact document name string\n'
        '- "doc_year": integer year\n\n'
        "Rules:\n"
        "1. If user specifies years (e.g., '2019 vs 2026'), create separate "
        "sub_queries for each year.\n"
        f"2. If user says 'latest'/'current', use {CURRENT_YEAR}.\n"
        f"3. If no year specified, use the latest ({CURRENT_YEAR}).\n"
        "4. Extract concise clinical keywords optimized for vector search.\n"
        "5. If the query does not match any available document, return "
        '{"sub_queries": []}.\n'
        "6. If the user specifies a year that is not in the available years "
        "for a document, return {\"sub_queries\": []} — do not substitute "
        "a different year."
    )
    history_text = format_conversation_history(state.get("conversation_history", []))
    user_prompt = (
        f"{history_text}"
        f"## Latest User Query\n\n{state['user_query']}\n\n"
        "Based on the conversation history and the latest user query above, "
        "output a JSON object with a 'sub_queries' array for the latest query."
    )

    # Thinking mode ON ensures accurate extraction of search arguments
    response = call_llm(system_prompt, user_prompt, thinking=True, temperature=0.0)


    def extract_json(text: str) -> dict | None:
        """Extract the first JSON object from text, even if surrounded by
           conversational wrapper text."""
        text = text.strip()
        # Strip markdown code fences
        if "```json" in text:
            text = text.split("```json")[1].split("```")[0].strip()
        elif "```" in text:
            text = text.split("```")[1].split("```")[0].strip()
        # Find the outermost { … }
        start = text.find("{")
        end = text.rfind("}")
        if start == -1 or end == -1 or end <= start:
            return None
        try:
            return json.loads(text[start : end + 1])
        except json.JSONDecodeError:
            return None

    parsed = extract_json(response)
    try:
        if parsed is None:
            raise ValueError
        sub_queries = parsed.get("sub_queries", [])
        if not isinstance(sub_queries, list) or not sub_queries:
            raise ValueError
    except (ValueError, TypeError, AttributeError):
        logger.info("[AGENT] retrieval_agent ← no sub_queries could be extracted")
        return {
            "final_response": "Sorry, I am unable to utilize my knowledge base.",
            "fetched_chunks": [],
            "retrieval_system_prompt": system_prompt,
            "retrieval_user_prompt": user_prompt,
            "retrieval_raw_response": response,
        }

    # Validate every sub-query against the available documents
    fallback_response = {
        "final_response": "Sorry, that is outside my current knowledge base.",
        "fetched_chunks": [],
        "retrieval_system_prompt": system_prompt,
        "retrieval_user_prompt": user_prompt,
        "retrieval_raw_response": response,
    }

    try:
        for sq in sub_queries:
            # If sq isn't a dict, .get() raises an AttributeError
            doc_name = sq.get("doc_name", "")
            doc_year = sq.get("doc_year", 0)
            
            # Just to check if query_text exists.
            query_text = sq.get("query_text")
            
            # If doc_name or doc_year don't exist, it raises a KeyError
            if doc_year not in AVAILABLE_DOCS[doc_name]:
                logger.info(
                    "[AGENT] retrieval_agent ← year %d not available for doc '%s'",
                    doc_year, doc_name,
                )
                return fallback_response
                
    except (AttributeError, KeyError):
        logger.info("[AGENT] retrieval_agent ← sub_query validation failed (unknown doc/field)")
        return fallback_response

    all_chunks = []
    for sq_idx, sq in enumerate(sub_queries):
        doc_name = sq.get("doc_name")
        doc_year = sq.get("doc_year")
        query_text = sq.get("query_text")
        results = query_chromadb(query_text, doc_name, doc_year, k=state.get("top_k_retrieval", 5))
        for r in results:
            r["sq_idx"] = sq_idx
            all_chunks.append(r)

    logger.info("[AGENT] retrieval_agent ← done (%d sub_queries, %d chunks)", len(sub_queries), len(all_chunks))
    return {
        "fetched_chunks": all_chunks,
        "retrieval_system_prompt": system_prompt,
        "retrieval_user_prompt": user_prompt,
        "retrieval_raw_response": response,
    }


def re_ranker(state: GraphState) -> dict:
    """Pure-Python re-ranker node (no LLM call).  Uses a CrossEncoder
       to score every fetched chunk against the original query, keeps
       only the top 3 most relevant results.
       Passes through if a final response was already set upstream."""
    logger.info("[AGENT] re_ranker → start (input: %d chunks)", len(state.get("fetched_chunks", [])))
    if state.get("final_response"):
        logger.info("[AGENT] re_ranker ← skipped (final_response already set)")
        return {}

    cross_encoder = load_cross_encoder()

    grouped: dict[int, list] = {}
    for chunk in state["fetched_chunks"]:
        grouped.setdefault(chunk["sq_idx"], []).append(chunk)

    selected = []
    for group in grouped.values():
        passages = [c["text"] for c in group]
        if not passages:
            continue
        scores = cross_encoder.predict([(state["user_query"], p) for p in passages])
        scored = list(zip(scores, group))
        scored.sort(key=lambda x: x[0], reverse=True)
        k_rerank = state.get("top_k_rerank", 3)
        selected.extend(c for _, c in scored[:k_rerank])

    logger.info("[AGENT] re_ranker ← done (%d chunks after filtering)", len(selected))
    return {"fetched_chunks": selected}


# =============================================================================
# 3. Workflow Compilation
# =============================================================================
#
#  LangGraph flow:
#
#    START
#      |
#      v
#   [guard_agent]
#      |
#      ├──[if relevant]──→ [retrieval_agent] → [re_ranker] → [main_agent] → END
#      └──[if irrelevant]─→ [end_inappropriate] → END


@st.cache_resource
def build_graph():
    """Build and compile the LangGraph state machine.
       Cached so the graph topology is only constructed once."""
    builder = StateGraph(GraphState)

    # Register all 5 nodes
    builder.add_node("main_agent", main_agent)
    builder.add_node("guard_agent", guard_agent)
    builder.add_node("end_inappropriate", end_inappropriate)
    builder.add_node("retrieval_agent", retrieval_agent)
    builder.add_node("re_ranker", re_ranker)

    # First node to execute
    builder.set_entry_point("guard_agent")

    # -- guard_agent routes --
    # If query is clinically relevant → proceed to retrieval_agent.
    # Otherwise → skip retrieval and go straight to end_inappropriate.
    def route_guard(state):
        return "retrieval_agent" if state["is_appropriate"] else "end_inappropriate"

    builder.add_conditional_edges(
        "guard_agent",
        route_guard,
        {
            "retrieval_agent": "retrieval_agent",
            "end_inappropriate": "end_inappropriate",
        },
    )

    # -- deterministic edges --
    builder.add_edge("retrieval_agent", "re_ranker")   # retrieve → re-rank
    builder.add_edge("re_ranker", "main_agent")        # re-rank → synthesise
    builder.add_edge("main_agent", END)                # synthesis → end
    builder.add_edge("end_inappropriate", END)         # terminal node

    return builder.compile()


# =============================================================================
# 4. Streamlit Application
# =============================================================================


def init_session_state():
    if "thread_id" not in st.session_state:
        st.session_state.thread_id = str(uuid.uuid4())
    if "messages" not in st.session_state:
        st.session_state.messages = []
    if "processing" not in st.session_state:
        st.session_state.processing = False
    if "debug_transcript" not in st.session_state:
        st.session_state.debug_transcript = []
    if "quota_exceeded" not in st.session_state:
        st.session_state.quota_exceeded = False
    if "_quota_error_detected" not in st.session_state:
        st.session_state._quota_error_detected = False
    if "_text_to_copy" not in st.session_state:
        st.session_state._text_to_copy = ""
    if "top_k_retrieval" not in st.session_state:
        st.session_state.top_k_retrieval = 16
    if "top_k_rerank" not in st.session_state:
        st.session_state.top_k_rerank = 8
    if "temperature" not in st.session_state:
        st.session_state.temperature = 0.1


def _request_copy(text: str):
    st.session_state._text_to_copy = text


def _copy_conversation():
    lines = []
    for msg in st.session_state.messages:
        role = "You" if msg["role"] == "user" else "Assistant"
        lines.append(f"{role}: {msg['content']}")
    _request_copy("\n\n---\n\n".join(lines))


def _render_copy_script():
    text = st.session_state.pop("_text_to_copy", "")
    if not text.strip():
        return
    import html
    srcdoc = f"""<html><body><script>
var text = {json.dumps(text)};
var input = parent.document.createElement('textarea');
input.innerHTML = text;
parent.document.body.appendChild(input);
input.select();
parent.document.execCommand('copy');
parent.document.body.removeChild(input);
</script></body></html>"""
    st.markdown(
        f'<iframe srcdoc="{html.escape(srcdoc, quote=True)}" width=0 height=0></iframe>',
        unsafe_allow_html=True,
    )
    st.toast("Copied!", icon="📋")


def reset_session():
    prefs = {
        "top_k_retrieval": st.session_state.get("top_k_retrieval", 16),
        "top_k_rerank": st.session_state.get("top_k_rerank", 8),
        "temperature": st.session_state.get("temperature", 0.1),
    }
    st.session_state.thread_id = str(uuid.uuid4())
    st.session_state.messages = []
    st.session_state.processing = False
    st.session_state.debug_transcript = []
    st.session_state.quota_exceeded = False
    st.session_state._quota_error_detected = False
    st.session_state.top_k_retrieval = prefs["top_k_retrieval"]
    st.session_state.top_k_rerank = prefs["top_k_rerank"]
    st.session_state.temperature = prefs["temperature"]
@st.dialog("📜 Debug Transcript", width="large")
def show_debug_transcript():
    for entry in st.session_state.debug_transcript:
        st.markdown(f"### Turn {entry['turn']}: {entry['user_query']}")
        guard_icon = "✅" if entry["guard_decision"] else "❌"
        st.markdown(f"**Guard:** {guard_icon} {'Relevant' if entry['guard_decision'] else 'Irrelevant'} | **Chunks:** {entry['num_chunks']}")
        st.markdown("**Retrieval Agent — System Prompt:**")
        st.code(entry["retrieval_system_prompt"])
        st.markdown("**Retrieval Agent — User Prompt:**")
        st.code(entry["retrieval_user_prompt"])
        st.markdown("**Retrieval Agent — Decomposed Sub-queries:**")
        st.code(entry["retrieval_raw_response"])
        st.markdown("**Main Agent — System Prompt:**")
        st.code(entry["main_system_prompt"])
        st.markdown("**Main Agent — User Prompt:**")
        st.code(entry["main_user_prompt"])
        st.markdown(f"**Response:** {entry['final_response']}")
        st.divider()
    if st.button("Close"):
        st.rerun()


@st.dialog("🏗️ Architecture", width="large")
def show_architecture():
    left, right = st.columns(2)
    with left:
        try:
            st.image("assets/architecture.png", width=350)
        except Exception:
            st.caption("Unable to render graph diagram.")
    with right:
        st.markdown("### ChromaDB Schema")
        st.markdown(
            "| Column | Type | Description |\n"
            "|---|---|---|\n"
            "| `id` | string | Auto-generated UUID |\n"
            "| `embedding` | `List[float]` | 384-dim vector from `BAAI/bge-small-en-v1.5` |\n"
            "| `document` | string | Enriched chunk text |\n"
            "| `metadata` | dict | `{\"doc_year\": int, \"doc_name\": string}` |\n"
        )
    if st.button("Close"):
        st.rerun()


SAMPLES = [
    {
        "desc": "Can't answer an out-of-scope question (guard fails)",
        "question": "Why did President Trump put tariff on everyone?",
    },
    {
        "desc": "Can't answer a clinically relevant question if it doesn't match available docs (retrieval fails)",
        "question": "Tell me about Antiplatelet Treatment for patients with minor Acute Ischemic Stroke in Guidelines 2015.",
    },
    {
        "desc": "Answer using one call to ChromaDB (vector database), filtering with doc_year=2019",
        "question": "Tell me about Antiplatelet Treatment for patients with minor Acute Ischemic Stroke in Guidelines 2019.",
    },
    {
        "desc": "Test follow-up question: 'Does the treatment have side effects?'",
        "question": "Tell me about Antiplatelet Treatment for patients with minor Acute Ischemic Stroke.",
    },
    {
        "desc": "Compare treatment between two medical documents (2019, 2026), filtering with doc_year=2019 and then doc_year=2026",
        "question": "How did blood pressure management recommendations change after thrombolysis or thrombectomy between 2026 and 2019 guidelines?",
    },
]


@st.dialog("❓ Help & Examples", width="large")
def show_help_dialog():
    st.markdown(
        "Each sample below demonstrates a different type of question "
        "the system can answer. Click **Use this question** to try it."
    )
    st.markdown(
        "📖 For full documentation, visit the "
        "[README](https://huggingface.co/spaces/tnt306/agentic_retrieval/blob/main/README.md)."
    )
    st.markdown(
        "<style>"
        'div[data-testid="stDialog"] hr {margin:0.1rem 0!important;}'
        'div[data-testid="stDialog"] .stButton button {margin-bottom:0!important;}'
        'div[data-testid="stDialog"] p {margin-bottom:0!important;}'
        "</style>",
        unsafe_allow_html=True,
    )
    st.markdown('<hr style="margin:0.1rem 0;">', unsafe_allow_html=True)
    for sample in SAMPLES:
        st.markdown(
            f"**{sample['question']}**<br>"
            f'<span style="color:#6b7280;">👉 '
            f'{sample["desc"]}.</span>',
            unsafe_allow_html=True,
        )
        if st.button("Use this question", key=f"sample_{hash(sample['question'])}"):
            reset_session()
            st.session_state.chat_input = sample["question"]
            st.rerun()
        st.markdown('<hr style="margin:0.1rem 0;">', unsafe_allow_html=True)
    if st.button("Close"):
        st.rerun()


def main():
    st.set_page_config(
        page_title="Clinical Decision Support - Healthcare Guidelines",
        page_icon="⚡",
        layout="wide",
    )

    st.markdown(
        "<style>"
        "[data-testid='stSidebarCollapseButton']{visibility:visible!important;}"
        "[data-testid='stCode'] pre,[data-testid='stCode'] code{white-space:pre-wrap!important;word-break:break-word!important;overflow-x:hidden!important;}"
        "[data-testid='stCode'] pre{max-height:200px;overflow-y:auto!important;}"
        "[data-testid='stCode']{border:1px solid rgba(0,0,0,0.15)!important;border-radius:6px!important;}"
        "[data-testid='stAppDeployButton']{display:none!important;}"
        "[data-testid='stMainMenu']{display:none!important;}"
        "</style>",
        unsafe_allow_html=True,
    )

    init_session_state()

    with st.sidebar:
        st.title("⚡ Agentic Retrieval")
        st.markdown("**Clinical Decision Support System**")
        st.divider()

        if st.button("🗑️ New Session", use_container_width=True, type="primary"):
            reset_session()
            st.rerun()

        col1, col2 = st.columns(2)
        with col1:
            if st.button("📜 Debug", use_container_width=True, disabled=len(st.session_state.debug_transcript) == 0):
                show_debug_transcript()
        with col2:
            if st.button("📋 Copy", use_container_width=True):
                _copy_conversation()

        col3, col4 = st.columns(2)
        with col3:
            if st.button("🏗️ Architecture", use_container_width=True):
                show_architecture()
        with col4:
            if st.button("❓ Examples", use_container_width=True):
                show_help_dialog()

        st.checkbox(
            "⚠️ Simulate quota exceeded",
            value=st.session_state.get("quota_exceeded", False),
            key="quota_exceeded",
            help="Show the quota-exceeded warning without actually hitting the API.",
        )

        with st.expander("⚙️ Settings"):
            st.number_input("Chunks to retrieve", 1, 20, key="top_k_retrieval")
            st.number_input("Chunks after re-ranking", min_value=1, max_value=st.session_state.top_k_retrieval, key="top_k_rerank")
            st.slider("LLM temperature", 0.0, 1.0, key="temperature", step=0.05)

        if not DEEPSEEK_API_KEY:
            st.warning(
                "⚠️ DEEPSEEK_API_KEY not set. Set it in your environment.",
                icon="⚠️",
            )

    st.markdown("### 🏥 Healthcare Guideline Assistant")
    st.caption(
        "Multi-Agent RAG | LangGraph · ChromaDB · BGE Reranker · DeepSeek V4"
    )

    with st.expander("📚 Medical Documents Covered"):
        docs = get_available_docs()
        for doc_name, doc_year in sorted(docs, key=lambda x: -x[1]):
            url = DOC_URLS.get(doc_name, "")
            if url:
                st.markdown(f"- [**{doc_name}** ({doc_year})]({url})")
            else:
                st.markdown(f"- **{doc_name}** ({doc_year})")

    for i, msg in enumerate(st.session_state.messages):
        with st.chat_message(msg["role"]):
            st.markdown(msg["content"])
        if msg["role"] == "assistant":
            label = f"📋 Copy  ⏱ {msg['response_time']:.1f}s" if msg.get("response_time") else "📋 Copy"
            st.button(label, key=f"copy_{i}", on_click=_request_copy, args=(msg["content"],))

    prompt = st.chat_input(
        "Ask a clinical question about healthcare guidelines...",
        key="chat_input",
        disabled=st.session_state.processing or st.session_state.get("quota_exceeded") or st.session_state.get("_quota_error_detected"),
    )

    if not DEEPSEEK_API_KEY:
        st.markdown(
            '<div style="'
            "background:#fef2f2;border:1px solid #fca5a5;"
            "border-radius:6px;padding:6px 12px;margin-bottom:8px;"
            '">'
            '<span style="color:#dc2626;font-size:0.85em;">'
            "⚠️ The API key is not configured."
            "</span></div>",
            unsafe_allow_html=True,
        )
    elif st.session_state.get("quota_exceeded") or st.session_state.get("_quota_error_detected"):
        st.markdown(
            '<div style="'
            "background:#fef2f2;border:1px solid #fca5a5;"
            "border-radius:6px;padding:6px 12px;margin-bottom:8px;"
            '">'
            '<span style="color:#dc2626;font-size:0.85em;">'
            "⚠️ The LLM is currently unavailable due to quota limits."
            "</span></div>",
            unsafe_allow_html=True,
        )

    if prompt:
        st.session_state.messages.append({"role": "user", "content": prompt})
        with st.chat_message("user"):
            st.markdown(prompt)

        st.session_state.processing = True
        t_start = time.time()

        graph = build_graph()
        initial_state: GraphState = {
            "user_query": prompt,
            "conversation_history": [
                {"role": m["role"], "content": m["content"]}
                for m in st.session_state.messages[:-1]
            ],
            "is_appropriate": False,
            "fetched_chunks": [],
            "final_response": "",
            "debug_system_prompt": "",
            "debug_user_prompt": "",
            "retrieval_system_prompt": "",
            "retrieval_user_prompt": "",
            "retrieval_raw_response": "",
            "top_k_retrieval": st.session_state.top_k_retrieval,
            "top_k_rerank": st.session_state.top_k_rerank,
            "temperature": st.session_state.temperature,
        }

        with st.chat_message("assistant"):
            response = None
            error_msg = None

            with st.status("Starting...", expanded=False) as status:
                try:
                    final_state = {}
                    for step in graph.stream(initial_state):
                        for node_name, node_data in step.items():
                            if node_data is None:
                                continue
                            if node_name == "guard_agent":
                                is_appropriate = node_data.get("is_appropriate", False)
                                if is_appropriate:
                                    status.update(label="✅ Query is clinically relevant", state="complete")
                                else:
                                    status.update(label="❌ Query is out of scope", state="error")
                            elif node_name == "retrieval_agent":
                                chunks = node_data.get("fetched_chunks", [])
                                status.update(label=f"🔍 Retrieved {len(chunks)} guideline chunks", state="complete")
                            elif node_name == "re_ranker":
                                chunks = node_data.get("fetched_chunks", [])
                                status.update(label=f"🎯 Selected top {len(chunks)} most relevant chunks", state="complete")
                            elif node_name == "main_agent":
                                status.update(label="🧠 Synthesizing clinical answer...", state="running")
                            elif node_name == "end_inappropriate":
                                status.update(label="❌ Cannot answer this query", state="error")
                            final_state.update(node_data)

                    st.session_state.debug_transcript.append({
                        "turn": len(st.session_state.debug_transcript) + 1,
                        "user_query": prompt,
                        "guard_decision": final_state.get("is_appropriate"),
                        "num_chunks": len(final_state.get("fetched_chunks", [])),
                        "main_system_prompt": final_state.get("debug_system_prompt", ""),
                        "main_user_prompt": final_state.get("debug_user_prompt", ""),
                        "retrieval_system_prompt": final_state.get("retrieval_system_prompt", ""),
                        "retrieval_user_prompt": final_state.get("retrieval_user_prompt", ""),
                        "retrieval_raw_response": final_state.get("retrieval_raw_response", ""),
                        "final_response": final_state.get("final_response", ""),
                    })

                    response = final_state.get(
                        "final_response",
                        "I'm sorry, I couldn't process that query.",
                    )
                    status.update(label="✅ Complete", state="complete")
                except openai.RateLimitError:
                    st.session_state._quota_error_detected = True
                    error_msg = "The LLM is currently unavailable due to quota limits."
                    status.update(label="❌ Quota exceeded", state="error")
                except Exception as e:
                    error_msg = (
                        "I encountered an error processing your request: "
                        f"{str(e)}"
                    )
                    status.update(label="❌ Error", state="error")

            elapsed = time.time() - t_start

            if response:
                def word_stream(text):
                    words = text.split(" ")
                    for i, word in enumerate(words):
                        yield word + (" " if i < len(words) - 1 else "")
                        time.sleep(0.02)
                st.write_stream(word_stream(response))
                st.session_state.messages.append(
                    {"role": "assistant", "content": response, "response_time": elapsed}
                )
            elif error_msg:
                st.error(error_msg)
                st.session_state.messages.append(
                    {"role": "assistant", "content": error_msg}
                )

        if response:
            st.button(f"📋 Copy  ⏱ {elapsed:.1f}s", key="copy_new", on_click=_request_copy, args=(response,))

        st.session_state.processing = False

        st.rerun()

    _render_copy_script()


if __name__ == "__main__":
    main()