import os

import re
import gradio as gr
from PIL import Image

from cpu_agent import (
    full_agent,
    save_retriever_from_pdf,
    load_persisted_retriever,
    VECTOR_STORE_PATH,
)


# ==========================================
# BACKEND: TAB 1 — ADMIN KNOWLEDGE BASE
# ==========================================
def initialize_knowledge_base(pdf_file, legal_consent):
    if not legal_consent:
        yield "⛔ Cannot proceed: you must accept the legal consent declaration before initialising the knowledge base."
        return
    if pdf_file is None:
        yield "⛔ No PDF uploaded. Please upload your institution's licensed oncology guideline document."
        return
    pdf_path = pdf_file if isinstance(pdf_file, str) else pdf_file.name
    filename = os.path.basename(pdf_path)
    yield f"📄 Received: {filename}\n⏳ Extracting text and building vector index — please wait..."
    success, message = save_retriever_from_pdf(pdf_path)
    if success:
        yield (
            f"{message}\n\n"
            f"✅ Knowledge base is ready.\n"
            f"Clinicians can now use the Patient Consultation tab."
        )
    else:
        yield (
            f"{message}\n\n"
            f"⚠️ Falling back to built-in NCCN/ESMO guideline excerpts."
        )


def get_kb_status():
    r, label = load_persisted_retriever()
    if r:
        return f"✅ Active: {label}"
    return "⚠️ No knowledge base configured. Upload a PDF in the Configuration tab first."


# ==========================================
# BACKEND: TAB 2 — PATIENT CONSULTATION
# ==========================================
def process_patient_data(user_query, lab_values, behaviour_changes, wsi_image):
    parts = []
    if lab_values.strip():
        parts.append(f"Lab Values:\n{lab_values.strip()}")
    if behaviour_changes.strip():
        parts.append(f"Behaviour / Symptoms:\n{behaviour_changes.strip()}")
    raw_clinical_text = "\n\n".join(parts) if parts else "No clinical data provided."

    wsi_path = ""
    if wsi_image is not None:
        wsi_path = "/tmp/uploaded_wsi.bmp"
        wsi_image.save(wsi_path)

    initial_state = {
        "patient_id":           "VAJRAM_UI",
        "user_query":           user_query.strip() or "Give me a full clinical workup and treatment plan.",
        "raw_clinical_text":    raw_clinical_text,
        "modules_queue":        [],
        "wsi_image_path":       wsi_path,
        "wsi_output_path":      "/tmp/annotated_wsi_output.png",
        "module2_risk_score":   "",
        "module3_wsi_analysis": "",
        "module4_progression":  "",
        "module5_guidelines":   "",
        "module5_raw_chunks":   "",
        "module5_source":       "",
        "final_recommendation": "",
    }

    final_state = full_agent.invoke(initial_state)

    ran_module2 = bool(final_state.get("module2_risk_score",   "").strip())
    ran_module3 = bool(final_state.get("module3_wsi_analysis", "").strip())
    ran_module4 = bool(final_state.get("module4_progression",  "").strip())
    ran_module5 = bool(final_state.get("module5_guidelines",   "").strip())

    selected = []
    if ran_module2: selected.append("Module 2 · Risk Assessment")
    if ran_module3: selected.append("Module 3 · Bone Marrow WSI")
    if ran_module4: selected.append("Module 4 · Progression Tracking")
    if ran_module5: selected.append("Module 5 · Guideline Retrieval")
    modules_selected_str = "  ›  ".join(selected) if selected else "None selected"

    annotated_img   = None
    malignancy_stat = gr.update(visible=False)

    if ran_module3 and wsi_path:
        try:
            annotated_img = Image.open("/tmp/annotated_wsi_output.png")
            wsi_summary   = final_state.get("module3_wsi_analysis", "")
            pct_str       = ""
            m = re.search(r'(\d+\.?\d*)\s*%\s*\)', wsi_summary)
            if m:
                pct_str = m.group(1)
            counts_m = re.search(r'(\d+)/(\d+) patches', wsi_summary)
            if counts_m:
                n_mal, n_total = counts_m.group(1), counts_m.group(2)
                malignancy_stat = gr.update(
                    value=(
                        f"Malignant: {n_mal} / {n_total} patches  ({pct_str}%)\n"
                        f"Red = Malignant  ·  Green = Normal  ·  Gray = Background  ·  Yellow = Unknown"
                    ),
                    visible=True
                )
            else:
                malignancy_stat = gr.update(value=wsi_summary, visible=True)
        except Exception:
            annotated_img = None

    m5_source = final_state.get("module5_source", "")
    m5_source_update = gr.update(
        value=f"Source: {m5_source}" if m5_source else "",
        visible=ran_module5
    )
    raw_chunks_update = gr.update(
        value=final_state.get("module5_raw_chunks", ""),
        visible=ran_module5
    )

    return (
        modules_selected_str,
        final_state.get("final_recommendation", "No recommendation generated."),
        gr.update(value=final_state.get("module2_risk_score",   ""), visible=ran_module2),
        gr.update(visible=ran_module2),
        gr.update(value=final_state.get("module3_wsi_analysis", ""), visible=ran_module3),
        gr.update(visible=ran_module3),
        gr.update(value=annotated_img, visible=annotated_img is not None),
        malignancy_stat,
        gr.update(value=final_state.get("module4_progression",  ""), visible=ran_module4),
        gr.update(visible=ran_module4),
        gr.update(value=final_state.get("module5_guidelines",   ""), visible=ran_module5),
        gr.update(visible=ran_module5),
        m5_source_update,
        raw_chunks_update,
    )


