import os
import re
import time
import gradio as gr
from PIL import Image

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

# Import the raw inference function so we can stream the orchestrator
from gguf_engine import (
    _load_text_model,
    _format_gemma_prompt,
    exclude_thinking_component,
    LLM_LORA_PATHS,
)
import psutil
import os

def check_memory():
    process = psutil.Process(os.getpid())
    # RAM usage in GB
    return process.memory_info().rss / 1024**3 




# ==========================================
# STREAMING HELPER
# Calls llama-cpp-python with stream=True so
# the orchestrator output appears token-by-token.
# ==========================================
def _stream_generate(prompt: str, max_tokens: int = 800):
    """
    Generator that yields incremental text from the base model (no LoRA).
    Used exclusively for the orchestrator's final answer.
    """
    print(f"[BEFORE MODEL LOAD] RAM Usage: {check_memory():.2f} GB")
    model = _load_text_model("default")
    print(f"[AFTER MODEL LOAD] RAM Usage: {check_memory():.2f} GB")
    formatted = _format_gemma_prompt(prompt)
    accumulated = ""
    for chunk in model(
        formatted,
        max_tokens  = max_tokens,
        stop        = ["<end_of_turn>", "<eos>"],
        echo        = False,
        temperature = 0.0,
        top_p       = 1.0,
        stream      = True,
    ):
        print(f"[STREAM] RAM Usage: {check_memory():.2f} GB")
        token = chunk["choices"][0]["text"]
        accumulated += token
        yield exclude_thinking_component(accumulated)


# ==========================================
# BACKEND: TAB 1 — ADMIN KNOWLEDGE BASE
# ==========================================
def initialize_knowledge_base(pdf_file, legal_consent):
    if not legal_consent:
        yield "⛔ Cannot proceed: legal consent is required."
        return
    if pdf_file is None:
        yield "⛔ No PDF uploaded."
        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⏳ Building vector index — please wait..."
    success, message = save_retriever_from_pdf(pdf_path)
    if success:
        yield f"{message}\n\n✅ Knowledge base ready."
    else:
        yield f"{message}\n\n⚠️ Falling back to built-in NCCN/ESMO excerpts."


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


