server/ui.py

# Copyright (c) Meta Platforms, Inc. and affiliates.
# All rights reserved.
#
# This source code is licensed under the BSD-style license found in the
# LICENSE file in the root directory of this source tree.

"""Interactive Gradio UI for WorkflowArena."""

from __future__ import annotations

import random
from types import SimpleNamespace
from typing import Any

import gradio as gr
import plotly.graph_objects as go

from workflow_arena.models import DifficultyPreset, TaskStatus, WorkflowActionType, WorkflowArenaAction
from workflow_arena.presets import get_preset_config
from workflow_arena.server.workflow_arena_environment import WorkflowArenaEnvironment


Session = dict[str, Any]

DETAIL_HEADERS = [
    "Task",
    "Priority",
    "Duration",
    "Deadline",
    "Criticality",
    "Slack",
    "Deps",
    "Downstream",
    "Attempts",
    "Start",
    "End",
]

PRESET_BRIEFS = {
    DifficultyPreset.EASY.value: {
        "label": "Warm-up Flow",
        "summary": "Small DAG, softer deadlines, and fewer traps. Good for learning how dispatch and wait interact.",
        "focus": "Keep workers busy, avoid empty waits, and build intuition for parallel batches.",
        "mechanics": "No hard time budget and no failure events.",
    },
    DifficultyPreset.MEDIUM.value: {
        "label": "Balanced Pressure",
        "summary": "Tighter dependencies and more timing pressure. Scheduling mistakes start to compound.",
        "focus": "Balance urgency, downstream unlocks, and worker utilization.",
        "mechanics": "Adds a fixed time budget and terminal penalty for unfinished work.",
    },
    DifficultyPreset.HARD.value: {
        "label": "Critical Path Sprint",
        "summary": "Dense DAGs, tighter deadlines, and much less room for idle capacity.",
        "focus": "Protect the critical path and use every free slot intentionally.",
        "mechanics": "Adds a tighter time budget plus seeded worker outages and task retry failures.",
    },
}