EXAMPLE_QUERY     = "What is the recommended treatment for this transplant-eligible myeloma patient with renal impairment?"
EXAMPLE_LABS      = (
    "Creatinine: 2.5 mg/dL\nCalcium: 10.5 mg/dL\nHemoglobin: 9.0 g/dL\n"
    "Beta-2 Microglobulin: 5.8 mg/L\nLDH: 280 U/L\n"
    "M-Spike: increased 1.2 g/dL over 3 months\n"
    "Patient: 65-year-old Male. Transplant eligible."
)
EXAMPLE_BEHAVIOUR = "Dizziness, Chest Pain, Mental Confusion"


# ==========================================
# CUSTOM CSS — Medical Luxury Dark Theme
# Fonts: Playfair Display (headers) + IBM Plex Mono (data)
# Palette: Deep navy slate · Warm ivory text · Amber accent
# ==========================================
CUSTOM_CSS = """
@import url('https://fonts.googleapis.com/css2?family=Playfair+Display:wght@400;600;700&family=IBM+Plex+Mono:wght@300;400;500&family=IBM+Plex+Sans:wght@300;400;500&display=swap');

/* ── Root palette ─────────────────────────────────────────────────────────── */
:root {
    --bg-void:        #080c14;
    --bg-deep:        #0d1320;
    --bg-panel:       #111827;
    --bg-card:        #161f30;
    --bg-input:       #1a2438;
    --bg-hover:       #1e2d45;
    --border-dim:     #1e2d45;
    --border-mid:     #2a3f5f;
    --border-bright:  #3d5a80;
    --text-ivory:     #f0ead8;
    --text-muted:     #8a9bb8;
    --text-faint:     #4a5d7a;
    --accent-amber:   #c8963c;
    --accent-amber-dim: #8a6422;
    --accent-teal:    #3d9e8c;
    --accent-teal-dim: #1e4f47;
    --danger:         #c84c3c;
    --success:        #3d9e5a;
    --font-display:   'Playfair Display', Georgia, serif;
    --font-data:      'IBM Plex Mono', 'Courier New', monospace;
    --font-body:      'IBM Plex Sans', system-ui, sans-serif;
}

/* ── Global reset ─────────────────────────────────────────────────────────── */
*, *::before, *::after { box-sizing: border-box; }

body, .gradio-container {
    background: var(--bg-void) !important;
    color: var(--text-ivory) !important;
    font-family: var(--font-body) !important;
    font-weight: 300 !important;
}

.gradio-container {
    max-width: 1400px !important;
    margin: 0 auto !important;
    padding: 0 24px 48px !important;
}

/* ── Header block ─────────────────────────────────────────────────────────── */
.header-block {
    border-bottom: 1px solid var(--border-mid);
    padding: 40px 0 28px;
    margin-bottom: 32px;
    position: relative;
}

.header-block::before {
    content: '';
    position: absolute;
    bottom: -1px; left: 0;
    width: 80px; height: 2px;
    background: var(--accent-amber);
}

/* ── Markdown headings ────────────────────────────────────────────────────── */
.gradio-container h1, .prose h1 {
    font-family: var(--font-display) !important;
    font-size: 2.4rem !important;
    font-weight: 700 !important;
    color: var(--text-ivory) !important;
    letter-spacing: -0.02em !important;
    line-height: 1.15 !important;
    margin: 0 0 8px !important;
}

.gradio-container h2, .prose h2 {
    font-family: var(--font-display) !important;
    font-size: 1.35rem !important;
    font-weight: 600 !important;
    color: var(--accent-amber) !important;
    letter-spacing: 0.04em !important;
    text-transform: uppercase !important;
    margin: 32px 0 16px !important;
    padding-bottom: 8px !important;
    border-bottom: 1px solid var(--border-dim) !important;
}

.gradio-container h3, .prose h3 {
    font-family: var(--font-body) !important;
    font-size: 0.78rem !important;
    font-weight: 500 !important;
    color: var(--text-muted) !important;
    letter-spacing: 0.12em !important;
    text-transform: uppercase !important;
    margin: 24px 0 12px !important;
}

.gradio-container h4, .prose h4 {
    font-family: var(--font-body) !important;
    font-size: 0.72rem !important;
    font-weight: 500 !important;
    color: var(--accent-teal) !important;
    letter-spacing: 0.14em !important;
    text-transform: uppercase !important;
    margin: 20px 0 10px !important;
}

.gradio-container p, .prose p {
    color: var(--text-muted) !important;
    font-size: 0.88rem !important;
    line-height: 1.7 !important;
}

.gradio-container strong, .prose strong {
    color: var(--text-ivory) !important;
    font-weight: 500 !important;
}

/* ── Tab navigation ───────────────────────────────────────────────────────── */
.tab-nav {
    background: var(--bg-deep) !important;
    border-bottom: 1px solid var(--border-dim) !important;
    padding: 0 !important;
    gap: 0 !important;
}

.tab-nav button {
    font-family: var(--font-body) !important;
    font-size: 0.75rem !important;
    font-weight: 500 !important;
    letter-spacing: 0.1em !important;
    text-transform: uppercase !important;
    color: var(--text-faint) !important;
    background: transparent !important;
    border: none !important;
    border-bottom: 2px solid transparent !important;
    padding: 16px 28px !important;
    transition: all 0.2s ease !important;
    border-radius: 0 !important;
}

.tab-nav button:hover {
    color: var(--text-muted) !important;
    background: var(--bg-hover) !important;
}

.tab-nav button.selected {
    color: var(--accent-amber) !important;
    border-bottom-color: var(--accent-amber) !important;
    background: transparent !important;
}

/* ── Panels and cards ─────────────────────────────────────────────────────── */
.panel, .gradio-group, .gr-group {
    background: var(--bg-panel) !important;
    border: 1px solid var(--border-dim) !important;
    border-radius: 4px !important;
    padding: 20px !important;
}

/* ── Labels ───────────────────────────────────────────────────────────────── */
label span, .gradio-container label {
    font-family: var(--font-body) !important;
    font-size: 0.72rem !important;
    font-weight: 500 !important;
    letter-spacing: 0.1em !important;
    text-transform: uppercase !important;
    color: var(--text-muted) !important;
}

/* ── Text inputs and textareas ────────────────────────────────────────────── */
textarea, input[type="text"], .gradio-textbox textarea {
    background: var(--bg-input) !important;
    border: 1px solid var(--border-dim) !important;
    border-radius: 3px !important;
    color: var(--text-ivory) !important;
    font-family: var(--font-data) !important;
    font-size: 0.82rem !important;
    font-weight: 300 !important;
    line-height: 1.6 !important;
    padding: 12px 14px !important;
    transition: border-color 0.2s ease !important;
    caret-color: var(--accent-amber) !important;
}

textarea:focus, input[type="text"]:focus {
    border-color: var(--border-bright) !important;
    outline: none !important;
    box-shadow: 0 0 0 1px var(--accent-amber-dim) !important;
}

textarea::placeholder, input::placeholder {
    color: var(--text-faint) !important;
    font-style: italic !important;
}

/* Output textareas — distinct from inputs */
.gradio-textbox[data-testid] textarea[readonly],
textarea[disabled] {
    background: var(--bg-card) !important;
    border-color: var(--border-dim) !important;
    color: var(--text-ivory) !important;
}

/* ── Buttons ──────────────────────────────────────────────────────────────── */
button.primary, .gr-button-primary {
    background: var(--accent-amber) !important;
    border: none !important;
    border-radius: 3px !important;
    color: var(--bg-void) !important;
    font-family: var(--font-body) !important;
    font-size: 0.78rem !important;
    font-weight: 600 !important;
    letter-spacing: 0.1em !important;
    text-transform: uppercase !important;
    padding: 14px 32px !important;
    transition: all 0.2s ease !important;
    cursor: pointer !important;
}

button.primary:hover, .gr-button-primary:hover {
    background: #d9a94d !important;
    transform: translateY(-1px) !important;
    box-shadow: 0 4px 16px rgba(200, 150, 60, 0.25) !important;
}

button.secondary, .gr-button-secondary {
    background: transparent !important;
    border: 1px solid var(--border-mid) !important;
    border-radius: 3px !important;
    color: var(--text-muted) !important;
    font-family: var(--font-body) !important;
    font-size: 0.75rem !important;
    font-weight: 400 !important;
    letter-spacing: 0.08em !important;
    text-transform: uppercase !important;
    padding: 12px 24px !important;
    transition: all 0.2s ease !important;
}

button.secondary:hover {
    border-color: var(--border-bright) !important;
    color: var(--text-ivory) !important;
}

/* ── Checkbox ─────────────────────────────────────────────────────────────── */
.gradio-checkbox label {
    color: var(--text-muted) !important;
    font-size: 0.82rem !important;
    font-weight: 300 !important;
    letter-spacing: 0.01em !important;
    text-transform: none !important;
    line-height: 1.6 !important;
}

input[type="checkbox"] {
    accent-color: var(--accent-amber) !important;
}

/* ── File upload ──────────────────────────────────────────────────────────── */
.gradio-file {
    background: var(--bg-input) !important;
    border: 1px dashed var(--border-mid) !important;
    border-radius: 4px !important;
    transition: border-color 0.2s !important;
}

.gradio-file:hover {
    border-color: var(--accent-amber-dim) !important;
}

/* ── Image upload ─────────────────────────────────────────────────────────── */
.gradio-image {
    background: var(--bg-input) !important;
    border: 1px solid var(--border-dim) !important;
    border-radius: 4px !important;
}

/* ── Accordion ────────────────────────────────────────────────────────────── */
.gradio-accordion > .label-wrap {
    background: var(--bg-card) !important;
    border: 1px solid var(--border-dim) !important;
    border-radius: 3px !important;
    padding: 10px 16px !important;
}

.gradio-accordion > .label-wrap span {
    font-family: var(--font-body) !important;
    font-size: 0.72rem !important;
    font-weight: 500 !important;
    letter-spacing: 0.1em !important;
    text-transform: uppercase !important;
    color: var(--text-muted) !important;
}

.gradio-accordion > .label-wrap:hover {
    border-color: var(--border-mid) !important;
}

/* ── Status / info boxes ──────────────────────────────────────────────────── */
.status-box {
    background: var(--bg-card) !important;
    border-left: 3px solid var(--accent-teal) !important;
    border-top: 1px solid var(--border-dim) !important;
    border-right: 1px solid var(--border-dim) !important;
    border-bottom: 1px solid var(--border-dim) !important;
    border-radius: 0 3px 3px 0 !important;
    padding: 12px 16px !important;
}

/* ── Dividers ─────────────────────────────────────────────────────────────── */
hr {
    border: none !important;
    border-top: 1px solid var(--border-dim) !important;
    margin: 28px 0 !important;
}

/* ── Scrollbars ───────────────────────────────────────────────────────────── */
::-webkit-scrollbar { width: 5px; height: 5px; }
::-webkit-scrollbar-track { background: var(--bg-deep); }
::-webkit-scrollbar-thumb { background: var(--border-mid); border-radius: 2px; }
::-webkit-scrollbar-thumb:hover { background: var(--border-bright); }

/* ── Examples row ─────────────────────────────────────────────────────────── */
.examples table {
    background: var(--bg-card) !important;
    border: 1px solid var(--border-dim) !important;
    border-radius: 3px !important;
}

.examples table td, .examples table th {
    color: var(--text-muted) !important;
    font-family: var(--font-data) !important;
    font-size: 0.78rem !important;
    border-color: var(--border-dim) !important;
    padding: 8px 12px !important;
}

.examples table tr:hover td {
    background: var(--bg-hover) !important;
    color: var(--text-ivory) !important;
}

/* ── Small helper text ────────────────────────────────────────────────────── */
small, .small-text {
    color: var(--text-faint) !important;
    font-size: 0.75rem !important;
    line-height: 1.5 !important;
}

/* ── Selection color ──────────────────────────────────────────────────────── */
::selection {
    background: var(--accent-amber-dim) !important;
    color: var(--text-ivory) !important;
}

/* ── Loading spinner ──────────────────────────────────────────────────────── */
.generating {
    border-color: var(--accent-amber) !important;
}

/* ── Blockquotes (used in description) ───────────────────────────────────── */
blockquote {
    border-left: 3px solid var(--accent-amber-dim) !important;
    background: var(--bg-card) !important;
    margin: 0 0 12px !important;
    padding: 10px 16px !important;
    border-radius: 0 3px 3px 0 !important;
}

blockquote p {
    color: var(--text-muted) !important;
    font-size: 0.82rem !important;
    margin: 0 !important;
}

/* ── Code / monospace in descriptions ────────────────────────────────────── */
code {
    font-family: var(--font-data) !important;
    background: var(--bg-input) !important;
    color: var(--accent-teal) !important;
    padding: 2px 6px !important;
    border-radius: 2px !important;
    font-size: 0.82em !important;
}
"""