# ==========================================
# BACKEND: TAB 2 — PATIENT CONSULTATION
# Streaming generator — yields partial UI
# updates as each LangGraph node completes,
# then streams the final answer token-by-token.
# ==========================================
def process_patient_data(
    user_query, lab_values, behaviour_changes, wsi_image,
    progress=gr.Progress(track_tqdm=True)
):
    print(f"[START] RAM Usage: {check_memory():.2f} GB")
    # ── Assemble clinical text ─────────────────────────────────────────────
    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)

    # ── Helper: build the full yield-tuple from current partial state ──────
    def _build_yield(
        modules_str="", recommendation="⏳ Analysing...",
        risk="", show_m2=False,
        wsi_text="", show_m3=False,
        img=None, mal_stat=gr.update(visible=False),
        prog="", show_m4=False,
        rag="", show_m5=False,
        m5_src=gr.update(visible=False),
        raw_chunks=gr.update(visible=False),
    ):
        return (
            modules_str, recommendation,
            gr.update(value=risk,     visible=show_m2), gr.update(visible=show_m2),
            gr.update(value=wsi_text, visible=show_m3), gr.update(visible=show_m3),
            gr.update(value=img,      visible=img is not None), mal_stat,
            gr.update(value=prog,     visible=show_m4), gr.update(visible=show_m4),
            gr.update(value=rag,      visible=show_m5), gr.update(visible=show_m5),
            m5_src, raw_chunks,
        )

    # ── Initial "running" state ────────────────────────────────────────────
    progress(0, desc="Initialising pipeline…")
    yield _build_yield(modules_str="Planning…", recommendation="⏳ Running diagnostic pipeline…")

    # ── Run the full agent (blocking — all modules run here) ───────────────
    progress(0.1, desc="Routing clinical query to modules…")

    initial_state = {
        "patient_id":           "VAJRAM_UI",
        "user_query":           user_query.strip() or "Give me a full clinical workup.",
        "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": "",
    }

    # Run everything up to (but not including) the orchestrator synthesis
    # by temporarily patching the orchestrator to return early.
    # Simpler: just run the full agent and grab intermediate results.
    progress(0.2, desc="Module 2 · Risk Stratification LoRA…")
    final_state = full_agent.invoke(initial_state)
    print(f"[AFTER AGENT] RAM Usage: {check_memory():.2f} GB")

    # ── Unpack results ─────────────────────────────────────────────────────
    ran_m2 = bool(final_state.get("module2_risk_score",   "").strip())
    ran_m3 = bool(final_state.get("module3_wsi_analysis", "").strip())
    ran_m4 = bool(final_state.get("module4_progression",  "").strip())
    ran_m5 = bool(final_state.get("module5_guidelines",   "").strip())

    selected = []
    if ran_m2: selected.append("Module 2 · Risk")
    if ran_m3: selected.append("Module 3 · WSI")
    if ran_m4: selected.append("Module 4 · Progression")
    if ran_m5: selected.append("Module 5 · RAG")
    modules_str = "  ›  ".join(selected) if selected else "None"

    # ── Annotated WSI image ────────────────────────────────────────────────
    annotated_img = None
    mal_stat      = gr.update(visible=False)
    if ran_m3 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)
                mal_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:
                mal_stat = gr.update(value=wsi_summary, visible=True)
        except Exception:
            annotated_img = None

    m5_src_update = gr.update(
        value=f"Source: {final_state.get('module5_source','')}" if ran_m5 else "",
        visible=ran_m5
    )
    raw_chunks_update = gr.update(
        value=final_state.get("module5_raw_chunks", ""),
        visible=ran_m5
    )

    # ── Yield intermediate state before streaming the final answer ─────────
    progress(0.7, desc="Modules complete — synthesising recommendation…")
    yield _build_yield(
        modules_str  = modules_str,
        recommendation = "⏳ Synthesising final recommendation…",
        risk         = final_state.get("module2_risk_score",   ""),
        show_m2      = ran_m2,
        wsi_text     = final_state.get("module3_wsi_analysis", ""),
        show_m3      = ran_m3,
        img          = annotated_img,
        mal_stat     = mal_stat,
        prog         = final_state.get("module4_progression",  ""),
        show_m4      = ran_m4,
        rag          = final_state.get("module5_guidelines",   ""),
        show_m5      = ran_m5,
        m5_src       = m5_src_update,
        raw_chunks   = raw_chunks_update,
    )

    # ── Stream the final recommendation token-by-token ─────────────────────
    # Build the same prompt the orchestrator would use, then stream it.
    profile_parts = []
    if final_state.get("module2_risk_score"):
        profile_parts.append(f"- Risk Score: {final_state['module2_risk_score']}")
    if final_state.get("module3_wsi_analysis"):
        profile_parts.append(f"- Bone Marrow WSI: {final_state['module3_wsi_analysis']}")
    if final_state.get("module4_progression"):
        profile_parts.append(f"- Progression: {final_state['module4_progression']}")
    profile_block = "\n".join(profile_parts) if profile_parts else "No module findings."

    if not ran_m5:
        stream_prompt = (
            f"You are a hematology AI assistant.\n"
            f"Answer the clinician's question concisely using only the findings below.\n"
            f"Do not speculate. Do not mention treatment guidelines.\n\n"
            f"[QUESTION]\n{user_query}\n\n"
            f"[FINDINGS]\n{profile_block}\n\n"
            f"Answer in 2-4 sentences:"
        )
        max_tok = 250
    else:
        guidelines_block = final_state.get("module5_guidelines", "")
        stream_prompt = (
            f"You are the Master Hematology Orchestrator.\n"
            f"Answer the clinician's question using ONLY the data below.\n"
            f"Do not speculate. Do not show your reasoning.\n\n"
            f"CITATION RULE: After every treatment or guideline recommendation,\n"
            f"append the exact [SOURCE: ... | PAGE: ...] tag from the guidelines.\n"
            f"If no supporting guideline exists for a statement, write (no guideline available).\n\n"
            f"[QUESTION]\n{user_query}\n\n"
            f"[PATIENT DATA]\n{raw_clinical_text}\n\n"
            f"[MODULE FINDINGS]\n{profile_block}\n\n"
            f"[RETRIEVED GUIDELINES — cite exactly]\n{guidelines_block}\n\n"
            f"Final Answer (with citations):"
        )
        max_tok = 800

    progress(0.85, desc="Streaming final recommendation…")

    for streamed_text in _stream_generate(stream_prompt, max_tokens=max_tok):
        yield _build_yield(
            modules_str  = modules_str,
            recommendation = streamed_text,
            risk         = final_state.get("module2_risk_score",   ""),
            show_m2      = ran_m2,
            wsi_text     = final_state.get("module3_wsi_analysis", ""),
            show_m3      = ran_m3,
            img          = annotated_img,
            mal_stat     = mal_stat,
            prog         = final_state.get("module4_progression",  ""),
            show_m4      = ran_m4,
            rag          = final_state.get("module5_guidelines",   ""),
            show_m5      = ran_m5,
            m5_src       = m5_src_update,
            raw_chunks   = raw_chunks_update,
        )

    progress(1.0, desc="Complete")


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
# ==========================================
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 {
    --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;
    --font-display:     'Playfair Display', Georgia, serif;
    --font-data:        'IBM Plex Mono', 'Courier New', monospace;
    --font-body:        'IBM Plex Sans', system-ui, sans-serif;
}