CSS = """
.gradio-container {
  background:
    radial-gradient(circle at top left, rgba(216, 116, 76, 0.14), transparent 28%),
    radial-gradient(circle at top right, rgba(201, 157, 92, 0.10), transparent 24%),
    linear-gradient(180deg, #fbf4ea 0%, #f4e7d6 100%);
  color: #2d241c;
  font-family: "IBM Plex Sans", "Avenir Next", "Segoe UI", sans-serif;
}
.wa-shell {max-width: 1380px; margin: 0 auto; padding: 10px 10px 30px;}
.wa-title {margin-bottom: 18px;}
.wa-title h1 {margin: 0; font-size: 2.6rem; line-height: 1; letter-spacing: -0.04em; color: #3e2618;}
.wa-title p {margin: 10px 0 0; max-width: 920px; font-size: 1rem; color: #745f50;}
.wa-hero {display: grid; grid-template-columns: 1.15fr 0.85fr; gap: 18px; margin-bottom: 18px;}
.wa-card {
  background: rgba(255, 252, 247, 0.92);
  border: 1px solid rgba(139, 110, 84, 0.12);
  border-radius: 24px;
  box-shadow: 0 18px 60px rgba(114, 84, 51, 0.10);
  backdrop-filter: blur(16px);
}
.wa-control-card {padding: 20px;}
.wa-control-card h3,
.wa-panel h3,
.wa-playbook h3 {
  margin: 0 0 8px;
  font-size: 0.8rem;
  letter-spacing: 0.12em;
  text-transform: uppercase;
  color: #9f5b33;
}
.wa-control-card p,
.wa-panel p,
.wa-playbook p {margin: 0; color: #715e50; line-height: 1.5;}
.wa-control-grid {display: grid; grid-template-columns: 1fr 1fr 1fr; gap: 12px; align-items: end; margin-top: 14px;}
.wa-control-buttons {display: flex; gap: 10px; align-items: center; margin-top: 12px;}
.wa-inline-buttons {display: flex; gap: 10px; align-items: center; flex-wrap: wrap;}
.wa-compact-accordion {margin-top: 14px;}
.wa-compact-accordion .label-wrap span {font-size: 0.85rem;}
.wa-problem-box {
  padding: 4px 2px 2px;
  border-radius: 18px;
}
.wa-problem-box strong {color: #8d4f2d;}
.wa-problem-box p {margin: 0 0 8px; color: #6d594c;}
.wa-problem-box p:last-child {margin-bottom: 0;}
.wa-preset-card {padding: 20px; min-height: 100%;}
.wa-preset-card .eyebrow {font-size: 0.75rem; letter-spacing: 0.14em; text-transform: uppercase; color: #b16d3d;}
.wa-preset-card .name {margin-top: 8px; font-size: 1.7rem; font-weight: 700; letter-spacing: -0.03em; color: #3e2618;}
.wa-preset-card .summary {margin-top: 10px; color: #6e594a; line-height: 1.55;}
.wa-preset-meta {margin-top: 12px; color: #7b6554; line-height: 1.5;}
.wa-preset-focus {
  margin-top: 14px;
  padding: 12px 14px;
  border-radius: 18px;
  background: linear-gradient(135deg, rgba(222, 143, 93, 0.12), rgba(242, 210, 171, 0.28));
  border: 1px solid rgba(174, 117, 72, 0.14);
}
.wa-topbar {display: grid; grid-template-columns: 1.2fr 1fr 1fr 1fr 1fr 1fr; gap: 12px; margin: 0 0 16px;}
.wa-stat {
  background: linear-gradient(180deg, #fff8ef 0%, #f9efe2 100%);
  border: 1px solid rgba(168, 130, 95, 0.16);
  border-radius: 20px;
  padding: 16px 18px;
}
.wa-stat .label {font-size: 0.7rem; letter-spacing: 0.12em; text-transform: uppercase; color: #a56c43;}
.wa-stat .value {margin-top: 6px; font-size: 1.65rem; font-weight: 700; color: #3e2618;}
.wa-stat .sub {margin-top: 4px; font-size: 0.82rem; color: #7a6657;}
.wa-banner {
  border-radius: 24px;
  padding: 18px 20px;
  border: 1px solid rgba(177, 107, 70, 0.18);
  background: linear-gradient(135deg, rgba(245, 186, 144, 0.35), rgba(255, 251, 245, 0.96));
  color: #35271f;
  margin-bottom: 16px;
}
.wa-banner.invalid {
  background: linear-gradient(135deg, rgba(247, 202, 196, 0.92), rgba(255, 246, 244, 0.98));
  border-color: rgba(180, 84, 69, 0.22);
}
.wa-banner.done {
  background: linear-gradient(135deg, rgba(219, 229, 195, 0.9), rgba(255, 251, 244, 0.98));
  border-color: rgba(112, 139, 90, 0.22);
}
.wa-banner-top {display: flex; justify-content: space-between; gap: 16px; align-items: flex-start;}
.wa-banner .status {
  display: inline-block;
  padding: 5px 10px;
  border-radius: 999px;
  font-size: 0.72rem;
  font-weight: 700;
  letter-spacing: 0.08em;
  text-transform: uppercase;
  background: rgba(164, 97, 59, 0.12);
}
.wa-banner .meta {margin-top: 8px; font-size: 0.9rem; color: #7f6654;}
.wa-banner .note {margin-top: 12px; font-size: 1rem; line-height: 1.5; color: #49372c;}
.wa-banner-grid {display: grid; grid-template-columns: repeat(4, minmax(0, 1fr)); gap: 10px; margin-top: 14px;}
.wa-banner-metric {
  padding: 12px 14px;
  border-radius: 18px;
  background: rgba(255, 253, 248, 0.76);
  border: 1px solid rgba(178, 132, 96, 0.12);
}
.wa-banner-metric span {display: block; font-size: 0.72rem; letter-spacing: 0.08em; text-transform: uppercase; color: #a16b45;}
.wa-banner-metric strong {display: block; margin-top: 6px; font-size: 1.1rem; color: #3d281c;}
.wa-progress {height: 10px; margin-top: 14px; border-radius: 999px; overflow: hidden; background: rgba(120, 83, 54, 0.10);}
.wa-progress-fill {height: 100%; background: linear-gradient(90deg, #e49157 0%, #f1c27a 100%);}
.wa-main {display: grid; grid-template-columns: 1.18fr 0.82fr; gap: 18px; align-items: start;}
.wa-left-stack,
.wa-right-stack {display: grid; gap: 18px;}
.wa-panel {padding: 18px 18px 16px;}
.wa-playbook {padding: 18px;}
.wa-playbook-header {display: flex; justify-content: space-between; gap: 12px; align-items: center; margin-bottom: 12px;}
.wa-playbook-title {font-size: 1.4rem; font-weight: 700; letter-spacing: -0.03em; color: #3e2618;}
.wa-chip-row {display: flex; flex-wrap: wrap; gap: 8px; margin-top: 12px;}
.wa-chip {
  display: inline-flex;
  align-items: center;
  padding: 7px 10px;
  border-radius: 999px;
  background: rgba(172, 121, 80, 0.10);
  color: #7b4f32;
  font-size: 0.84rem;
  font-weight: 600;
}
.wa-lane-header {display: flex; justify-content: space-between; gap: 12px; align-items: flex-start; margin-bottom: 8px;}
.wa-lane-title {font-size: 1.32rem; font-weight: 700; letter-spacing: -0.03em; color: #3e2618;}
.wa-lane-copy {font-size: 0.96rem; color: #78624f;}
.wa-hint {
  margin-bottom: 14px;
  padding: 12px 14px;
  border-radius: 18px;
  background: rgba(174, 126, 88, 0.08);
  border: 1px solid rgba(174, 126, 88, 0.12);
  color: #6d4a32;
}
.wa-card-grid {display: grid; grid-template-columns: repeat(auto-fit, minmax(235px, 1fr)); gap: 12px;}
.wa-task-card {
  background: linear-gradient(180deg, #fffaf2 0%, #f7ecde 100%);
  border: 1px solid rgba(175, 135, 100, 0.18);
  border-radius: 22px;
  padding: 14px;
  color: #34261d;
}
.wa-task-card.running {background: linear-gradient(180deg, #f4ebe0 0%, #ecdcca 100%);}
.wa-task-card.recommended {outline: 2px solid rgba(226, 145, 87, 0.9); outline-offset: 2px;}
.wa-task-head {display: flex; justify-content: space-between; gap: 10px; align-items: flex-start; margin-bottom: 10px;}
.wa-task-name {font-size: 1.08rem; font-weight: 700; color: #3c271b;}
.wa-badge {
  display: inline-flex;
  align-items: center;
  padding: 4px 8px;
  border-radius: 999px;
  background: rgba(170, 123, 84, 0.12);
  color: #825336;
  font-size: 0.68rem;
  font-weight: 700;
  letter-spacing: 0.08em;
  text-transform: uppercase;
}
.wa-badge.urgent {background: rgba(208, 108, 97, 0.16); color: #8c3c35;}
.wa-badge.active {background: rgba(151, 179, 120, 0.18); color: #5f7142;}
.wa-badge.recommended {background: rgba(229, 166, 93, 0.20); color: #86501f;}
.wa-badge.retry {background: rgba(176, 141, 78, 0.18); color: #7a5a22;}
.wa-task-meta {display: flex; flex-wrap: wrap; gap: 8px; margin-bottom: 10px;}
.wa-task-meta span {
  padding: 5px 8px;
  border-radius: 999px;
  background: rgba(179, 139, 104, 0.10);
  font-size: 0.76rem;
  color: #77553b;
}
.wa-metrics {display: grid; grid-template-columns: repeat(2, minmax(0, 1fr)); gap: 10px 12px;}
.wa-metric span {display: block; font-size: 0.68rem; letter-spacing: 0.08em; text-transform: uppercase; color: #a26c45;}
.wa-metric strong {display: block; margin-top: 4px; font-size: 0.98rem; color: #35261d;}
.wa-empty {
  padding: 20px;
  border-radius: 20px;
  border: 1px dashed rgba(176, 133, 98, 0.26);
  background: rgba(178, 143, 112, 0.06);
  color: #7b6656;
  text-align: center;
}
.wa-action-row {display: flex; flex-wrap: wrap; gap: 10px; margin-top: 14px;}
.wa-button-primary button {background: linear-gradient(135deg, #d97b4b, #c95f34) !important; color: #fff7f0 !important; border: none !important;}
.wa-button-secondary button {background: #8f5b3b !important; color: #fff8f2 !important; border: none !important;}
.wa-button-ghost button {background: rgba(180, 132, 96, 0.08) !important; color: #7a4d31 !important; border: 1px solid rgba(180, 132, 96, 0.16) !important;}
.wa-plot-wrap {padding: 10px 10px 2px;}
.wa-footer-stack {display: grid; gap: 18px; margin-top: 18px;}
.wa-accordion {border-radius: 20px !important; overflow: hidden;}
@media (max-width: 1080px) {
  .wa-hero,
  .wa-main {grid-template-columns: 1fr;}
  .wa-topbar {grid-template-columns: repeat(2, minmax(0, 1fr));}
}
@media (max-width: 760px) {
  .wa-control-grid {grid-template-columns: 1fr;}
  .wa-banner-grid {grid-template-columns: repeat(2, minmax(0, 1fr));}
  .wa-topbar {grid-template-columns: 1fr;}
}
"""


