ui/layout.py

"""Gradio cockpit layout for the car diagnostic agent."""

from __future__ import annotations

from collections.abc import Callable, Iterator

import gradio as gr

from fault_simulation import describe_fault, list_fault_choices
from obd_connection import VITALS_POLL_INTERVAL
from obd_vitals import format_vitals_hud_cached, refresh_vitals_hud
from trace_store import TRACE_ENABLED
from ui.header import render_cockpit_header
from ui.theme import load_cockpit_css


def build_ui(
    *,
    default_thinking_enabled: bool,
    active_fault_status_fn: Callable[[], str],
    on_fault_select_fn: Callable[[str | None], str],
    on_apply_fault_fn: Callable[..., Iterator],
    on_clear_fault_fn: Callable[[], tuple],
    chat_submit_fn: Callable[..., Iterator],
    download_traces_fn: Callable[[], str | None],
) -> gr.Blocks:
    """Build the futuristic cockpit Gradio interface."""
    with gr.Blocks(title="Car Diagnostic Agent", fill_height=True) as demo:
        gr.HTML(f"<style>{load_cockpit_css()}</style>", visible=False)
        gr.HTML(render_cockpit_header())

        with gr.Row(elem_classes=["cockpit-main-row"]):
            with gr.Column(scale=5, elem_id="command-center", elem_classes=["hud-panel"]):
                enable_thinking = gr.Checkbox(
                    label="Enable thinking",
                    value=default_thinking_enabled,
                    info="Reasoning is enabled server-side but not shown in chat.",
                )
                chatbot = gr.Chatbot(
                    label="Diagnostic chat",
                    height=460,
                    elem_classes=["cockpit-chat"],
                )
                with gr.Row():
                    msg_input = gr.Textbox(
                        label="Message",
                        placeholder="Describe symptoms…",
                        scale=4,
                        show_label=False,
                    )
                    send_btn = gr.Button(
                        "Send",
                        variant="primary",
                        scale=1,
                        elem_id="send-btn",
                    )
                gr.Examples(
                    examples=[
                        "My check engine light is on and the car idles rough.",
                        "I hear a grinding noise when I brake.",
                        "The car won't start but the battery seems fine.",
                    ],
                    inputs=msg_input,
                )

            with gr.Column(scale=3, elem_classes=["instrument-bay"]):
                vitals_panel = gr.HTML(
                    value=format_vitals_hud_cached(),
                    elem_id="vitals-hud",
                    elem_classes=["hud-panel"],
                )
                refresh_btn = gr.Button(
                    "Refresh vitals",
                    variant="secondary",
                    elem_id="refresh-vitals-btn",
                )
                fault_status = gr.Markdown(
                    value=active_fault_status_fn(),
                    elem_id="fault-status-strip",
                    elem_classes=["hud-panel"],
                )
                trace_panel = gr.Markdown(
                    value="*Tool calls and results will stream here…*",
                    label="Telemetry log",
                    elem_id="telemetry-log",
                    elem_classes=["hud-panel"],
                )
                if TRACE_ENABLED:
                    download_traces = gr.DownloadButton("Download agent traces (JSONL)")
                    download_traces.click(download_traces_fn, outputs=download_traces)

                with gr.Accordion("Simulate fault", open=False, elem_id="fault-bay"):
                    fault_select = gr.Dropdown(
                        choices=list_fault_choices(),
                        value="none",
                        label="Fault profile",
                    )
                    fault_info = gr.Markdown(value=describe_fault("none"))
                    with gr.Row():
                        apply_fault_btn = gr.Button("Apply fault", variant="primary")
                        clear_fault_btn = gr.Button("Clear / reset", variant="secondary")

        vitals_timer = gr.Timer(VITALS_POLL_INTERVAL, active=True)
        vitals_timer.tick(format_vitals_hud_cached, outputs=vitals_panel)
        vitals_timer.tick(active_fault_status_fn, outputs=fault_status)
        refresh_btn.click(refresh_vitals_hud, outputs=vitals_panel)
        demo.load(refresh_vitals_hud, outputs=vitals_panel)
        demo.load(active_fault_status_fn, outputs=fault_status)
        fault_select.change(on_fault_select_fn, inputs=fault_select, outputs=fault_info)
        apply_fault_btn.click(
            on_apply_fault_fn,
            inputs=[fault_select, chatbot, enable_thinking],
            outputs=[fault_status, vitals_panel, chatbot, trace_panel],
        ).then(on_fault_select_fn, inputs=fault_select, outputs=fault_info)
        clear_fault_btn.click(
            on_clear_fault_fn,
            outputs=[fault_status, vitals_panel, fault_select],
        ).then(lambda: describe_fault("none"), outputs=fault_info)

        chat_inputs = [msg_input, chatbot, enable_thinking]
        chat_outputs = [chatbot, trace_panel]

        send_btn.click(chat_submit_fn, chat_inputs, chat_outputs).then(
            lambda: "", outputs=msg_input
        )
        msg_input.submit(chat_submit_fn, chat_inputs, chat_outputs).then(
            lambda: "", outputs=msg_input
        )

    return demo