"""
Premium Gradio UI for the AI Executive Assistant Simulator.

Features a glassmorphic dark theme with smooth animations,
interactive agent comparison, and real-time schedule visualization.
"""

import sys
import os
import json

sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))

import gradio as gr
import plotly.graph_objects as go

from env.assistant_env import ExecutiveAssistantEnv
from agents.random_agent import RandomAgent
from agents.rule_based_agent import RuleBasedAgent
from agents.rl_agent import RLAgent
from ui.timeline import create_timeline, create_inbox_summary


# ─── Global State ────────────────────────────────────────────────────────────

env = None
agent = None
episode_rewards = []
step_log = []
cumulative_reward = 0.0


# ─── Premium CSS ─────────────────────────────────────────────────────────────

CUSTOM_CSS = """
@import url('https://fonts.googleapis.com/css2?family=Inter:wght@300;400;500;600;700;800&display=swap');

/* ── Base ── */
.gradio-container {
    background: linear-gradient(135deg, #0a0a1a 0%, #1a1040 30%, #0d1b3e 60%, #0a0a1a 100%) !important;
    font-family: 'Inter', -apple-system, BlinkMacSystemFont, sans-serif !important;
    min-height: 100vh;
    --background-fill-primary: rgba(30, 30, 45, 0.4) !important;
    --background-fill-secondary: rgba(20, 20, 30, 0.6) !important;
    --border-color-primary: rgba(255, 255, 255, 0.1) !important;
    --block-background-fill: rgba(255, 255, 255, 0.03) !important;
/* ── Glass Cards ── */
.glass-card {
    background: rgba(255, 255, 255, 0.03) !important;
    border: 1px solid rgba(255, 255, 255, 0.08) !important;
    border-radius: 16px !important;
    backdrop-filter: blur(20px) !important;
    -webkit-backdrop-filter: blur(20px) !important;
    padding: 20px !important;
    transition: all 0.3s ease !important;
}
.glass-card:hover {
    border-color: rgba(139, 92, 246, 0.3) !important;
    background: rgba(255, 255, 255, 0.05) !important;
    box-shadow: 0 8px 32px rgba(139, 92, 246, 0.1) !important;
}

/* ── Headings ── */
h1, .title-text {
    background: linear-gradient(135deg, #a78bfa, #818cf8, #6366f1, #8b5cf6) !important;
    -webkit-background-clip: text !important;
    -webkit-text-fill-color: transparent !important;
    background-clip: text !important;
    font-weight: 800 !important;
    letter-spacing: -0.5px !important;
}

/* ── Buttons ── */
.primary-btn {
    background: linear-gradient(135deg, #7c3aed, #6366f1, #8b5cf6) !important;
    border: none !important;
    color: white !important;
    font-weight: 600 !important;
    font-size: 15px !important;
    padding: 12px 28px !important;
    border-radius: 12px !important;
    transition: all 0.3s ease !important;
    box-shadow: 0 4px 15px rgba(124, 58, 237, 0.3) !important;
    letter-spacing: 0.3px !important;
}
.primary-btn:hover {
    transform: translateY(-2px) !important;
    box-shadow: 0 8px 25px rgba(124, 58, 237, 0.5) !important;
}

.success-btn {
    background: linear-gradient(135deg, #059669, #10b981, #34d399) !important;
    border: none !important;
    color: white !important;
    font-weight: 600 !important;
    font-size: 15px !important;
    padding: 12px 28px !important;
    border-radius: 12px !important;
    transition: all 0.3s ease !important;
    box-shadow: 0 4px 15px rgba(5, 150, 105, 0.3) !important;
}
.success-btn:hover {
    transform: translateY(-2px) !important;
    box-shadow: 0 8px 25px rgba(5, 150, 105, 0.5) !important;
}

.danger-btn {
    background: linear-gradient(135deg, #dc2626, #ef4444, #f87171) !important;
    border: none !important;
    color: white !important;
    font-weight: 600 !important;
    font-size: 15px !important;
    padding: 12px 28px !important;
    border-radius: 12px !important;
    transition: all 0.3s ease !important;
    box-shadow: 0 4px 15px rgba(220, 38, 38, 0.3) !important;
}
.danger-btn:hover {
    transform: translateY(-2px) !important;
    box-shadow: 0 8px 25px rgba(220, 38, 38, 0.5) !important;
}

/* ── Dropdowns ── */
/* Handled by global dark mode */

/* ── Labels ── */
label, span[data-testid="block-info"] {
    color: #a5b4fc !important;
    font-weight: 600 !important;
    font-size: 13px !important;
    text-transform: uppercase !important;
    letter-spacing: 0.8px !important;
}

/* ── Markdown ── */
.markdown-text, .prose {
    color: #c7d2fe !important;
}
.prose h1, .prose h2, .prose h3 {
    color: #e0e7ff !important;
}
.prose table {
    border-collapse: separate !important;
    border-spacing: 0 !important;
    border-radius: 12px !important;
    overflow: hidden !important;
}
.prose th {
    background: rgba(139, 92, 246, 0.15) !important;
    color: #c4b5fd !important;
    padding: 10px 16px !important;
    font-size: 12px !important;
    text-transform: uppercase !important;
    letter-spacing: 0.5px !important;
    border-bottom: 1px solid rgba(139, 92, 246, 0.2) !important;
}
.prose td {
    background: rgba(255, 255, 255, 0.02) !important;
    color: #e0e7ff !important;
    padding: 10px 16px !important;
    border-bottom: 1px solid rgba(255, 255, 255, 0.05) !important;
}
.prose code {
    background: rgba(139, 92, 246, 0.15) !important;
    color: #c4b5fd !important;
    padding: 2px 8px !important;
    border-radius: 6px !important;
    font-size: 13px !important;
}

/* ── Code Block ── */
.code-wrap, .cm-editor {
    background: rgba(0, 0, 0, 0.3) !important;
    border: 1px solid rgba(255, 255, 255, 0.06) !important;
    border-radius: 12px !important;
}

/* ── Plot containers ── */
.plot-container {
    background: transparent !important;
    border: 1px solid rgba(255, 255, 255, 0.06) !important;
    border-radius: 16px !important;
    overflow: hidden !important;
}

/* ── Hide footer ── */
footer { display: none !important; }

/* ── Scrollbar ── */
::-webkit-scrollbar { width: 6px; }
::-webkit-scrollbar-track { background: rgba(0,0,0,0.2); }
::-webkit-scrollbar-thumb { background: rgba(139, 92, 246, 0.3); border-radius: 3px; }
::-webkit-scrollbar-thumb:hover { background: rgba(139, 92, 246, 0.5); }

/* ── Tab styling ── */
button[role="tab"] {
    background: transparent !important;
    border: none !important;
    border-bottom: 2px solid transparent !important;
    color: #94a3b8 !important;
    border-radius: 0 !important;
    font-weight: 600 !important;
    padding: 12px 24px !important;
    font-size: 14px !important;
    transition: all 0.3s ease !important;
}
button[role="tab"].selected {
    background: rgba(139, 92, 246, 0.05) !important;
    border-bottom: 2px solid #a78bfa !important;
    color: #e0e7ff !important;
}
button[role="tab"]:hover:not(.selected) {
    background: rgba(255, 255, 255, 0.02) !important;
    color: #c7d2fe !important;
}
div[role="tabpanel"] {
    background: rgba(255, 255, 255, 0.02) !important;
    border: 1px solid rgba(255, 255, 255, 0.06) !important;
    border-radius: 12px !important;
    padding: 12px !important;
    margin-top: 8px !important;
}
"""