def _blank_session() -> Session:
    return {"env": WorkflowArenaEnvironment(), "observation": None, "history": []}


def _fmt_num(value: Any, digits: int = 3) -> str:
    if value is None:
        return "—"
    if isinstance(value, float):
        return f"{value:.{digits}f}"
    return str(value)


def _preset_html(preset: str) -> str:
    brief = PRESET_BRIEFS.get(preset, PRESET_BRIEFS[DifficultyPreset.EASY.value])
    preset_config = get_preset_config(DifficultyPreset(preset))
    budget_note = (
        "No fixed time budget."
        if preset_config.time_budget_multiplier is None
        else f"Time budget uses {preset_config.time_budget_multiplier:.2f}x the lower-bound makespan."
    )
    return (
        '<div class="wa-preset-card wa-card">'
        '<div class="eyebrow">Preset brief</div>'
        f'<div class="name">{brief["label"]}</div>'
        f'<div class="summary">{brief["summary"]}</div>'
        f'<div class="wa-preset-meta">{budget_note} {brief["mechanics"]}</div>'
        f'<div class="wa-preset-focus"><strong>What matters now:</strong> {brief["focus"]}</div>'
        "</div>"
    )


def _status_text(observation: Any) -> tuple[str, str]:
    if observation.validation_error:
        return "Invalid action", "bad"
    if observation.done and observation.termination_reason:
        return "Episode terminated", "bad"
    if observation.done:
        return "Workflow completed", "good"
    if observation.free_workers == 0 and observation.running_tasks:
        return "Wait required", ""
    if observation.ready_tasks:
        return "Ready to dispatch", ""
    return "Waiting on completions", ""


def _recommended_task_ids(observation: Any) -> list[str]:
    if observation is None or observation.done or observation.free_workers <= 0:
        return []
    ready_tasks = list(observation.ready_tasks)
    if not ready_tasks:
        return []
    time_remaining = observation.time_remaining
    ranked = sorted(
        ready_tasks,
        key=lambda task: (
            time_remaining is not None and task.duration > time_remaining,
            max(0, task.duration - time_remaining) if time_remaining is not None else 0,
            task.slack if task.slack is not None else 1_000_000,
            task.deadline if task.deadline is not None else 1_000_000,
            -(task.criticality or 0.0),
            -task.priority,
            task.duration,
            task.task_id,
        ),
    )
    return [task.task_id for task in ranked[: observation.free_workers]]


def _dispatch_window(observation: Any) -> tuple[int, int, int]:
    ready_count = len(observation.ready_tasks)
    free_workers = max(0, observation.free_workers)
    dispatchable_now = min(ready_count, free_workers)
    overflow_ready = max(0, ready_count - free_workers)
    return ready_count, dispatchable_now, overflow_ready


def _topbar_html(observation: Any) -> str:
    completed = observation.progress.completed
    total = max(1, observation.progress.total)
    score = observation.benchmark_score
    if score is None:
        score = observation.success_metrics.benchmark_score
    time_sub = (
        f"{observation.time_remaining} remaining"
        if observation.time_remaining is not None
        else "simulation clock"
    )
    worker_sub = (
        f"idle / usable of {observation.total_workers}"
        if getattr(observation, "effective_workers", observation.total_workers) != observation.total_workers
        else "free / total"
    )
    cards = [
        ("State", _status_text(observation)[0], f"{observation.progress.ready} ready / {observation.progress.running} running"),
        (
            "Workers",
            f"{observation.free_workers}/{getattr(observation, 'effective_workers', observation.total_workers)}",
            worker_sub,
        ),
        ("Completed", f"{completed}/{total}", f"{round(100 * completed / total, 1)}% finished"),
        ("Reward", _fmt_num(observation.cumulative_reward, 3), "cumulative"),
        ("Time", observation.current_time, time_sub),
        ("Score", _fmt_num(score, 3), "terminal if done"),
    ]
    return '<div class="wa-topbar">' + "".join(
        f'<div class="wa-stat"><div class="label">{label}</div><div class="value">{value}</div><div class="sub">{sub}</div></div>'
        for label, value, sub in cards
    ) + "</div>"