*, *::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;
}

/* ── Hardware banner ──────────────────────────────────────────────────────── */
.hw-banner {
    background: linear-gradient(135deg, #0d1a2e 0%, #0a1520 100%) !important;
    border: 1px solid var(--accent-amber-dim) !important;
    border-left: 3px solid var(--accent-amber) !important;
    border-radius: 0 4px 4px 0 !important;
    padding: 12px 18px !important;
    margin-bottom: 20px !important;
}

.hw-banner p {
    color: var(--text-muted) !important;
    font-size: 0.82rem !important;
    line-height: 1.5 !important;
    margin: 0 !important;
}

.hw-banner strong {
    color: var(--accent-amber) !important;
    font-weight: 500 !important;
}

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

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

/* ── Progress bar ─────────────────────────────────────────────────────────── */
.progress-bar-wrap {
    background: var(--bg-card) !important;
    border: 1px solid var(--border-dim) !important;
    border-radius: 3px !important;
    overflow: hidden !important;
}

.progress-bar {
    background: var(--accent-amber) !important;
    height: 3px !important;
    transition: width 0.4s ease !important;
}

/* ── 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;
}

/* ── Tabs ─────────────────────────────────────────────────────────────────── */
.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 ───────────────────────────────────────────────────────────────── */
.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;
}

/* ── Inputs ───────────────────────────────────────────────────────────────── */
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; }

/* ── 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;
    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 / Image 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; }
.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;
}

/* ── Status box ───────────────────────────────────────────────────────────── */
.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;
}

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

::-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 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 { color: var(--text-faint) !important; font-size: 0.75rem !important; }
::selection { background: var(--accent-amber-dim) !important; color: var(--text-ivory) !important; }

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 {
    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 · Air-gapped deployment
        """)

    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 becomes 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 institutional AI use, and "
                            "(3) 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. Split into 1000-char chunks with overlap
3. Embedded via local sentence-transformer
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"):

            # ── Hardware disclaimer banner ────────────────────────────────
            with gr.Column(elem_classes=["hw-banner"]):
                gr.Markdown(
                    "**Live demo running on a throttled 2-core CPU** to demonstrate edge-deployment capability. "
                    "Inference takes ~2 minutes on this hardware. "
                    "Enterprise GPU deployments process in **< 5 seconds**. "
                    "The progress bar and streaming output below keep you informed throughout."
                )

            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 generates 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 Recommendation  (streaming · 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 WSI  (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",
                        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,
                ],
            )

demo.queue().launch(
    server_name="0.0.0.0",
    server_port=7860,
    share=False,
    ssr_mode=False,
)