# ─── Helper: Agent factory ───────────────────────────────────────────────────

def _get_agent(agent_type: str):
    agents = {
        "Random Agent": RandomAgent(),
        "Rule-Based Agent": RuleBasedAgent(),
        "Q-Learning Agent": RLAgent(epsilon=0.1),
    }
    return agents.get(agent_type, RandomAgent())


AGENT_DESCRIPTIONS = {
    "Random Agent": "Picks random valid actions. Performance lower-bound baseline.",
    "Rule-Based Agent": "Priority-first heuristic: high tasks > urgent emails > medium tasks.",
    "Q-Learning Agent": "Learns from experience via tabular Q-learning with exploration.",
}

DIFFICULTY_INFO = {
    "Easy": "3-5 tasks, few conflicts, no hidden items",
    "Medium": "5-8 tasks, some conflicts, hidden tasks appear mid-day",
    "Hard": "8-12 tasks, many conflicts, lots of surprises",
}


# ─── Core Functions ──────────────────────────────────────────────────────────

def initialize(agent_type: str, difficulty: str):
    global env, agent, episode_rewards, step_log, cumulative_reward

    env = ExecutiveAssistantEnv(difficulty=difficulty.lower())
    agent = _get_agent(agent_type)
    episode_rewards = []
    step_log = []
    cumulative_reward = 0.0

    state = env.reset()

    state_json = json.dumps(state, indent=2, default=str)

    timeline_fig = create_timeline(
        state.get("tasks", []),
        current_time=state.get("time", "08:00"),
    )
    inbox_fig = create_inbox_summary(state.get("inbox", []))

    dashboard = _build_dashboard(state, 0, 0.0)
    tasks_html = _build_tasks_html(state)
    inbox_html = _build_inbox_html(state)

    log_msg = _format_log(
        "SYSTEM", "Episode Initialized",
        f"Agent: {agent_type} | Difficulty: {difficulty} | "
        f"Tasks: {len(state.get('tasks', []))} | Messages: {len(state.get('inbox', []))}",
        "info"
    )

    return state_json, timeline_fig, inbox_fig, dashboard, tasks_html, inbox_html, log_msg