def _banner_html(observation: Any) -> str:
    completed = observation.progress.completed
    total = max(1, observation.progress.total)
    progress_pct = round(100 * completed / total, 1)
    status_text, status_kind = _status_text(observation)
    banner_class = "wa-banner"
    if status_kind == "bad":
        banner_class += " invalid"
    elif status_kind == "good":
        banner_class += " done"

    failure_note = _failure_summary(observation)
    note = observation.note or "No environment note."
    if observation.validation_error:
        note = f"{note} {observation.validation_error}"
    if failure_note:
        note = f"{note} {failure_note}"

    score = observation.benchmark_score
    if score is None:
        score = observation.success_metrics.benchmark_score

    metric_cards = [
        ("Ready", observation.progress.ready),
        ("Running", observation.progress.running),
        ("Workers", f"{getattr(observation, 'effective_workers', observation.total_workers)}/{observation.total_workers}"),
        (
            "Time Left",
            observation.time_remaining if observation.time_remaining is not None else "—",
        ),
    ]

    return (
        f'<div class="{banner_class}">'
        '<div class="wa-banner-top">'
        '<div>'
        f'<span class="status">{status_text}</span>'
        f'<div class="meta">Preset: {observation.config.preset.value} • Seed: {observation.config.seed} • Workers: {observation.total_workers}</div>'
        f'<div class="note">{note}</div>'
        "</div>"
        f'<div class="meta">{completed}/{total} complete</div>'
        "</div>"
        '<div class="wa-banner-grid">'
        + "".join(
            f'<div class="wa-banner-metric"><span>{label}</span><strong>{value}</strong></div>'
            for label, value in metric_cards
        )
        + "</div>"
        f'<div class="wa-progress"><div class="wa-progress-fill" style="width:{progress_pct:.1f}%"></div></div>'
        "</div>"
    )


def _planner_html(observation: Any) -> str:
    recommended = _recommended_task_ids(observation)
    ready_count, dispatchable_now, overflow_ready = _dispatch_window(observation)
    if observation.done:
        title = "Episode finished"
        body = "Reset for another episode or inspect the final timeline and reward trace below."
    elif observation.validation_error:
        title = "Fix the last move"
        body = observation.validation_error
    elif observation.free_workers == 0 and observation.running_tasks:
        title = "Advance time"
        body = "All workers are occupied. Waiting is the only legal move until the next task completes."
    elif recommended:
        title = f"Dispatch {', '.join(recommended)}"
        body = (
            "These tasks minimize slack first, then prefer tighter deadlines, stronger criticality, and higher priority. "
            f"The recommendation is capped at `{dispatchable_now}` because only `{observation.free_workers}` worker"
            f"{'s are' if observation.free_workers != 1 else ' is'} free right now."
        )
    else:
        title = "Hold for completions"
        body = "No ready work is available. Wait until dependencies unlock new tasks."

    chips = [
        f"free workers: {observation.free_workers}",
        f"usable workers: {getattr(observation, 'effective_workers', observation.total_workers)}",
        f"ready queue: {ready_count}",
        f"dispatchable now: {dispatchable_now}",
        f"last reward: {_fmt_num(observation.reward if hasattr(observation, 'reward') else 0.0, 3)}",
    ]
    if observation.time_remaining is not None:
        chips.append(f"time remaining: {observation.time_remaining}")
    if overflow_ready:
        chips.append(f"queued beyond capacity: {overflow_ready}")
    if observation.running_tasks:
        next_finish = min(task.end_time or observation.current_time for task in observation.running_tasks)
        chips.append(f"next completion: t={next_finish}")
    if observation.degraded_workers:
        chips.append(f"worker outage: -{observation.degraded_workers} usable")

    return (
        '<div class="wa-playbook wa-card">'
        '<div class="wa-playbook-header">'
        '<div>'
        '<h3>Decision support</h3>'
        f'<div class="wa-playbook-title">{title}</div>'
        "</div>"
        f'<div class="wa-chip">{_status_text(observation)[0]}</div>'
        "</div>"
        f'<p>{body}</p>'
        '<div class="wa-chip-row">'
        + "".join(f'<div class="wa-chip">{chip}</div>' for chip in chips)
        + "</div>"
        "</div>"
    )


def _capacity_hint(observation: Any) -> str:
    ready_count, dispatchable_now, overflow_ready = _dispatch_window(observation)
    if observation.done:
        return "Episode finished. Review the schedule or reset to try another seed."
    if observation.validation_error:
        return (
            f"Last action was rejected. Select at most {observation.free_workers} ready "
            f"task{'s' if observation.free_workers != 1 else ''}."
        )
    if observation.free_workers == 0 and observation.running_tasks:
        return "All workers are busy. Use Wait to jump to the next completion."
    if not observation.ready_tasks:
        return "No ready tasks available right now. Wait until dependencies unlock more work."
    overflow_suffix = f" `{overflow_ready}` ready task(s) will stay queued." if overflow_ready else ""
    return (
        f"{ready_count} ready task{'s' if ready_count != 1 else ''}. "
        f"You can dispatch up to {dispatchable_now} right now.{overflow_suffix}"
    )


def _task_badges(task: Any, *, running: bool = False, recommended: bool = False) -> str:
    badges: list[str] = []
    if task.deadline is not None and task.slack is not None and task.slack <= 1:
        badges.append('<span class="wa-badge urgent">Urgent</span>')
    if getattr(task, "attempt_count", 0) > 0:
        badges.append(
            f'<span class="wa-badge retry">Retry {getattr(task, "attempt_count", 0) + 1}</span>'
        )
    if recommended:
        badges.append('<span class="wa-badge recommended">Recommended</span>')
    if running:
        badges.append('<span class="wa-badge active">Running</span>')
    if not badges:
        badges.append('<span class="wa-badge">Ready</span>')
    return "".join(badges)


def _task_card(task: Any, *, running: bool = False, recommended: bool = False) -> str:
    deps = ", ".join(task.dependencies) if task.dependencies else "None"
    deadline = task.deadline if task.deadline is not None else "—"
    start = task.start_time if task.start_time is not None else "—"
    end = task.end_time if task.end_time is not None else "—"
    classes = ["wa-task-card"]
    if running:
        classes.append("running")
    if recommended:
        classes.append("recommended")
    return (
        f'<div class="{" ".join(classes)}">'
        '<div class="wa-task-head">'
        f'<div class="wa-task-name">{task.task_id}</div>'
        f'<div>{_task_badges(task, running=running, recommended=recommended)}</div>'
        "</div>"
        f'<div class="wa-task-meta"><span>deps: {deps}</span><span>downstream: {task.downstream_count}</span><span>attempts: {getattr(task, "attempt_count", 0) + 1}</span></div>'
        '<div class="wa-metrics">'
        f'<div class="wa-metric"><span>Deadline</span><strong>{deadline}</strong></div>'
        f'<div class="wa-metric"><span>Duration</span><strong>{task.duration}</strong></div>'
        f'<div class="wa-metric"><span>Priority</span><strong>{task.priority}</strong></div>'
        f'<div class="wa-metric"><span>Criticality</span><strong>{_fmt_num(task.criticality, 3)}</strong></div>'
        f'<div class="wa-metric"><span>Slack</span><strong>{_fmt_num(task.slack, 1)}</strong></div>'
        f'<div class="wa-metric"><span>{"Finish" if running else "Start"}</span><strong>{end if running else start}</strong></div>'
        "</div>"
        "</div>"
    )


