server/ui.py

import gradio as gr
import json

CUSTOM_CSS = """
.metric-card {
    background: rgba(255, 255, 255, 0.85);
    backdrop-filter: blur(10px);
    border: 1px solid rgba(226, 232, 240, 0.8);
    border-radius: 12px;
    padding: 1.5rem;
    margin-bottom: 1rem;
    box-shadow: 0 4px 6px -1px rgba(0, 0, 0, 0.05);
}
.dark .metric-card {
    background: rgba(30, 41, 59, 0.85);
    border-color: rgba(51, 65, 85, 0.8);
}
.status-pill {
    display: inline-block;
    padding: 0.25rem 0.75rem;
    border-radius: 9999px;
    font-size: 0.875rem;
    font-weight: 600;
}
.status-pill.success { background: #dcfce7; color: #166534; }
.status-pill.error { background: #fee2e2; color: #991b1b; }
.status-pill.info { background: #dbeafe; color: #1e3a8a; }
.dark .status-pill.success { background: #14532d; color: #bbf7d0; }
.dark .status-pill.error { background: #7f1d1d; color: #fecaca; }
.dark .status-pill.info { background: #1e3a8a; color: #bfdbfe; }
"""


def build_triage_ui(web_manager, action_fields, metadata, is_chat_env, title, quick_start_md):
    with gr.Blocks(css=CUSTOM_CSS) as blocks:
        gr.Markdown(f"# 🏥 {title} — Clinical Dashboard")
        gr.Markdown(
            "Patients arrive one at a time. **Classify** each patient and **reorder the priority queue**. "
            "Use ESCALATE for patients with missing critical data."
        )

        current_state = gr.State({})

        with gr.Row():
            task_dropdown = gr.Dropdown(
                choices=["basic-triage", "incomplete-records-triage", "mass-casualty-triage"],
                label="Scenario",
                value="basic-triage",
                interactive=True,
            )
            reset_btn = gr.Button("🔄 Initialize Scenario", variant="primary")

        with gr.Row():
            with gr.Column(scale=2):
                with gr.Group(elem_classes="metric-card"):
                    gr.Markdown("### 🫀 Incoming Patient")
                    patient_md = gr.Markdown("*Press Initialize to start*")

                with gr.Group(elem_classes="metric-card"):
                    gr.Markdown("### 📊 Current Priority Queue")
                    queue_md = gr.Markdown("*Empty*")
                    feedback_md = gr.Markdown("")

            with gr.Column(scale=1):
                with gr.Group(elem_classes="metric-card"):
                    gr.Markdown("### ⚡ Triage Decision")

                    classification = gr.Dropdown(
                        choices=["immediate", "urgent", "less_urgent", "non_urgent", "escalate"],
                        label="Classification",
                        value="urgent",
                    )

                    queue_input = gr.Textbox(
                        label="Reordered Queue (comma-separated patient IDs)",
                        placeholder="P001, P003, P002",
                    )

                    submit_btn = gr.Button("📋 Submit", variant="primary")

                with gr.Group(elem_classes="metric-card"):
                    gr.Markdown("### 📜 Event Log")
                    log_md = gr.Markdown("No actions yet.")

        def format_patient(obs):
            pt = obs.get("incoming_patient", {})
            if not pt:
                return "**No more patients.**"
            vitals = pt.get("vitals")
            v_str = ", ".join(f"**{k}**: {v}" for k, v in vitals.items()) if vitals else "⚠️ **MISSING**"
            step = obs.get("step_number", 0) + 1
            total = obs.get("total_expected_patients", 0)
            return (
                f"**Step {step}/{total}** — ID: `{pt.get('patient_id')}` | Age: {pt.get('age')}\n\n"
                f"**Complaint:** {pt.get('chief_complaint')}\n\n"
                f"**Symptoms:** {', '.join(pt.get('symptoms', []))}\n\n"
                f"**Vitals:** {v_str}\n\n"
                f"**History:** {pt.get('history', 'N/A')} | **Meds:** {pt.get('medications', 'N/A')}"
            )

        def format_queue(queue):
            if not queue:
                return "📭 *Empty queue*"
            return " → ".join(f"`{pid}`" for pid in queue)

        async def ui_reset(task_name):
            try:
                res = await web_manager.reset_environment({"task_name": task_name})
                obs = res.get("observation", {})
                q = obs.get("current_queue", [])
                fb = obs.get("previous_feedback", "")
                return [res, format_patient(obs), format_queue(q), fb, "No actions yet."]
            except Exception as e:
                return [{}, f"Error: {e}", "*Empty*", "", ""]

        async def ui_step(cls, q_text, st):
            obs = st.get("observation", {})
            pt = obs.get("incoming_patient", {})
            pid = pt.get("patient_id", "") if pt else ""

            queue_ids = [x.strip() for x in q_text.split(",") if x.strip()] if q_text.strip() else []

            payload = {"classification": cls, "reordered_queue": queue_ids}

            try:
                res = await web_manager.step_environment(payload)
                obs2 = res.get("observation", {})
                q = obs2.get("current_queue", [])
                fb = obs2.get("previous_feedback", "")

                logs = getattr(web_manager.episode_state, "action_logs", [])
                log_str = ""
                for log in reversed(logs[-8:]):
                    r = log.reward
                    d = log.done
                    a_cls = log.action.get("classification", "?")
                    st_str = "🏁 Done" if d else f"R={r:.2f}"
                    log_str += f"- **{a_cls}** → `{pid}` | {st_str}\n"

                return [res, format_patient(obs2), format_queue(q), fb, log_str or "No actions yet."]
            except Exception as e:
                return [st, f"Error: {e}", "", "", ""]

        outputs = [current_state, patient_md, queue_md, feedback_md, log_md]
        reset_btn.click(fn=ui_reset, inputs=[task_dropdown], outputs=outputs)
        submit_btn.click(fn=ui_step, inputs=[classification, queue_input, current_state], outputs=outputs)

    return blocks