def take_step():
    global env, agent, episode_rewards, step_log, cumulative_reward

    if env is None:
        empty = _empty_outputs()
        return (*empty, _format_log("ERROR", "Not Initialized", "Click Initialize first.", "error"))

    state = env.get_state()
    action = agent.act(state)
    next_state, reward, done, info = env.step(action)
    cumulative_reward += reward
    episode_rewards.append(reward)
    step_log.append(info)

    state_json = json.dumps(next_state, indent=2, default=str)
    timeline_fig = create_timeline(
        next_state.get("tasks", []),
        current_time=next_state.get("time", "08:00"),
    )
    inbox_fig = create_inbox_summary(next_state.get("inbox", []))
    dashboard = _build_dashboard(next_state, len(step_log), cumulative_reward)
    tasks_html = _build_tasks_html(next_state)
    inbox_html = _build_inbox_html(next_state)

    if done:
        metrics = info.get("metrics", {})
        log_msg = _format_log(
            "COMPLETE", f"Episode Finished in {len(step_log)} Steps",
            f"Total Reward: {cumulative_reward:+.1f} | "
            f"Completion: {metrics.get('task_completion_rate', 0):.0%} | "
            f"Efficiency: {metrics.get('efficiency_score', 0):.0f}/100",
            "success"
        )
    else:
        icon = "+" if info.get("action_success") else "x"
        log_msg = _format_log(
            f"STEP {len(step_log)}",
            f"{info.get('action_type', '?')} -> target #{info.get('target_id', '?')}",
            f"[{icon}] {info.get('action_detail', '')} | Reward: {reward:+.2f} | Total: {cumulative_reward:+.1f}",
            "success" if info.get("action_success") else "warning"
        )

    return state_json, timeline_fig, inbox_fig, dashboard, tasks_html, inbox_html, log_msg


