app.py

"""Sokrates — Gradio entrypoint.

Three-column clinical UI:
  1. Live transcript (uploaded audio, live microphone, or a sample dialogue)
  2. Clinical intake form filling in (green = filled, grey = missing)
  3. Suggested questions for the doctor (missing mandatory fields + follow-ups)

The flow is: audio -> ASR -> orchestration loop
(structured extraction + gap engine + question generation) -> this UI.
Sokrates does NOT make diagnoses; it structures data collection.
"""

from __future__ import annotations

import glob
import os

import gradio as gr

# Work around a gradio_client 1.3.0 bug: building the API schema crashes with
# "argument of type 'bool' is not iterable" when a JSON schema contains boolean
# values (e.g. additionalProperties: true). That crash makes Gradio's launch
# health check think localhost is unreachable and abort. Tolerate bool schemas.
import gradio_client.utils as _gcu

_orig_json_to_pytype = _gcu._json_schema_to_python_type
_orig_get_type = _gcu.get_type


def _safe_json_to_pytype(schema, defs=None):
    if isinstance(schema, bool):
        return "Any"
    return _orig_json_to_pytype(schema, defs)


def _safe_get_type(schema):
    if not isinstance(schema, dict):
        return "Any"
    return _orig_get_type(schema)


_gcu._json_schema_to_python_type = _safe_json_to_pytype
_gcu.get_type = _safe_get_type

from sokrates import asr
from sokrates.orchestrator import process_transcript
from sokrates.schema import (
    FIELD_META,
    IntakeForm,
    coerce_form,
    format_value,
    is_empty,
)

DATA_DIR = os.path.join(os.path.dirname(__file__), "data")