def _cards_html(tasks: list[Any], *, running: bool = False, recommended_ids: set[str] | None = None) -> str:
    if not tasks:
        message = "No tasks in this lane yet." if running else "No ready tasks available."
        return f'<div class="wa-empty">{message}</div>'
    recommended_ids = recommended_ids or set()
    return '<div class="wa-card-grid">' + "".join(
        _task_card(task, running=running, recommended=task.task_id in recommended_ids)
        for task in tasks
    ) + "</div>"


def _timeline_figure(observation: Any) -> go.Figure:
    fig = go.Figure()

    completed = sorted(observation.completed_tasks, key=lambda task: task.task_id)
    running = sorted(observation.running_tasks, key=lambda task: task.task_id)
    ready = sorted(observation.ready_tasks, key=lambda task: task.task_id)

    timeline_tasks = completed + running
    task_ids = [task.task_id for task in timeline_tasks] + [task.task_id for task in ready]

    if completed:
        fig.add_trace(
            go.Bar(
                x=[max(0, (task.end_time or 0) - (task.start_time or 0)) for task in completed],
                y=[task.task_id for task in completed],
                base=[task.start_time or 0 for task in completed],
                orientation="h",
                name="Completed",
                marker_color="#85c88a",
                hovertemplate=(
                    "<b>%{y}</b><br>Status: Completed<br>Start: %{base}<br>"
                    "Duration: %{x}<extra></extra>"
                ),
            )
        )

    if running:
        fig.add_trace(
            go.Bar(
                x=[
                    max(
                        0,
                        (task.end_time or observation.current_time) - (task.start_time or observation.current_time),
                    )
                    for task in running
                ],
                y=[task.task_id for task in running],
                base=[task.start_time or observation.current_time for task in running],
                orientation="h",
                name="Running",
                marker_color="#d88a5b",
                hovertemplate=(
                    "<b>%{y}</b><br>Status: Running<br>Start: %{base}<br>"
                    "Allocated span: %{x}<extra></extra>"
                ),
            )
        )

    if ready:
        fig.add_trace(
            go.Scatter(
                x=[observation.current_time] * len(ready),
                y=[task.task_id for task in ready],
                mode="markers",
                name="Ready",
                marker=dict(color="#9e6a43", size=11, symbol="diamond"),
                customdata=[[task.deadline, task.priority, task.duration] for task in ready],
                hovertemplate=(
                    "<b>%{y}</b><br>Status: Ready<br>Current time: %{x}<br>"
                    "Deadline: %{customdata[0]}<br>Priority: %{customdata[1]}<br>"
                    "Duration: %{customdata[2]}<extra></extra>"
                ),
            )
        )

    if not task_ids:
        task_ids = ["No tasks yet"]
        fig.add_annotation(
            text="Reset an episode to populate the workflow timeline.",
            x=0.5,
            y=0.5,
            xref="paper",
            yref="paper",
            showarrow=False,
            font=dict(color="#8c6f58", size=14),
        )

    horizon_candidates = [observation.current_time + 1]
    horizon_candidates.extend(task.end_time or 0 for task in completed)
    horizon_candidates.extend(task.end_time or observation.current_time for task in running)
    x_max = max(horizon_candidates) + 1

    fig.add_vline(
        x=observation.current_time,
        line_width=2,
        line_dash="dash",
        line_color="#9f6b48",
        annotation_text="Now",
        annotation_position="top left",
    )

    fig.update_layout(
        barmode="overlay",
        height=max(280, 90 + 34 * len(task_ids)),
        margin=dict(l=10, r=10, t=44, b=18),
        paper_bgcolor="#ffffff",
        plot_bgcolor="#ffffff",
        font=dict(color="#4d382b", family="IBM Plex Sans, Arial, sans-serif"),
        legend=dict(orientation="h", yanchor="bottom", y=1.02, x=0),
        title=dict(text="Workflow Timeline", x=0.02, font=dict(size=18)),
        xaxis=dict(
            title="Simulated Time",
            range=[0, x_max],
            gridcolor="#eadfce",
            zeroline=False,
            linecolor="#cfb79c",
            title_font=dict(color="#6c4a33"),
            tickfont=dict(color="#6c4a33"),
        ),
        yaxis=dict(
            title="Tasks",
            categoryorder="array",
            categoryarray=list(reversed(task_ids)),
            gridcolor="#f2e8da",
            linecolor="#cfb79c",
            title_font=dict(color="#6c4a33"),
            tickfont=dict(color="#6c4a33"),
        ),
    )
    return fig


def _detail_rows(tasks: list[Any]) -> list[list[Any]]:
    return [
        [
            task.task_id,
            task.priority,
            task.duration,
            task.deadline if task.deadline is not None else "—",
            _fmt_num(task.criticality, 3),
            _fmt_num(task.slack, 1),
            len(task.dependencies),
            task.downstream_count,
            getattr(task, "attempt_count", 0) + 1,
            task.start_time if task.start_time is not None else "—",
            task.end_time if task.end_time is not None else "—",
        ]
        for task in tasks
    ]


def _failure_summary(observation: Any) -> str:
    events = getattr(observation, "recent_failure_events", []) or []
    if not events:
        return ""
    return " ".join(event.detail for event in events if getattr(event, "detail", ""))