def run_full_episode():
    global env, agent, episode_rewards, step_log, cumulative_reward

    if env is None:
        empty = _empty_outputs()
        return (*empty, _format_log("ERROR", "Not Initialized", "Click Initialize first.", "error"))

    done = False
    while not done:
        state = env.get_state()
        action = agent.act(state)
        next_state, reward, done, info = env.step(action)
        cumulative_reward += reward
        episode_rewards.append(reward)
        step_log.append(info)

    final_state = env.get_state()
    state_json = json.dumps(final_state, indent=2, default=str)
    timeline_fig = create_timeline(
        final_state.get("tasks", []),
        current_time=final_state.get("time", "08:00"),
    )
    inbox_fig = create_inbox_summary(final_state.get("inbox", []))
    dashboard = _build_dashboard(final_state, len(step_log), cumulative_reward)
    tasks_html = _build_tasks_html(final_state)
    inbox_html = _build_inbox_html(final_state)

    metrics = info.get("metrics", {})
    log_msg = _format_log(
        "AUTO-RUN COMPLETE",
        f"Finished in {len(step_log)} Steps",
        f"Total Reward: {cumulative_reward:+.1f} | "
        f"Completion: {metrics.get('task_completion_rate', 0):.0%} | "
        f"Hi-Priority: {metrics.get('high_priority_completion', 0):.0%} | "
        f"Msg Response: {metrics.get('message_response_rate', 0):.0%} | "
        f"Efficiency: {metrics.get('efficiency_score', 0):.0f}/100 | "
        f"Conflicts: {metrics.get('conflict_count', 0)}",
        "success"
    )

    return state_json, timeline_fig, inbox_fig, dashboard, tasks_html, inbox_html, log_msg


def _empty_outputs():
    return (
        "{}",
        create_timeline([]),
        create_inbox_summary([]),
        _build_dashboard({}, 0, 0.0),
        "_No tasks yet._",
        "_No messages yet._",
    )


# ─── Formatters ──────────────────────────────────────────────────────────────

def _format_log(tag: str, title: str, detail: str, level: str = "info") -> str:
    colors = {
        "info": "#818cf8",
        "success": "#34d399",
        "warning": "#fbbf24",
        "error": "#f87171",
    }
    color = colors.get(level, "#818cf8")
    return (
        f"### <span style='color:{color}'>[{tag}]</span> {title}\n\n"
        f"<span style='color:#94a3b8'>{detail}</span>"
    )