CUSTOM_CSS = """
:root {
    --sok-primary: #0e7490; --sok-primary-dark: #0b5566;
    --sok-accent: #2563eb; --sok-green: #16a34a; --sok-red: #dc2626;
    --sok-bg: #eef2f6; --sok-card: #ffffff; --sok-border: #e2e8f0;
    --sok-text: #1e293b; --sok-muted: #64748b;
}
.gradio-container { background: var(--sok-bg) !important; max-width: 1400px !important; }
footer { display: none !important; }

/* Header banner */
.sok-banner {
    background: linear-gradient(120deg, #0e7490 0%, #0891b2 50%, #2563eb 100%);
    border-radius: 16px; padding: 1.1rem 1.4rem; margin-bottom: 0.4rem;
    color: #fff; box-shadow: 0 6px 20px rgba(8,89,110,0.25);
    display: flex; align-items: center; gap: 1rem;
}
.sok-logo {
    width: 52px; height: 52px; border-radius: 14px; flex: 0 0 auto;
    background: rgba(255,255,255,0.18); display: flex; align-items: center;
    justify-content: center; font-size: 1.7rem;
}
.sok-banner h1 { margin: 0; font-size: 1.7rem; font-weight: 800; letter-spacing: 0.3px; }
.sok-banner .sub { opacity: 0.92; font-size: 0.92rem; margin-top: 0.15rem; }
.sok-badge {
    margin-left: auto; background: rgba(255,255,255,0.16);
    border: 1px solid rgba(255,255,255,0.3); border-radius: 999px;
    padding: 0.35rem 0.85rem; font-size: 0.78rem; font-weight: 600; white-space: nowrap;
}
.sok-steps { display: flex; gap: 0.5rem; margin: 0.5rem 0 0.2rem 0; }
.sok-step {
    flex: 1; background: var(--sok-card); border: 1px solid var(--sok-border);
    border-radius: 10px; padding: 0.45rem 0.7rem; font-size: 0.8rem; color: var(--sok-muted);
}
.sok-step b { color: var(--sok-primary); }

/* Cards */
.panel-card {
    background: var(--sok-card); border: 1px solid var(--sok-border); border-radius: 16px;
    padding: 1rem 1.1rem; box-shadow: 0 2px 10px rgba(16,40,70,0.06); min-height: 480px;
}
.panel-title {
    font-weight: 700; color: var(--sok-text); margin-bottom: 0.7rem; font-size: 1.02rem;
    display: flex; align-items: center; gap: 0.5rem;
}
.panel-title .chip {
    width: 30px; height: 30px; border-radius: 9px; display: inline-flex;
    align-items: center; justify-content: center; font-size: 1rem;
    background: #ecfeff; color: var(--sok-primary);
}

/* Progress */
.progress-wrap { margin: 0.1rem 0 0.9rem 0; }
.progress-bar { height: 9px; border-radius: 999px; background: #eef2f6; overflow: hidden; }
.progress-fill {
    height: 100%; border-radius: 999px;
    background: linear-gradient(90deg, #22c55e, #16a34a); transition: width 0.5s ease;
}
.progress-label {
    font-size: 0.78rem; color: var(--sok-muted); margin-top: 0.3rem;
    display: flex; justify-content: space-between;
}
.progress-label b { color: var(--sok-text); }

/* Form rows */
.form-row {
    display: flex; align-items: center; gap: 0.6rem;
    padding: 0.45rem 0.55rem; border-radius: 10px; margin-bottom: 0.32rem;
    border: 1px solid transparent; transition: transform 0.12s ease, box-shadow 0.12s ease;
}
.form-row:hover { transform: translateX(2px); }
.form-row .ic {
    width: 20px; height: 20px; border-radius: 50%; flex: 0 0 auto; color: #fff;
    display: inline-flex; align-items: center; justify-content: center; font-size: 0.7rem;
}
.form-row .label { flex: 1; color: var(--sok-text); font-size: 0.9rem; font-weight: 500; }
.form-row .value {
    color: var(--sok-text); font-size: 0.85rem; font-weight: 600;
    background: #f1f5f9; border-radius: 999px; padding: 0.15rem 0.6rem; max-width: 55%;
    overflow: hidden; text-overflow: ellipsis; white-space: nowrap;
}
.form-row.filled  { background: #f0fdf4; border-color: #dcfce7; }
.form-row.filled  .ic { background: var(--sok-green); }
.form-row.filled  .value { background: #dcfce7; color: #166534; }
.form-row.missing { background: #f8fafc; }
.form-row.missing .ic { background: #cbd5e1; }
.form-row.missing.req-missing { background: #fef2f2; border-color: #fee2e2; }
.form-row.missing.req-missing .ic { background: var(--sok-red); }
.form-row .req { color: var(--sok-red); font-size: 0.7rem; margin-left: 3px; font-weight: 700; }

/* Questions */
.q-item {
    background: #f8fbff; border: 1px solid #e3edfb; border-left: 4px solid var(--sok-accent);
    border-radius: 10px; padding: 0.65rem 0.8rem; margin-bottom: 0.55rem; color: #1e3a5f;
    font-size: 0.92rem; display: flex; align-items: flex-start; gap: 0.6rem;
    box-shadow: 0 1px 4px rgba(37,99,235,0.06); transition: transform 0.12s ease;
}
.q-item:hover { transform: translateY(-1px); box-shadow: 0 3px 10px rgba(37,99,235,0.12); }
.q-num {
    flex: 0 0 auto; min-width: 1.5rem; height: 1.5rem; line-height: 1.5rem; text-align: center;
    background: var(--sok-accent); color: #fff; border-radius: 50%; font-size: 0.78rem; font-weight: 700;
}
.gaps-note {
    background: #fffbeb; border: 1px solid #fde68a; color: #92400e; font-size: 0.8rem;
    border-radius: 9px; padding: 0.45rem 0.65rem; margin-bottom: 0.7rem;
}
.gaps-note b { color: #b45309; }
.empty-state { color: var(--sok-muted); font-size: 0.88rem; text-align: center; padding: 1.5rem 0.5rem; }
.disclaimer {
    color: #94a3b8; font-size: 0.76rem; margin-top: 0.7rem; font-style: italic;
    border-top: 1px dashed var(--sok-border); padding-top: 0.5rem;
}
/* Status pill — always clearly visible */
.sok-status {
    background: #ecfeff; border: 1px solid #a5f3fc; border-radius: 10px;
    padding: 0.55rem 0.85rem; margin-top: 0.55rem;
}
.sok-status p, .sok-status span, .sok-status {
    margin: 0 !important; color: #0b5566 !important;
    font-weight: 600 !important; font-size: 0.92rem !important;
}

/* Transcript box — force readable light styling even in OS dark mode */
.sok-transcript textarea {
    background: #ffffff !important; color: #1e293b !important;
    border: 1px solid var(--sok-border) !important; border-radius: 10px !important;
    font-size: 0.9rem !important; line-height: 1.4 !important;
}
.sok-transcript textarea::placeholder { color: #94a3b8 !important; }
"""