# ==========================================
# GRADIO UI
# ==========================================
with gr.Blocks(
    theme=gr.themes.Base(
        primary_hue=gr.themes.colors.slate,
        secondary_hue=gr.themes.colors.slate,
        neutral_hue=gr.themes.colors.slate,
        font=[gr.themes.GoogleFont("IBM Plex Sans"), "system-ui", "sans-serif"],
        font_mono=[gr.themes.GoogleFont("IBM Plex Mono"), "monospace"],
    ),
    css=CUSTOM_CSS,
    title="VAJRAM — Clinical Decision Support"
) as demo:

    # ── Header ─────────────────────────────────────────────────────────────────
    with gr.Column(elem_classes=["header-block"]):
        gr.Markdown("""
# VAJRAM
**V**irtual **A**gent for **J**oint **R**isk **A**ssessment of **M**ultiple Myeloma  
MedGemma 1.5 · Mixture of Adapters · CPU-native · 
        """)

    with gr.Tabs():

        # ======================================================
        # TAB 1: CLINIC CONFIGURATION
        # ======================================================
        with gr.Tab("⚙  Configuration"):

            gr.Markdown("## Knowledge Base Initialization")
            gr.Markdown(
                "Upload your institution's **legally licensed** oncology guideline document. "
                "This document will serve as the evidence base for all treatment recommendations."
            )
            gr.Markdown(
                "> **Supported formats:** Text-based PDF only. Scanned documents are not supported.  \n"
                "> **Recommended sources:** ESMO, NCCN, ICMR, or your institution's local protocol."
            )

            kb_status_box = gr.Textbox(
                label="Knowledge Base Status",
                value=get_kb_status(),
                interactive=False,
                lines=1,
                elem_classes=["status-box"],
            )

            gr.Markdown("---")

            with gr.Row():
                with gr.Column(scale=2):
                    admin_pdf_upload = gr.File(
                        label="Oncology Guideline Document",
                        file_types=[".pdf"],
                        file_count="single",
                    )
                    legal_checkbox = gr.Checkbox(
                        label=(
                            "I confirm: (1) my institution holds a valid license for this document, "
                            "(2) I am authorised to upload it for AI use within this institution, and "
                            "(3) I understand that VAJRAM outputs are for clinical decision support only "
                            "and must be verified by a licensed physician before any clinical action."
                        ),
                        value=False,
                    )
                    init_btn = gr.Button(
                        "Initialize Knowledge Base",
                        variant="primary",
                        size="lg",
                    )

                with gr.Column(scale=1):
                    gr.Markdown("""
### Process Overview
1. Text extracted page-by-page
2. Content split into 1000-char chunks with overlap
3. Embedded via local sentence-transformer model
4. FAISS index saved to `./local_vector_store/`
5. All future consultations load instantly

### Estimated Processing Time
- 10-page document: 30 – 60 seconds
- 50-page document: 3 – 5 minutes
- 200-page document: 10 – 15 minutes

*Runs once. No cloud calls.*
                    """)

            admin_status = gr.Textbox(
                label="Initialization Log",
                interactive=False,
                lines=6,
            )

            init_btn.click(
                fn=initialize_knowledge_base,
                inputs=[admin_pdf_upload, legal_checkbox],
                outputs=admin_status,
            ).then(
                fn=get_kb_status,
                inputs=None,
                outputs=kb_status_box,
            )

        # ======================================================
        # TAB 2: PATIENT CONSULTATION
        # ======================================================
        with gr.Tab("🩺  Patient Consultation"):

            gr.Markdown("## Patient Data Entry")

            with gr.Row():
                with gr.Column(scale=2):
                    input_query = gr.Textbox(
                        lines=2,
                        label="Clinical Question",
                        placeholder="e.g. What is the recommended treatment for this transplant-eligible patient with renal impairment?",
                    )
                    with gr.Row():
                        input_labs = gr.Textbox(
                            lines=9,
                            label="Laboratory Values & Clinical Data",
                            placeholder="Creatinine, Calcium, Haemoglobin, Beta-2 Microglobulin, LDH, M-Spike, biopsy findings, staging...",
                        )
                        input_behaviour = gr.Textbox(
                            lines=9,
                            label="Behavioural Symptoms",
                            placeholder="e.g. Dizziness, Chest Pain, Mental Confusion, Fatigue, Peripheral Neuropathy...",
                        )

                with gr.Column(scale=1):
                    gr.Markdown("### Bone Marrow Biopsy")
                    input_wsi = gr.Image(
                        type="pil",
                        label="Whole Slide Image  (.bmp / .png / .tiff)",
                        height=220,
                    )
                    gr.Markdown(
                        "<small>If no image is uploaded, Module 3 will generate a "
                        "text-based WSI summary from clinical notes.</small>"
                    )

            with gr.Row():
                gr.Examples(
                    examples=[[EXAMPLE_QUERY, EXAMPLE_LABS, EXAMPLE_BEHAVIOUR, None]],
                    inputs=[input_query, input_labs, input_behaviour, input_wsi],
                    label="Load Example Patient Record",
                )
                submit_btn = gr.Button(
                    "Run Analysis",
                    variant="primary",
                    scale=0,
                    min_width=180,
                )

            # ── OUTPUT ────────────────────────────────────────────────────────
            gr.Markdown("---")
            gr.Markdown("## Analysis Output")

            with gr.Group():
                gr.Markdown("### Orchestrator — Module Pipeline & Final Recommendation")
                out_modules_selected = gr.Textbox(
                    label="Active Modules",
                    interactive=False,
                    lines=1,
                )
                out_final = gr.Textbox(
                    label="Final Treatment Recommendation  (citations in [SOURCE | PAGE] format)",
                    interactive=False,
                    lines=10,
                )

            gr.Markdown("---")
            gr.Markdown("### Specialist Module Outputs")

            with gr.Group(visible=True) as grp_m2:
                gr.Markdown("#### Module II — Risk Stratification")
                out_risk = gr.Textbox(
                    label="Risk Profile",
                    interactive=False, lines=3, visible=False,
                )

            with gr.Group(visible=True) as grp_m3:
                gr.Markdown("#### Module III — Bone Marrow Pathology")
                out_wsi_text = gr.Textbox(
                    label="WSI Analysis Summary",
                    interactive=False, lines=3, visible=False,
                )
                out_wsi_img = gr.Image(
                    label="AI-Annotated Whole Slide Image  (Red = Malignant · Green = Normal · Gray = Background · Yellow = Unknown)",
                    type="pil", height=420, interactive=False, visible=False,
                )
                out_malignancy_stat = gr.Textbox(
                    label="Patch Classification Statistics",
                    interactive=False, lines=2, visible=False,
                )

            with gr.Group(visible=True) as grp_m4:
                gr.Markdown("#### Module IV — Disease Progression")
                out_prog = gr.Textbox(
                    label="Progression Summary",
                    interactive=False, lines=3, visible=False,
                )

            with gr.Group(visible=True) as grp_m5:
                gr.Markdown("#### Module V — Guideline Retrieval (RAG)")
                out_m5_source = gr.Textbox(
                    label="Evidence Source",
                    interactive=False, lines=1, visible=False,
                )
                out_rag = gr.Textbox(
                    label="Retrieved Guideline Passages  (with source citations)",
                    interactive=False, lines=6, visible=False,
                )
                with gr.Accordion("View Raw Reference Chunks", open=False):
                    out_raw_chunks = gr.Textbox(
                        label="Raw knowledge base excerpts — verify AI citations against these passages",
                        interactive=False, lines=12, visible=False,
                    )

            submit_btn.click(
                fn=process_patient_data,
                inputs=[input_query, input_labs, input_behaviour, input_wsi],
                outputs=[
                    out_modules_selected, out_final,
                    out_risk,       grp_m2,
                    out_wsi_text,   grp_m3,
                    out_wsi_img,    out_malignancy_stat,
                    out_prog,       grp_m4,
                    out_rag,        grp_m5,
                    out_m5_source,  out_raw_chunks,
                ],
            )

if __name__ == "__main__":
    demo.queue().launch(share=False)