def _build_dashboard(state: dict, step: int, reward: float) -> str:
    tasks = state.get("tasks", [])
    inbox = state.get("inbox", [])

    total = len(tasks)
    pending = sum(1 for t in tasks if t.get("status") == "pending")
    completed = sum(1 for t in tasks if t.get("status") == "completed")
    missed = sum(1 for t in tasks if t.get("status") == "missed")
    total_msgs = len(inbox)
    unreplied = sum(1 for m in inbox if not m.get("replied", False))
    replied = total_msgs - unreplied

    time_val = state.get("time", "--:--")

    # Build metrics cards
    return f"""
<div style="display:grid; grid-template-columns: repeat(4, 1fr); gap:12px; margin:8px 0;">
  <div style="background:rgba(99,102,241,0.12); border:1px solid rgba(99,102,241,0.25); border-radius:14px; padding:16px; text-align:center;">
    <div style="font-size:12px; color:#a5b4fc; font-weight:600; text-transform:uppercase; letter-spacing:1px;">Time</div>
    <div style="font-size:28px; font-weight:800; color:#c7d2fe; margin-top:4px;">{time_val}</div>
  </div>
  <div style="background:rgba(16,185,129,0.12); border:1px solid rgba(16,185,129,0.25); border-radius:14px; padding:16px; text-align:center;">
    <div style="font-size:12px; color:#6ee7b7; font-weight:600; text-transform:uppercase; letter-spacing:1px;">Step</div>
    <div style="font-size:28px; font-weight:800; color:#a7f3d0; margin-top:4px;">{step}</div>
  </div>
  <div style="background:rgba(251,191,36,0.12); border:1px solid rgba(251,191,36,0.25); border-radius:14px; padding:16px; text-align:center;">
    <div style="font-size:12px; color:#fcd34d; font-weight:600; text-transform:uppercase; letter-spacing:1px;">Reward</div>
    <div style="font-size:28px; font-weight:800; color:#fef08a; margin-top:4px;">{reward:+.1f}</div>
  </div>
  <div style="background:rgba(244,63,94,0.12); border:1px solid rgba(244,63,94,0.25); border-radius:14px; padding:16px; text-align:center;">
    <div style="font-size:12px; color:#fda4af; font-weight:600; text-transform:uppercase; letter-spacing:1px;">Difficulty</div>
    <div style="font-size:28px; font-weight:800; color:#fecdd3; margin-top:4px;">{state.get('difficulty', 'N/A').upper() if 'difficulty' in state else (env.difficulty.upper() if env else 'N/A')}</div>
  </div>
</div>

<div style="display:grid; grid-template-columns: repeat(3, 1fr); gap:12px; margin-top:12px;">
  <div style="background:rgba(255,255,255,0.03); border:1px solid rgba(255,255,255,0.06); border-radius:12px; padding:14px;">
    <div style="font-size:11px; color:#94a3b8; font-weight:600; text-transform:uppercase; letter-spacing:1px;">Tasks</div>
    <div style="display:flex; justify-content:space-between; margin-top:10px;">
      <span style="color:#fbbf24; font-weight:700;">{pending} pending</span>
      <span style="color:#34d399; font-weight:700;">{completed} done</span>
      <span style="color:#f87171; font-weight:700;">{missed} missed</span>
    </div>
    <div style="background:rgba(255,255,255,0.06); border-radius:8px; height:6px; margin-top:10px; overflow:hidden;">
      <div style="background:linear-gradient(90deg,#34d399,#10b981); height:100%; width:{int(completed/max(total,1)*100)}%; border-radius:8px;"></div>
    </div>
  </div>
  <div style="background:rgba(255,255,255,0.03); border:1px solid rgba(255,255,255,0.06); border-radius:12px; padding:14px;">
    <div style="font-size:11px; color:#94a3b8; font-weight:600; text-transform:uppercase; letter-spacing:1px;">Inbox</div>
    <div style="display:flex; justify-content:space-between; margin-top:10px;">
      <span style="color:#818cf8; font-weight:700;">{total_msgs} total</span>
      <span style="color:#34d399; font-weight:700;">{replied} replied</span>
      <span style="color:#f87171; font-weight:700;">{unreplied} waiting</span>
    </div>
    <div style="background:rgba(255,255,255,0.06); border-radius:8px; height:6px; margin-top:10px; overflow:hidden;">
      <div style="background:linear-gradient(90deg,#818cf8,#6366f1); height:100%; width:{int(replied/max(total_msgs,1)*100)}%; border-radius:8px;"></div>
    </div>
  </div>
  <div style="background:rgba(255,255,255,0.03); border:1px solid rgba(255,255,255,0.06); border-radius:12px; padding:14px;">
    <div style="font-size:11px; color:#94a3b8; font-weight:600; text-transform:uppercase; letter-spacing:1px;">Efficiency</div>
    <div style="font-size:24px; font-weight:800; color:#c7d2fe; margin-top:6px; text-align:center;">{int(completed/max(total,1)*100)}%</div>
    <div style="background:rgba(255,255,255,0.06); border-radius:8px; height:6px; margin-top:10px; overflow:hidden;">
      <div style="background:linear-gradient(90deg,#fbbf24,#f59e0b); height:100%; width:{int(completed/max(total,1)*100)}%; border-radius:8px;"></div>
    </div>
  </div>
</div>
"""


def _build_tasks_html(state: dict) -> str:
    tasks = state.get("tasks", [])
    if not tasks:
        return "_No tasks yet._"

    priority_colors = {"high": "#f87171", "medium": "#fbbf24", "low": "#60a5fa"}
    status_colors = {
        "pending": "#fbbf24", "completed": "#34d399", "scheduled": "#a78bfa",
        "missed": "#f87171", "deferred": "#94a3b8", "rejected": "#6b7280",
    }
    type_icons = {"meeting": "VIDEO", "work": "CODE", "personal": "USER"}

    rows = ""
    for t in sorted(tasks, key=lambda x: x.get("time", "")):
        p_color = priority_colors.get(t.get("priority", ""), "#94a3b8")
        s_color = status_colors.get(t.get("status", ""), "#94a3b8")
        icon = type_icons.get(t.get("type", ""), "DOT")
        dur = t.get("duration", 30)

        rows += f"""
    <div style="display:flex; align-items:center; padding:10px 14px; background:rgba(255,255,255,0.02);
                border-radius:10px; margin-bottom:6px; border-left:3px solid {p_color};
                transition:all 0.2s ease;">
      <div style="min-width:55px; color:#94a3b8; font-weight:600; font-size:13px; font-family:monospace;">{t.get('time','')}</div>
      <div style="min-width:40px; color:#64748b; font-size:10px; font-weight:600;">{dur}m</div>
      <div style="flex:1; color:#e0e7ff; font-weight:500; font-size:13px;">{t.get('title','Task')}</div>
      <div style="min-width:50px; text-align:center;">
        <span style="background:{p_color}22; color:{p_color}; padding:2px 10px; border-radius:10px;
                     font-size:10px; font-weight:700; text-transform:uppercase;">{t.get('priority','')}</span>
      </div>
      <div style="min-width:70px; text-align:center;">
        <span style="background:{s_color}22; color:{s_color}; padding:2px 10px; border-radius:10px;
                     font-size:10px; font-weight:700; text-transform:uppercase;">{t.get('status','')}</span>
      </div>
    </div>"""

    return f"""<div style="max-height:380px; overflow-y:auto; padding-right:4px;">{rows}</div>"""