def render_form_html(form: IntakeForm) -> str:
    """Render the intake form: green = filled, grey = optional-missing,
    red-tinted = mandatory-missing. Includes a mandatory-completeness bar."""
    rows = []
    filled_mand = total_mand = 0
    for meta in FIELD_META:
        value = getattr(form, meta.name)
        filled = not is_empty(value)
        if meta.mandatory:
            total_mand += 1
            filled_mand += int(filled)
        css = "filled" if filled else "missing"
        if not filled and meta.mandatory:
            css += " req-missing"
        req = '<span class="req">*</span>' if meta.mandatory else ""
        icon = "✓" if filled else ("!" if meta.mandatory else "")
        value_html = (
            f'<span class="value">{format_value(value)}</span>' if filled else ""
        )
        rows.append(
            f'<div class="form-row {css}">'
            f'<span class="ic">{icon}</span>'
            f'<span class="label">{meta.label}{req}</span>'
            f"{value_html}"
            f"</div>"
        )
    pct = int(100 * filled_mand / total_mand) if total_mand else 0
    progress = (
        '<div class="progress-wrap">'
        f'<div class="progress-bar"><div class="progress-fill" '
        f'style="width:{pct}%"></div></div>'
        f'<div class="progress-label"><span>Mandatory fields</span>'
        f"<b>{filled_mand}/{total_mand} complete</b></div></div>"
    )
    return (
        '<div class="panel-card">'
        '<div class="panel-title"><span class="chip">🩺</span>'
        "Clinical intake form</div>"
        + progress
        + "".join(rows)
        + '<div class="disclaimer">* mandatory field. '
        "Sokrates structures data collection — it does not make diagnoses.</div>"
        "</div>"
    )


def render_questions_html(questions: list[str], missing: list[str] | None = None) -> str:
    """Render the suggested-questions panel with optional missing-field note."""
    missing = missing or []
    note = ""
    if missing:
        note = (
            '<div class="gaps-note">📋 <b>Still missing (mandatory):</b> '
            + ", ".join(missing)
            + "</div>"
        )
    if not questions:
        items = (
            '<div class="empty-state">💬 No suggested questions yet.<br>'
            "Analyze the transcript to generate them.</div>"
        )
    else:
        items = "".join(
            f'<div class="q-item"><span class="q-num">{i}</span>'
            f"<span>{q}</span></div>"
            for i, q in enumerate(questions, start=1)
        )
    return (
        '<div class="panel-card">'
        '<div class="panel-title"><span class="chip">💡</span>'
        "Suggested questions</div>"
        + note
        + items
        + "</div>"
    )


# --- Sample dialogues (synthetic, English) -------------------------------

def list_sample_dialogues() -> dict[str, str]:
    """Map a friendly label -> file path for each sample dialogue in data/."""
    samples: dict[str, str] = {}
    for path in sorted(glob.glob(os.path.join(DATA_DIR, "sample_dialogue_*.txt"))):
        label = (
            os.path.splitext(os.path.basename(path))[0]
            .replace("sample_dialogue_", "Sample dialogue ")
            .replace("_", " ")
        )
        samples[label] = path
    return samples


SAMPLE_DIALOGUES = list_sample_dialogues()


def load_sample_dialogue(label: str) -> str:
    """Return the text of a sample dialogue (used in place of ASR for the demo)."""
    path = SAMPLE_DIALOGUES.get(label)
    if not path or not os.path.exists(path):
        return ""
    with open(path, "r", encoding="utf-8") as fh:
        return fh.read().strip()


def transcribe_audio(audio_path: str | None) -> str:
    """Run ASR on an uploaded/recorded audio file and return the transcript."""
    if not audio_path:
        return ""
    try:
        return asr.transcribe_file(audio_path)
    except Exception as exc:  # surface ASR/setup errors in the UI, don't crash
        return f"[ASR error] {exc}"


# --- Orchestration (Step 5) ----------------------------------------------