def _selection_markdown(selected_task_ids: list[str], observation: Any) -> str:
    if observation is None:
        return "No episode yet. Reset an episode to start building a dispatch batch."

    task_map = {task.task_id: task for task in observation.ready_tasks}
    selected_tasks = [task_map[task_id] for task_id in selected_task_ids if task_id in task_map]
    capacity = max(0, observation.free_workers)
    ready_count, dispatchable_now, overflow_ready = _dispatch_window(observation)

    if not selected_tasks:
        recommended = _recommended_task_ids(observation)
        if not recommended:
            return (
                f"**Dispatch builder**\n\nReady queue: `{ready_count}`. "
                f"Dispatchable now: `{dispatchable_now}`."
            )
        overflow_suffix = f" `{overflow_ready}` ready task(s) stay queued after dispatch." if overflow_ready else ""
        return (
            f"**Dispatch builder**\n\nNo tasks selected yet. Ready queue: `{ready_count}`. "
            f"Recommended batch: `{', '.join(recommended)}`. Dispatch cap now: `{dispatchable_now}`.{overflow_suffix}"
        )

    total_priority = sum(task.priority for task in selected_tasks)
    shortest_finish = observation.current_time + min(task.duration for task in selected_tasks)
    longest_finish = observation.current_time + max(task.duration for task in selected_tasks)
    warnings: list[str] = []
    if len(selected_tasks) > capacity:
        warnings.append(f"Selection exceeds capacity by `{len(selected_tasks) - capacity}`.")

    warning_text = "\n\n" + " ".join(warnings) if warnings else ""
    return (
        f"**Dispatch builder**\n\nSelected `{len(selected_tasks)}` task(s) for `{capacity}` free slot(s). "
        f"Priority sum: `{total_priority}`. Earliest completion: `t={shortest_finish}`. "
        f"Longest in-flight span: `t={longest_finish}`.{warning_text}"
    )


def _reward_markdown(observation: Any) -> str:
    breakdown = observation.last_reward_breakdown
    rows = [
        ("completion", breakdown.completion_reward),
        ("utilization", breakdown.utilization_reward),
        ("deadline", breakdown.deadline_reward),
        ("criticality", breakdown.criticality_reward),
        ("idle", breakdown.idle_penalty),
        ("invalid", breakdown.invalid_action_penalty),
        ("terminal", breakdown.terminal_makespan_score),
        ("unfinished", breakdown.unfinished_task_penalty),
    ]
    lines = ["| Channel | Value |", "| --- | ---: |"]
    lines.extend(f"| {label} | {value:.3f} |" for label, value in rows)
    return "\n".join(lines)


def _history_markdown(history: list[dict[str, Any]]) -> str:
    if not history:
        return "No actions yet."
    lines: list[str] = []
    for item in history[-12:]:
        reward = _fmt_num(item.get("reward"), 3)
        suffix = f" • error: `{item['error']}`" if item.get("error") else ""
        note = item.get("note") or ""
        lines.append(f"**{item['label']}** at `t={item['time']}` • reward `{reward}`{suffix}  \n{note}")
    return "\n\n".join(reversed(lines))


def _blank_observation_view() -> Any:
    return SimpleNamespace(
        reward=0.0,
        progress=SimpleNamespace(completed=0, ready=0, running=0, blocked=0, total=1),
        benchmark_score=None,
        success_metrics=SimpleNamespace(benchmark_score=None, unfinished_task_count=0),
        free_workers=0,
        effective_workers=0,
        degraded_workers=0,
        total_workers=0,
        time_budget=None,
        time_remaining=None,
        cumulative_reward=0.0,
        current_time=0,
        done=False,
        termination_reason=None,
        validation_error=None,
        completed_tasks=[],
        ready_tasks=[],
        running_tasks=[],
        blocked_tasks=[],
        recent_failure_events=[],
        note="Reset an episode to start scheduling.",
        config=SimpleNamespace(preset=SimpleNamespace(value=DifficultyPreset.EASY.value), seed=0),
    )


def _empty_updates(session: Session):
    empty_rows: list[list[Any]] = []
    blank = _blank_observation_view()
    return (
        session,
        _preset_html(DifficultyPreset.EASY.value),
        _topbar_html(blank),
        _banner_html(blank),
        _planner_html(blank),
        "No episode yet.",
        _selection_markdown([], blank),
        '<div class="wa-empty">Reset an episode to see ready tasks.</div>',
        gr.update(choices=[], value=[]),
        gr.update(interactive=False),
        gr.update(interactive=False),
        gr.update(interactive=False),
        gr.update(interactive=False),
        gr.update(interactive=False),
        '<div class="wa-empty">Running tasks will appear here after dispatch.</div>',
        _timeline_figure(blank),
        "| Channel | Value |\n| --- | ---: |\n| — | — |",
        "No actions yet.",
        empty_rows,
        empty_rows,
        empty_rows,
    )


def _render(session: Session):
    observation = session.get("observation")
    env = session.get("env")
    history = session.get("history", [])
    if observation is None:
        return _empty_updates(session)

    if env is None:
        completed_rows = _detail_rows(observation.completed_tasks)
        blocked_rows = _detail_rows(observation.blocked_tasks)
    else:
        completed_rows = _detail_rows(env.debug_task_views_for_status(TaskStatus.COMPLETED))
        blocked_rows = _detail_rows(env.debug_task_views_for_status(TaskStatus.BLOCKED))

    ready_choices = [task.task_id for task in observation.ready_tasks]
    recommended_ids = _recommended_task_ids(observation)
    can_recommend = bool(recommended_ids) and not observation.done
    can_wait = bool(observation.running_tasks) and not observation.done
    can_clear = bool(ready_choices) and not observation.done

    return (
        session,
        _preset_html(observation.config.preset.value),
        _topbar_html(observation),
        _banner_html(observation),
        _planner_html(observation),
        _capacity_hint(observation),
        _selection_markdown([], observation),
        _cards_html(observation.ready_tasks, running=False, recommended_ids=set(recommended_ids)),
        gr.update(choices=ready_choices, value=[]),
        gr.update(interactive=False),
        gr.update(interactive=can_wait),
        gr.update(interactive=can_recommend),
        gr.update(interactive=can_recommend),
        gr.update(interactive=can_clear),
        _cards_html(observation.running_tasks, running=True),
        _timeline_figure(observation),
        _reward_markdown(observation),
        _history_markdown(history),
        empty_rows := _detail_rows([]),
        completed_rows,
        blocked_rows,
    )