def _build_inbox_html(state: dict) -> str:
    inbox = state.get("inbox", [])
    if not inbox:
        return "_No messages yet._"

    rows = ""
    for m in inbox:
        urgency = m.get("urgency", "low")
        replied = m.get("replied", False)
        u_color = "#f87171" if urgency == "high" else ("#fbbf24" if urgency == "medium" else "#60a5fa")
        r_color = "#34d399" if replied else "#64748b"
        r_text = "REPLIED" if replied else "PENDING"

        rows += f"""
    <div style="display:flex; align-items:center; padding:10px 14px; background:rgba(255,255,255,0.02);
                border-radius:10px; margin-bottom:6px; border-left:3px solid {u_color};
                transition:all 0.2s ease;">
      <div style="min-width:120px; color:#c7d2fe; font-weight:600; font-size:13px;">{m.get('sender','Unknown')}</div>
      <div style="flex:1; color:#94a3b8; font-size:12px; overflow:hidden; white-space:nowrap; text-overflow:ellipsis;">{m.get('content','')}</div>
      <div style="min-width:50px; text-align:center;">
        <span style="background:{u_color}22; color:{u_color}; padding:2px 10px; border-radius:10px;
                     font-size:10px; font-weight:700; text-transform:uppercase;">{urgency}</span>
      </div>
      <div style="min-width:65px; text-align:center;">
        <span style="background:{r_color}22; color:{r_color}; padding:2px 10px; border-radius:10px;
                     font-size:10px; font-weight:700;">{r_text}</span>
      </div>
    </div>"""

    return f"""<div style="max-height:300px; overflow-y:auto; padding-right:4px;">{rows}</div>"""


# ─── Build UI ────────────────────────────────────────────────────────────────