def analyze_transcript(transcript: str, form: IntakeForm):
    """Run the full loop on the current transcript and repaint all panels.

    Implemented as a generator so the UI shows instant "Analyzing…" feedback
    before the (slower) model round-trip completes. Yields tuples of
    (form_state, form_html, questions_html, status_markdown).
    """
    form = coerce_form(form)  # Gradio State may hand back a dict
    if not transcript or not transcript.strip():
        yield (
            form,
            render_form_html(form),
            render_questions_html([]),
            "⚠️ Transcript is empty — load a sample or transcribe audio first.",
        )
        return
    # Immediate feedback while the model runs.
    yield (
        form,
        render_form_html(form),
        render_questions_html([]),
        "⏳ Analyzing transcript with the model… (first call may be slower)",
    )
    try:
        result = process_transcript(transcript, form)
    except Exception as exc:  # never surface a raw 'Error' badge to the user
        import traceback

        traceback.print_exc()  # full detail in the terminal
        yield (
            form,
            render_form_html(form),
            render_questions_html([]),
            f"⚠️ Error: {exc}",
        )
        return
    yield (
        result.form,
        render_form_html(result.form),
        render_questions_html(result.questions, result.missing_mandatory),
        "✅ " + result.status,
    )


def reset_session():
    """Clear the form, questions and transcript to start a fresh visit."""
    empty = IntakeForm()
    return (
        empty,
        render_form_html(empty),
        render_questions_html([]),
        "",
        "Session reset.",
    )


# --- Live microphone: record a segment, transcribe + analyze on stop ---
#
# Gradio 4.x's streaming-microphone (`gr.Audio(streaming=True)` + `.stream()`)
# is unreliable — chunks are often never delivered to the callback, so live
# streaming silently captures nothing. We instead record a clip and process it
# on `stop_recording`, reusing the same file-transcription path that the
# Upload/sample tab already relies on. The user records a segment, pauses, and
# Sokrates transcribes it and updates the form; recording again appends more.

def record_and_analyze(audio_path: str | None, transcript: str, form: IntakeForm):
    """Transcribe a recorded mic clip, append it, and run the full loop.

    Implemented as a generator so the UI shows instant feedback before the
    (slower) ASR + model round-trip completes. Yields tuples of
    (transcript, form_html, questions_html, status_markdown, form_state).
    """
    form = coerce_form(form)  # Gradio State may hand back a dict
    if not audio_path:
        yield (
            transcript,
            render_form_html(form),
            render_questions_html([]),
            "🎙️ No audio captured — click record, speak, then stop.",
            form,
        )
        return
    yield (
        transcript,
        render_form_html(form),
        render_questions_html([]),
        "⏳ Transcribing the recording…",
        form,
    )
    try:
        text = asr.transcribe_file(audio_path)
    except Exception as exc:
        import traceback

        traceback.print_exc()
        yield (
            transcript,
            render_form_html(form),
            render_questions_html([]),
            f"⚠️ ASR error: {exc}",
            form,
        )
        return
    if not text:
        yield (
            transcript,
            render_form_html(form),
            render_questions_html([]),
            "🎙️ No speech detected in that recording — try again.",
            form,
        )
        return
    new_transcript = (transcript + "\n" + text).strip() if transcript else text
    try:
        result = process_transcript(new_transcript, form)
    except Exception as exc:
        import traceback

        traceback.print_exc()
        yield (
            new_transcript,
            render_form_html(form),
            render_questions_html([]),
            f"⚠️ Error: {exc}",
            form,
        )
        return
    yield (
        new_transcript,
        render_form_html(result.form),
        render_questions_html(result.questions, result.missing_mandatory),
        "✅ " + result.status,
        result.form,
    )


DEFAULT_TRANSCRIPT = (
    "Load a sample dialogue, upload an audio file, or use the microphone, "
    "then click “Analyze transcript”."
)

# Force light mode regardless of the OS/browser theme, so the clinical look is
# consistent and every component stays readable (reloads once with __theme=light).
FORCE_LIGHT_JS = """
() => {
  const url = new URL(window.location.href);
  if (url.searchParams.get('__theme') !== 'light') {
    url.searchParams.set('__theme', 'light');
    window.location.replace(url.href);
  }
}
"""