def _append_history(
    session: Session,
    label: str,
    observation: Any,
    *,
    reward: float | None = None,
    error: str | None = None,
) -> Session:
    history = list(session.get("history", []))
    history.append(
        {
            "label": label,
            "time": observation.current_time,
            "reward": reward,
            "error": error,
            "note": observation.note,
        }
    )
    session["history"] = history
    return session


def _reset(preset: str, seed: float, worker_count: float, session: Session):
    env = session.get("env") or WorkflowArenaEnvironment()
    observation = env.reset(
        preset=preset,
        seed=int(seed),
        worker_count=int(worker_count),
    )
    next_session = {"env": env, "observation": observation, "history": []}
    next_session = _append_history(
        next_session,
        f"reset • preset `{preset}` • workers `{int(worker_count)}`",
        observation,
        reward=0.0,
        error=None,
    )
    return _render(next_session)


def _dispatch(selected_task_ids: list[str], session: Session):
    observation = session.get("observation")
    env = session.get("env")
    if env is None or observation is None:
        return _render(_blank_session())

    action = WorkflowArenaAction(
        action_type=WorkflowActionType.DISPATCH,
        task_ids=selected_task_ids,
    )
    next_observation = env.step(action)
    next_session = {
        "env": env,
        "observation": next_observation,
        "history": session.get("history", []),
    }
    label = "dispatch " + (", ".join(selected_task_ids) if selected_task_ids else "(none)")
    next_session = _append_history(
        next_session,
        label,
        next_observation,
        reward=next_observation.reward,
        error=next_observation.validation_error,
    )
    return _render(next_session)


def _dispatch_recommended(session: Session):
    observation = session.get("observation")
    if observation is None:
        return _render(_blank_session())
    return _dispatch(_recommended_task_ids(observation), session)


def _wait(session: Session):
    observation = session.get("observation")
    env = session.get("env")
    if env is None or observation is None:
        return _render(_blank_session())

    action = WorkflowArenaAction(
        action_type=WorkflowActionType.WAIT,
        task_ids=[],
    )
    next_observation = env.step(action)
    next_session = {
        "env": env,
        "observation": next_observation,
        "history": session.get("history", []),
    }
    next_session = _append_history(
        next_session,
        "wait",
        next_observation,
        reward=next_observation.reward,
        error=next_observation.validation_error,
    )
    return _render(next_session)


def _update_selection(selected_task_ids: list[str], session: Session):
    observation = session.get("observation")
    if observation is None:
        return "No episode yet.", [], gr.update(interactive=False)

    task_map = {task.task_id: task for task in observation.ready_tasks}
    selected_tasks = [task_map[task_id] for task_id in selected_task_ids if task_id in task_map]
    can_dispatch = (
        bool(selected_tasks)
        and len(selected_tasks) == len(selected_task_ids)
        and len(selected_task_ids) <= observation.free_workers
        and not observation.done
    )
    return (
        _selection_markdown(selected_task_ids, observation),
        _detail_rows(selected_tasks),
        gr.update(interactive=can_dispatch),
    )


def _select_recommended(session: Session):
    observation = session.get("observation")
    if observation is None:
        return gr.update(value=[]), "No episode yet.", [], gr.update(interactive=False)
    recommended_ids = _recommended_task_ids(observation)
    task_map = {task.task_id: task for task in observation.ready_tasks}
    selected_tasks = [task_map[task_id] for task_id in recommended_ids if task_id in task_map]
    return (
        gr.update(value=recommended_ids),
        _selection_markdown(recommended_ids, observation),
        _detail_rows(selected_tasks),
        gr.update(interactive=bool(recommended_ids) and not observation.done),
    )


def _clear_selection(session: Session):
    observation = session.get("observation")
    return (
        gr.update(value=[]),
        _selection_markdown([], observation),
        [],
        gr.update(interactive=False),
    )


def _random_seed() -> int:
    return random.randint(0, 999_999)


def _preset_controls_update(preset: str):
    preset_config = get_preset_config(DifficultyPreset(preset))
    return _preset_html(preset), gr.update(value=preset_config.worker_count)