def build_ui():
    with gr.Blocks(title="AI Executive Assistant Simulator") as demo:

        # ── Header ──
        gr.HTML("""
        <div style="text-align:center; padding:30px 20px 10px;">
            <div style="font-size:48px; margin-bottom:8px;">🤖</div>
            <h1 style="font-size:36px; font-weight:800; margin:0;
                        background:linear-gradient(135deg,#a78bfa,#818cf8,#6366f1);
                        -webkit-background-clip:text; -webkit-text-fill-color:transparent;">
                AI Executive Assistant
            </h1>
            <p style="color:#94a3b8; font-size:15px; margin-top:6px; font-weight:400;">
                OpenEnv RL Environment &mdash; Watch AI agents learn to manage schedules, emails & priorities
            </p>
        </div>
        """)

        # ── Control Panel ──
        with gr.Row(equal_height=True):
            with gr.Column(scale=1, min_width=300, elem_classes=["glass-card"]):
                gr.HTML("""
                <div style="display:flex; align-items:center; margin-bottom:16px;">
                    <div style="background:rgba(139,92,246,0.2); padding:8px; border-radius:8px; margin-right:12px;">
                        ⚙️
                    </div>
                    <div style="font-size:14px; color:#c4b5fd; font-weight:700; text-transform:uppercase; letter-spacing:1.5px;">
                        Control Panel
                    </div>
                </div>
                """)

                agent_selector = gr.Dropdown(
                    choices=["Random Agent", "Rule-Based Agent", "Q-Learning Agent"],
                    value="Rule-Based Agent",
                    label="🤖 Agent Type",
                    info="Select which AI agent controls the assistant",
                )
                
                gr.HTML("<div style='height:4px'></div>")
                
                difficulty_selector = gr.Dropdown(
                    choices=["Easy", "Medium", "Hard"],
                    value="Medium",
                    label="🎯 Difficulty",
                    info="Controls task count, conflicts, and hidden items",
                )

                gr.HTML("<div style='height:8px'></div>")

                init_btn = gr.Button(
                    "Initialize New Episode",
                    variant="primary",
                    elem_classes=["primary-btn"],
                )
                with gr.Row():
                    step_btn = gr.Button(
                        "Step Forward",
                        variant="secondary",
                        elem_classes=["success-btn"],
                    )
                    run_btn = gr.Button(
                        "Run All Steps",
                        variant="stop",
                        elem_classes=["danger-btn"],
                    )

            with gr.Column(scale=3):
                action_log = gr.Markdown(
                    value=_format_log(
                        "WELCOME",
                        "Ready to simulate",
                        "Select an agent and difficulty, then click Initialize to begin.",
                        "info"
                    ),
                )
                dashboard = gr.HTML(
                    value=_build_dashboard({}, 0, 0.0),
                )

        # ── Main Content Tabs ──
        with gr.Tabs():
            with gr.Tab("Schedule Timeline"):
                timeline_plot = gr.Plot(label="Schedule Timeline", show_label=False)

            with gr.Tab("Task List"):
                tasks_display = gr.HTML(value="_No tasks yet._")

            with gr.Tab("Inbox"):
                with gr.Row():
                    with gr.Column(scale=2):
                        inbox_display = gr.HTML(value="_No messages yet._")
                    with gr.Column(scale=1):
                        inbox_plot = gr.Plot(label="Inbox Summary", show_label=False)

            with gr.Tab("Raw State (JSON)"):
                state_display = gr.Code(
                    value="{}",
                    language="json",
                    label="Observation Space",
                    lines=20,
                )

        # ── Footer ──
        gr.HTML("""
        <div style="text-align:center; padding:24px 0 12px; border-top:1px solid rgba(255,255,255,0.05); margin-top:24px;">
            <span style="color:#4b5563; font-size:12px;">
                Built with OpenEnv + Gradio + Plotly &nbsp;|&nbsp;
                Reinforcement Learning Executive Assistant Simulator
            </span>
        </div>
        """)

        # ── Event Handlers ──
        outputs = [state_display, timeline_plot, inbox_plot, dashboard, tasks_display, inbox_display, action_log]

        init_btn.click(fn=initialize, inputs=[agent_selector, difficulty_selector], outputs=outputs)
        step_btn.click(fn=take_step, inputs=[], outputs=outputs)
        run_btn.click(fn=run_full_episode, inputs=[], outputs=outputs)

    return demo


# ─── Entry Point ─────────────────────────────────────────────────────────────

if __name__ == "__main__":
    demo = build_ui()
    my_theme = gr.themes.Soft(
        primary_hue="violet",
        secondary_hue="indigo",
        neutral_hue="slate",
    ).set(
        body_background_fill="*background_fill_primary",
        block_background_fill="rgba(255, 255, 255, 0.03)",
        block_border_color="rgba(255, 255, 255, 0.08)",
        input_background_fill="rgba(255, 255, 255, 0.05)",
        input_background_fill_hover="rgba(255, 255, 255, 0.08)",
        border_color_primary="rgba(255, 255, 255, 0.1)",
        body_text_color="#e0e7ff",
        color_accent_soft="rgba(139, 92, 246, 0.15)",
        checkbox_background_color="rgba(255, 255, 255, 0.05)",
        slider_color="*primary_600",
        shadow_drop="none",
        border_color_accent="rgba(139, 92, 246, 0.5)",
    )

    demo.launch(
        server_name="0.0.0.0",
        share=False,
        theme=my_theme,
        css=CUSTOM_CSS,
    )