def build_demo() -> gr.Blocks:
    theme = gr.themes.Soft(
        primary_hue=gr.themes.colors.cyan,
        secondary_hue=gr.themes.colors.blue,
        font=[gr.themes.GoogleFont("Inter"), "system-ui", "sans-serif"],
    )
    with gr.Blocks(
        css=CUSTOM_CSS, title="Sokrates", theme=theme, js=FORCE_LIGHT_JS
    ) as demo:
        gr.HTML(
            '<div class="sok-banner">'
            '<div class="sok-logo">⚕️</div>'
            "<div><h1>Sokrates</h1>"
            '<div class="sub">AI-assisted clinical intake — listens, structures '
            "the form, and suggests follow-up questions.</div></div>"
            '<div class="sok-badge">🔒 No diagnoses · data collection only</div>'
            "</div>"
            '<div class="sok-steps">'
            '<div class="sok-step"><b>1 · Listen</b> — transcribe the visit</div>'
            '<div class="sok-step"><b>2 · Structure</b> — fill the intake form</div>'
            '<div class="sok-step"><b>3 · Ask</b> — suggest the next questions</div>'
            "</div>"
        )

        form_state = gr.State(IntakeForm())

        with gr.Row(equal_height=False):
            # --- Column 1: transcript & input ---
            with gr.Column(scale=1):
                gr.HTML('<div class="panel-title">🎙️ Live transcript</div>')
                with gr.Tab("Upload / sample"):
                    audio_in = gr.Audio(
                        sources=["upload", "microphone"],
                        type="filepath",
                        label="Upload an audio file (or record)",
                    )
                    transcribe_btn = gr.Button(
                        "Transcribe audio", variant="secondary"
                    )
                    with gr.Row():
                        sample_dd = gr.Dropdown(
                            choices=list(SAMPLE_DIALOGUES.keys()),
                            label="Or load a sample dialogue",
                            scale=3,
                        )
                        sample_btn = gr.Button("Load sample", scale=1)
                with gr.Tab("Live microphone"):
                    mic_in = gr.Audio(
                        sources=["microphone"],
                        type="filepath",
                        label="Record the visit — transcribes when you stop",
                    )
                    gr.Markdown(
                        "Click record and speak; click stop when you pause. "
                        "Sokrates transcribes the clip and updates the form. "
                        "Record again to add more to the same visit."
                    )

                transcript_box = gr.Textbox(
                    value="",
                    placeholder=DEFAULT_TRANSCRIPT,
                    lines=12,
                    show_label=False,
                    interactive=True,
                    container=False,
                    elem_classes="sok-transcript",
                )
                with gr.Row():
                    analyze_btn = gr.Button(
                        "Analyze transcript → update", variant="primary", scale=3
                    )
                    reset_btn = gr.Button("Reset", scale=1)
                status_box = gr.Markdown(
                    "Ready — load a transcript and click **Analyze**.",
                    elem_classes="sok-status",
                )

            # --- Column 2: clinical form ---
            with gr.Column(scale=1):
                form_html = gr.HTML(render_form_html(IntakeForm()))

            # --- Column 3: suggested questions ---
            with gr.Column(scale=1):
                questions_html = gr.HTML(render_questions_html([]))

        # --- Wiring ---
        transcribe_btn.click(
            fn=transcribe_audio, inputs=audio_in, outputs=transcript_box
        )
        sample_btn.click(
            fn=load_sample_dialogue, inputs=sample_dd, outputs=transcript_box
        )
        analyze_btn.click(
            fn=analyze_transcript,
            inputs=[transcript_box, form_state],
            outputs=[form_state, form_html, questions_html, status_box],
        )
        reset_btn.click(
            fn=reset_session,
            inputs=None,
            outputs=[
                form_state,
                form_html,
                questions_html,
                transcript_box,
                status_box,
            ],
        )
        mic_in.stop_recording(
            fn=record_and_analyze,
            inputs=[mic_in, transcript_box, form_state],
            outputs=[
                transcript_box,
                form_html,
                questions_html,
                status_box,
                form_state,
            ],
        )

    return demo


if __name__ == "__main__":
    # Ensure localhost isn't routed through an HTTP proxy, which can otherwise
    # make Gradio's post-launch health check fail ("localhost not accessible").
    for _var in ("no_proxy", "NO_PROXY"):
        os.environ[_var] = "localhost,127.0.0.1,0.0.0.0," + os.environ.get(_var, "")
    build_demo().launch(server_name="0.0.0.0", server_port=7860)