def create_gradio_app() -> gr.Blocks:
    with gr.Blocks(title="WorkflowArena") as demo:
        session = gr.State(_blank_session())

        with gr.Column(elem_classes=["wa-shell"]):
            gr.HTML(f"<style>{CSS}</style>")
            gr.HTML(
                """
                <div class="wa-title">
                  <h1>WorkflowArena</h1>
                  <p>Run a workflow episode like a control room instead of a raw form. Reset a seeded DAG, inspect urgency and capacity, build a legal dispatch batch, then advance time when workers are saturated.</p>
                </div>
                """
            )

            with gr.Row(elem_classes=["wa-hero"]):
                with gr.Column(elem_classes=["wa-control-card", "wa-card"]):
                    gr.HTML("<h3>Episode controls</h3><p>Change the preset, seed, or worker count, then reset to generate a new scheduling problem.</p>")
                    with gr.Accordion("Problem Framing", open=False, elem_classes=["wa-accordion", "wa-compact-accordion"]):
                        gr.HTML(
                            """
                            <div class="wa-problem-box">
                              <p><strong>What this problem is:</strong> You are scheduling a workflow where tasks depend on each other and workers are limited. At every step, the legal move is either to dispatch ready tasks to free workers or wait for the next completion event.</p>
                              <p><strong>What good play looks like:</strong> Finish urgent and high-value work on time, keep workers utilized, and avoid delaying the critical path. Higher difficulties add a time budget and, in hard mode, failure events that reduce usable capacity or force retries.</p>
                            </div>
                            """
                        )
                    with gr.Row(elem_classes=["wa-control-grid"]):
                        preset = gr.Dropdown(
                            label="Preset",
                            choices=[preset.value for preset in DifficultyPreset],
                            value=DifficultyPreset.EASY.value,
                            interactive=True,
                        )
                        seed = gr.Number(label="Seed", value=0, precision=0, minimum=0)
                        workers = gr.Slider(
                            minimum=1,
                            maximum=6,
                            step=1,
                            value=3,
                            label="Workers",
                            interactive=True,
                        )
                    with gr.Row(elem_classes=["wa-control-buttons", "wa-inline-buttons"]):
                        reset_button = gr.Button(
                            "Reset Episode",
                            variant="primary",
                            elem_classes=["wa-button-primary"],
                        )
                        random_seed_button = gr.Button(
                            "Random Seed",
                            variant="secondary",
                            elem_classes=["wa-button-ghost"],
                        )

                preset_brief = gr.HTML(value=_preset_html(DifficultyPreset.EASY.value))

            topbar = gr.HTML()
            banner = gr.HTML()
            planner = gr.HTML()

            with gr.Row(elem_classes=["wa-main"]):
                with gr.Column(elem_classes=["wa-left-stack"]):
                    with gr.Column(elem_classes=["wa-panel", "wa-card"]):
                        gr.HTML(
                            """
                            <div class="wa-lane-header">
                              <div>
                                <div class="wa-lane-title">Dispatch Lane</div>
                                <div class="wa-lane-copy">Inspect ready tasks, build a batch, and send work only when the action is legal.</div>
                              </div>
                            </div>
                            """
                        )
                        selection_summary = gr.Markdown(elem_classes=["wa-hint"])
                        ready_cards = gr.HTML()
                        ready_selector = gr.CheckboxGroup(
                            label="Build dispatch batch",
                            info="Choose up to the number of currently free workers.",
                        )
                        with gr.Row(elem_classes=["wa-action-row"]):
                            select_recommended_button = gr.Button(
                                "Select Recommended",
                                variant="secondary",
                                interactive=False,
                                elem_classes=["wa-button-secondary"],
                            )
                            dispatch_recommended_button = gr.Button(
                                "Dispatch Recommended",
                                variant="primary",
                                interactive=False,
                                elem_classes=["wa-button-primary"],
                            )
                            clear_selection_button = gr.Button(
                                "Clear Selection",
                                variant="secondary",
                                interactive=False,
                                elem_classes=["wa-button-ghost"],
                            )
                            dispatch_button = gr.Button(
                                "Dispatch Selected",
                                variant="primary",
                                interactive=False,
                                elem_classes=["wa-button-primary"],
                            )
                            wait_button = gr.Button(
                                "Wait",
                                variant="secondary",
                                interactive=False,
                                elem_classes=["wa-button-secondary"],
                            )

                    with gr.Column(elem_classes=["wa-panel", "wa-card"]):
                        gr.HTML(
                            """
                            <div class="wa-lane-header">
                              <div>
                                <div class="wa-lane-title">Selected Batch</div>
                                <div class="wa-lane-copy">This preview updates as you pick tasks from the current ready queue.</div>
                              </div>
                            </div>
                            """
                        )
                        selected_table = gr.Dataframe(
                            headers=DETAIL_HEADERS,
                            value=[],
                            interactive=False,
                            wrap=True,
                            label="Dispatch preview",
                        )

                with gr.Column(elem_classes=["wa-right-stack"]):
                    with gr.Column(elem_classes=["wa-panel", "wa-card"]):
                        gr.HTML(
                            """
                            <div class="wa-lane-header">
                              <div>
                                <div class="wa-lane-title">Mission Control</div>
                                <div class="wa-lane-copy">Use the recommendation as a guide, not a rule. The reward trace below tells you whether the choice paid off.</div>
                              </div>
                            </div>
                            """
                        )
                        gr.Markdown(elem_classes=["wa-hint"], value="No episode yet.")
                        decision_hint = gr.Markdown(elem_classes=["wa-hint"])
                        running_cards = gr.HTML()

                    with gr.Column(elem_classes=["wa-plot-wrap", "wa-card"]):
                        timeline_plot = gr.Plot(label="Workflow Timeline")

                    with gr.Accordion("Reward Breakdown", open=False, elem_classes=["wa-accordion"]):
                        reward_markdown = gr.Markdown()

                    with gr.Accordion("Action History", open=False, elem_classes=["wa-accordion"]):
                        history_markdown = gr.Markdown()

            with gr.Row(elem_classes=["wa-footer-stack"]):
                with gr.Accordion("Completed Tasks", open=False, elem_classes=["wa-accordion"]):
                    completed_table = gr.Dataframe(
                        headers=DETAIL_HEADERS,
                        value=[],
                        interactive=False,
                        wrap=True,
                        label="Completed",
                    )

                with gr.Accordion("Blocked Tasks", open=False, elem_classes=["wa-accordion"]):
                    blocked_table = gr.Dataframe(
                        headers=DETAIL_HEADERS,
                        value=[],
                        interactive=False,
                        wrap=True,
                        label="Blocked",
                    )

        outputs = [
            session,
            preset_brief,
            topbar,
            banner,
            planner,
            decision_hint,
            selection_summary,
            ready_cards,
            ready_selector,
            dispatch_button,
            wait_button,
            select_recommended_button,
            dispatch_recommended_button,
            clear_selection_button,
            running_cards,
            timeline_plot,
            reward_markdown,
            history_markdown,
            selected_table,
            completed_table,
            blocked_table,
        ]

        random_seed_button.click(_random_seed, outputs=[seed])
        preset.change(_preset_controls_update, inputs=[preset], outputs=[preset_brief, workers])

        reset_button.click(
            _reset,
            inputs=[preset, seed, workers, session],
            outputs=outputs,
        )
        dispatch_button.click(
            _dispatch,
            inputs=[ready_selector, session],
            outputs=outputs,
        )
        dispatch_recommended_button.click(
            _dispatch_recommended,
            inputs=[session],
            outputs=outputs,
        )
        wait_button.click(
            _wait,
            inputs=[session],
            outputs=outputs,
        )

        ready_selector.change(
            _update_selection,
            inputs=[ready_selector, session],
            outputs=[selection_summary, selected_table, dispatch_button],
        )
        select_recommended_button.click(
            _select_recommended,
            inputs=[session],
            outputs=[ready_selector, selection_summary, selected_table, dispatch_button],
        )
        clear_selection_button.click(
            _clear_selection,
            inputs=[session],
            outputs=[ready_selector, selection_summary, selected_table, dispatch_button],
        )

        demo.load(lambda: _empty_updates(_blank_session()), outputs=outputs)

    return demo