server/app.py

"""
FastAPI application for the AETHER-TaskFlow Environment.

Endpoints:
    GET  /        - Interactive Mission Control Dashboard
    POST /reset   - Reset the environment, return initial observation
    POST /step    - Execute an action, return next observation
    GET  /state   - Return current internal state
    GET  /schema  - Return action/observation/state JSON schemas
    GET  /health  - Health check
    WS   /ws      - WebSocket for persistent sessions
    GET  /docs    - Swagger UI
"""

from __future__ import annotations

import os
import sys
from pathlib import Path
from typing import Any, Dict

from fastapi import Body, HTTPException, status
from fastapi.responses import HTMLResponse
from openenv.core.env_server.http_server import create_app
from openenv.core.env_server.types import ResetRequest, ResetResponse, SchemaResponse, StepResponse
from openenv.core.env_server.web_interface import WebInterfaceManager
from pydantic import ValidationError

_REPO_ROOT = Path(__file__).resolve().parents[1]
sys.path.insert(0, str(_REPO_ROOT))

try:
    from models import AetherTaskFlowAction, AetherTaskFlowObservation, AetherTaskFlowState
    from env.aether_env import AetherTaskFlowEnvironment
except ModuleNotFoundError:
    sys.path.insert(0, str(_REPO_ROOT))
    from models import AetherTaskFlowAction, AetherTaskFlowObservation, AetherTaskFlowState
    from env.aether_env import AetherTaskFlowEnvironment

_DIFFICULTY = os.getenv("AETHER_DIFFICULTY", "easy")


def _make_env() -> AetherTaskFlowEnvironment:
    difficulty = os.getenv("AETHER_DIFFICULTY", _DIFFICULTY)
    return AetherTaskFlowEnvironment(difficulty=difficulty)


def _create_persistent_manager() -> WebInterfaceManager:
    temp_env = _make_env()
    metadata: dict = {}
    try:
        if hasattr(temp_env, "get_metadata"):
            metadata = temp_env.get_metadata()
    except Exception:
        pass
    finally:
        if hasattr(temp_env, "close"):
            try:
                temp_env.close()
            except Exception:
                pass
    return WebInterfaceManager(_make_env, AetherTaskFlowAction, AetherTaskFlowObservation, metadata=metadata)


app = create_app(
    _make_env, AetherTaskFlowAction, AetherTaskFlowObservation,
    env_name="aether_taskflow", max_concurrent_envs=4,
)

_persistent_manager = _create_persistent_manager()


def _remove_route(path: str, method: str) -> None:
    app.router.routes = [
        route for route in app.router.routes
        if not (getattr(route, "path", None) == path and method in (getattr(route, "methods", set()) or set()))
    ]


for _path, _method in (("/reset", "POST"), ("/step", "POST"), ("/state", "GET"), ("/schema", "GET")):
    _remove_route(_path, _method)


def _action_to_payload(action: Any) -> Dict[str, Any]:
    if hasattr(action, "model_dump"):
        return action.model_dump(exclude={"metadata"})
    if isinstance(action, dict):
        return action
    raise HTTPException(status_code=status.HTTP_422_UNPROCESSABLE_ENTITY, detail="Invalid action")


def _extract_action_payload(payload: Dict[str, Any]) -> Dict[str, Any]:
    env = _persistent_manager.env
    if "message" in payload and isinstance(payload["message"], str):
        if hasattr(env, "message_to_action"):
            return _action_to_payload(env.message_to_action(payload["message"]))
        return {"message": payload["message"]}
    action_payload = payload.get("action", payload)
    if isinstance(action_payload, str):
        if hasattr(env, "message_to_action"):
            return _action_to_payload(env.message_to_action(action_payload))
        return {"message": action_payload}
    if isinstance(action_payload, dict):
        if "message" in action_payload and isinstance(action_payload["message"], str):
            if hasattr(env, "message_to_action"):
                return _action_to_payload(env.message_to_action(action_payload["message"]))
            return {"message": action_payload["message"]}
        return action_payload
    raise HTTPException(status_code=status.HTTP_422_UNPROCESSABLE_ENTITY, detail="Invalid payload")



@app.get("/", response_class=HTMLResponse, tags=["UI"], summary="Mission Control Dashboard")
async def root() -> HTMLResponse:
    difficulty = os.getenv("AETHER_DIFFICULTY", "easy")
    html = f"""<!DOCTYPE html>
<html lang="en">
<head>
  <meta charset="UTF-8"/>
  <meta name="viewport" content="width=device-width, initial-scale=1.0"/>
  <title>AETHER-TaskFlow · Mission Control</title>
  <link rel="preconnect" href="https://fonts.googleapis.com"/>
  <link href="https://fonts.googleapis.com/css2?family=Space+Mono:ital,wght@0,400;0,700;1,400&family=Syne:wght@400;600;700;800&display=swap" rel="stylesheet"/>
  <style>
    :root {{
      --bg:       #050810;
      --panel:    #0b0f1e;
      --border:   #1a2040;
      --border2:  #252d50;
      --accent:   #3d7eff;
      --accent2:  #6b3dff;
      --cyan:     #00d4ff;
      --green:    #00e676;
      --amber:    #ffab00;
      --red:      #ff3d3d;
      --text:     #cdd6f4;
      --muted:    #6c7a9c;
      --glow:     rgba(61,126,255,0.18);
    }}
    *, *::before, *::after {{ box-sizing: border-box; margin: 0; padding: 0; }}

    body {{
      font-family: 'Syne', sans-serif;
      background: var(--bg);
      color: var(--text);
      min-height: 100vh;
      overflow-x: hidden;
    }}

    /* Grid scanline overlay */
    body::before {{
      content: '';
      position: fixed; inset: 0;
      background:
        repeating-linear-gradient(0deg, transparent, transparent 39px, rgba(61,126,255,0.03) 40px),
        repeating-linear-gradient(90deg, transparent, transparent 39px, rgba(61,126,255,0.03) 40px);
      pointer-events: none; z-index: 0;
    }}

    /* ── TOP BAR ── */
    header {{
      position: relative; z-index: 10;
      display: flex; align-items: center; justify-content: space-between;
      padding: 0 28px;
      height: 60px;
      background: rgba(11,15,30,0.95);
      border-bottom: 1px solid var(--border2);
      backdrop-filter: blur(12px);
    }}
    .logo {{
      display: flex; align-items: center; gap: 12px;
    }}
    .logo-icon {{
      width: 32px; height: 32px;
      background: linear-gradient(135deg, var(--accent), var(--accent2));
      border-radius: 8px;
      display: flex; align-items: center; justify-content: center;
      font-size: 16px;
      box-shadow: 0 0 20px rgba(61,126,255,0.4);
    }}
    .logo-text {{ font-size: 1rem; font-weight: 800; letter-spacing: -0.02em; }}
    .logo-sub {{ font-size: 0.7rem; color: var(--muted); font-family: 'Space Mono', monospace; margin-top: 1px; }}
    .header-links {{ display: flex; gap: 6px; }}
    .hlink {{
      font-family: 'Space Mono', monospace;
      font-size: 0.68rem; color: var(--muted);
      text-decoration: none; padding: 5px 10px;
      border: 1px solid var(--border2); border-radius: 4px;
      transition: all .15s;
    }}
    .hlink:hover {{ color: var(--cyan); border-color: var(--cyan); }}
    .live-indicator {{
      display: flex; align-items: center; gap: 6px;
      font-family: 'Space Mono', monospace; font-size: 0.68rem; color: var(--green);
    }}
    .live-dot {{
      width: 6px; height: 6px; border-radius: 50%;
      background: var(--green);
      box-shadow: 0 0 6px var(--green);
      animation: pulse-dot 1.8s ease-in-out infinite;
    }}
    @keyframes pulse-dot {{
      0%,100% {{ opacity: 1; transform: scale(1); }}
      50% {{ opacity: 0.5; transform: scale(0.7); }}
    }}

    /* ── LAYOUT ── */
    .workspace {{
      position: relative; z-index: 1;
      display: grid;
      grid-template-columns: 300px 1fr 260px;
      grid-template-rows: auto 1fr auto;
      gap: 0;
      height: calc(100vh - 60px);
      overflow: hidden;
    }}

    /* ── COMMAND BAR ── */
    .cmd-bar {{
      grid-column: 1 / -1;
      display: flex; align-items: center; gap: 10px;
      padding: 10px 20px;
      background: rgba(11,15,30,0.8);
      border-bottom: 1px solid var(--border);
    }}
    .cmd-label {{
      font-family: 'Space Mono', monospace;
      font-size: 0.65rem; color: var(--muted);
      text-transform: uppercase; letter-spacing: .1em;
      white-space: nowrap;
    }}
    select.diff-select {{
      background: var(--panel); border: 1px solid var(--border2);
      color: var(--text); padding: 7px 14px; border-radius: 6px;
      font-family: 'Space Mono', monospace; font-size: 0.75rem;
      cursor: pointer; outline: none;
      transition: border-color .15s;
    }}
    select.diff-select:focus {{ border-color: var(--accent); }}
    .cmd-btn {{
      padding: 7px 18px; border: none; border-radius: 6px;
      font-family: 'Syne', sans-serif; font-size: 0.78rem; font-weight: 700;
      cursor: pointer; transition: all .15s; white-space: nowrap;
    }}
    .cmd-btn:hover {{ transform: translateY(-1px); filter: brightness(1.1); }}
    .cmd-btn:active {{ transform: translateY(0); }}
    .cmd-btn:disabled {{ opacity: 0.35; cursor: not-allowed; transform: none; }}
    .btn-reset {{ background: linear-gradient(135deg, var(--accent), var(--accent2)); color: white; }}
    .btn-step  {{ background: var(--green); color: #000; }}
    .btn-auto  {{ background: var(--amber); color: #000; }}
    .btn-stop  {{ background: var(--red); color: white; }}
    .status-pill {{
      margin-left: auto;
      font-family: 'Space Mono', monospace; font-size: 0.68rem;
      padding: 4px 12px; border-radius: 20px;
      border: 1px solid var(--border2); color: var(--muted);
      transition: all .3s;
    }}
    .status-pill.ready {{ border-color: var(--green); color: var(--green); }}
    .status-pill.done  {{ border-color: var(--amber); color: var(--amber); }}

    /* ── LEFT PANEL: Task Queue ── */
    .panel-left {{
      grid-row: 2;
      background: var(--panel);
      border-right: 1px solid var(--border);
      display: flex; flex-direction: column;
      overflow: hidden;
    }}
    .panel-title {{
      padding: 12px 16px 8px;
      font-size: 0.65rem; font-weight: 700;
      text-transform: uppercase; letter-spacing: .12em;
      color: var(--muted);
      border-bottom: 1px solid var(--border);
      display: flex; align-items: center; justify-content: space-between;
    }}
    .task-scroll {{
      flex: 1; overflow-y: auto; padding: 10px;
      scrollbar-width: thin; scrollbar-color: var(--border2) transparent;
    }}
    .task-card {{
      border: 1px solid var(--border2);
      border-radius: 8px; padding: 10px 12px;
      margin-bottom: 8px; cursor: pointer;
      transition: all .15s;
      position: relative; overflow: hidden;
    }}
    .task-card::before {{
      content: '';
      position: absolute; left: 0; top: 0; bottom: 0;
      width: 3px;
      background: var(--muted);
      transition: background .15s;
    }}
    .task-card:hover {{ border-color: var(--accent); background: rgba(61,126,255,0.05); }}
    .task-card:hover::before {{ background: var(--accent); }}
    .task-card.selected {{ border-color: var(--accent); background: rgba(61,126,255,0.08); box-shadow: 0 0 0 1px rgba(61,126,255,0.2); }}
    .task-card.selected::before {{ background: var(--accent); box-shadow: 0 0 8px var(--accent); }}
    .task-card.urgent {{ border-color: rgba(255,61,61,0.4); }}
    .task-card.urgent::before {{ background: var(--red); }}
    .task-name {{ font-size: 0.78rem; font-weight: 600; margin-bottom: 6px; line-height: 1.3; padding-left: 6px; }}
    .task-chips {{ display: flex; flex-wrap: wrap; gap: 4px; padding-left: 6px; }}
    .chip {{
      font-family: 'Space Mono', monospace;
      font-size: 0.6rem; padding: 2px 5px;
      border-radius: 3px; border: 1px solid var(--border2);
      color: var(--muted);
    }}
    .chip.hi  {{ color: var(--green);  border-color: rgba(0,230,118,0.3); }}
    .chip.mid {{ color: var(--amber);  border-color: rgba(255,171,0,0.3); }}
    .chip.lo  {{ color: var(--red);    border-color: rgba(255,61,61,0.3); }}
    .chip.val {{ color: var(--cyan);   border-color: rgba(0,212,255,0.3); }}
    .task-empty {{
      text-align: center; padding: 40px 16px;
      color: var(--muted); font-size: 0.8rem; line-height: 1.6;
    }}

    /* ── CENTRE: Visualisation ── */
    .panel-centre {{
      grid-row: 2;
      display: flex; flex-direction: column;
      overflow: hidden;
      border-right: 1px solid var(--border);
    }}

    /* Resource meters */
    .resource-row {{
      display: grid; grid-template-columns: repeat(4, 1fr);
      gap: 10px; padding: 14px 16px;
      border-bottom: 1px solid var(--border);
      background: rgba(11,15,30,0.5);
    }}
    .res-cell {{
      background: var(--panel);
      border: 1px solid var(--border2);
      border-radius: 8px; padding: 10px 14px;
      position: relative; overflow: hidden;
    }}
    .res-cell::after {{
      content: '';
      position: absolute; bottom: 0; left: 0;
      height: 2px; width: var(--fill, 0%);
      background: var(--fill-color, var(--accent));
      transition: width .5s ease, background .5s ease;
    }}
    .res-label {{
      font-family: 'Space Mono', monospace;
      font-size: 0.6rem; color: var(--muted);
      text-transform: uppercase; letter-spacing: .08em;
    }}
    .res-value {{
      font-family: 'Space Mono', monospace;
      font-size: 1.25rem; font-weight: 700;
      color: var(--text); margin: 4px 0 0;
    }}

    /* Action picker */
    .action-panel {{
      padding: 12px 16px;
      border-bottom: 1px solid var(--border);
      background: rgba(11,15,30,0.3);
    }}
    .action-grid {{ display: grid; grid-template-columns: repeat(4,1fr); gap: 8px; margin-top: 8px; }}
    .act-btn {{
      border: 1px solid var(--border2);
      border-radius: 8px; padding: 10px 6px;
      cursor: pointer; text-align: center;
      transition: all .15s; background: var(--panel);
      color: var(--muted); font-family: 'Syne', sans-serif;
    }}
    .act-btn:hover {{ border-color: var(--accent); color: var(--text); background: rgba(61,126,255,0.07); }}
    .act-btn.active {{
      border-color: var(--accent); color: var(--accent);
      background: rgba(61,126,255,0.12);
      box-shadow: 0 0 12px rgba(61,126,255,0.15);
    }}
    .act-btn .act-icon {{ font-size: 1.2rem; display: block; margin-bottom: 3px; }}
    .act-btn .act-name {{ font-size: 0.72rem; font-weight: 700; }}
    .act-btn .act-sub  {{ font-size: 0.6rem; color: var(--muted); margin-top: 1px; font-family: 'Space Mono', monospace; }}

    /* Outcome text */
    .outcome-bar {{
      padding: 10px 16px;
      border-bottom: 1px solid var(--border);
      font-family: 'Space Mono', monospace;
      font-size: 0.72rem; color: var(--muted);
      background: rgba(5,8,16,0.5);
      min-height: 38px;
    }}
    .outcome-bar span {{ color: var(--cyan); }}

    /* Reward chart area */
    .chart-area {{
      flex: 1; padding: 14px 16px;
      display: flex; flex-direction: column; gap: 10px;
    }}
    .chart-title {{
      font-size: 0.62rem; font-weight: 700;
      text-transform: uppercase; letter-spacing: .1em;
      color: var(--muted);
    }}
    .reward-chart {{
      flex: 1;
      display: flex; align-items: flex-end;
      gap: 4px; border-bottom: 1px solid var(--border2);
      padding-bottom: 4px; min-height: 80px;
    }}
    .reward-bar {{
      flex: 1; min-width: 8px; border-radius: 3px 3px 0 0;
      transition: height .4s ease, background .4s;
      position: relative;
    }}
    .reward-bar::after {{
      content: attr(data-val);
      position: absolute; bottom: calc(100% + 2px); left: 50%;
      transform: translateX(-50%);
      font-family: 'Space Mono', monospace; font-size: 0.5rem;
      color: var(--muted); white-space: nowrap;
      opacity: 0; transition: opacity .2s;
    }}
    .reward-bar:hover::after {{ opacity: 1; }}

    /* ── RIGHT PANEL: Telemetry ── */
    .panel-right {{
      grid-row: 2;
      background: var(--panel);
      display: flex; flex-direction: column;
      overflow-y: auto;
      scrollbar-width: thin; scrollbar-color: var(--border2) transparent;
    }}
    .score-hero {{
      padding: 20px 16px 14px;
      text-align: center;
      border-bottom: 1px solid var(--border);
      position: relative;
    }}
    .score-ring {{
      width: 110px; height: 110px;
      margin: 0 auto 10px;
      position: relative;
    }}
    .score-ring svg {{ transform: rotate(-90deg); }}
    .score-ring .ring-bg {{ fill: none; stroke: var(--border2); stroke-width: 8; }}
    .score-ring .ring-fill {{
      fill: none; stroke-width: 8;
      stroke-linecap: round;
      stroke-dasharray: 283;
      stroke-dashoffset: 283;
      transition: stroke-dashoffset .8s cubic-bezier(.4,0,.2,1), stroke .4s;
      stroke: var(--accent);
    }}
    .score-num {{
      position: absolute; top: 50%; left: 50%;
      transform: translate(-50%,-50%);
      font-family: 'Space Mono', monospace;
      font-size: 1.5rem; font-weight: 700;
    }}
    .score-label {{ font-size: 0.62rem; color: var(--muted); text-transform: uppercase; letter-spacing: .1em; }}

    .stat-grid {{
      display: grid; grid-template-columns: 1fr 1fr;
      gap: 8px; padding: 12px;
      border-bottom: 1px solid var(--border);
    }}
    .stat-box {{
      background: rgba(5,8,16,0.6);
      border: 1px solid var(--border2);
      border-radius: 8px; padding: 10px;
      text-align: center;
    }}
    .stat-box .sn {{
      font-family: 'Space Mono', monospace;
      font-size: 1.3rem; font-weight: 700;
    }}
    .stat-box .sl {{ font-size: 0.6rem; color: var(--muted); margin-top: 2px; text-transform: uppercase; letter-spacing: .07em; }}
    .stat-box.green .sn {{ color: var(--green); }}
    .stat-box.red   .sn {{ color: var(--red);   }}
    .stat-box.cyan  .sn {{ color: var(--cyan);  }}
    .stat-box.amber .sn {{ color: var(--amber); }}

    /* Health bar */
    .health-section {{ padding: 12px; border-bottom: 1px solid var(--border); }}
    .health-label {{ font-size: 0.62rem; color: var(--muted); text-transform: uppercase; letter-spacing: .1em; margin-bottom: 8px; display: flex; justify-content: space-between; }}
    .health-track {{
      height: 8px; background: var(--border2); border-radius: 4px; overflow: hidden;
    }}
    .health-fill {{
      height: 100%; border-radius: 4px;
      transition: width .5s ease, background .5s ease;
      background: var(--green);
    }}

    /* Log */
    .log-panel {{ flex: 1; padding: 10px; }}
    .log-title {{ font-size: 0.62rem; font-weight: 700; text-transform: uppercase; letter-spacing: .1em; color: var(--muted); margin-bottom: 8px; }}
    .log-scroll {{
      height: 200px; overflow-y: auto;
      scrollbar-width: thin; scrollbar-color: var(--border2) transparent;
    }}
    .log-entry {{
      font-family: 'Space Mono', monospace;
      font-size: 0.62rem; line-height: 1.7;
      padding: 1px 0; border-bottom: 1px solid rgba(26,32,64,0.4);
      color: var(--muted);
    }}
    .log-entry.pos {{ color: var(--green); }}
    .log-entry.neg {{ color: var(--red);   }}
    .log-entry.inf {{ color: var(--cyan);  }}
    .log-entry.wrn {{ color: var(--amber); }}

    /* ── STATUS BAR ── */
    .status-bar {{
      grid-column: 1 / -1;
      display: flex; align-items: center; gap: 16px;
      padding: 6px 20px;
      background: rgba(11,15,30,0.95);
      border-top: 1px solid var(--border);
      font-family: 'Space Mono', monospace;
      font-size: 0.62rem; color: var(--muted);
    }}
    .sb-item {{ display: flex; align-items: center; gap: 6px; }}
    .sb-dot {{ width: 5px; height: 5px; border-radius: 50%; background: var(--muted); }}
    .sb-dot.ok  {{ background: var(--green); box-shadow: 0 0 5px var(--green); }}
    .sb-dot.bad {{ background: var(--red); }}

    /* Spinner */
    .spinner {{
      display: inline-block; width: 12px; height: 12px;
      border: 2px solid rgba(255,255,255,0.2);
      border-top-color: white; border-radius: 50%;
      animation: spin .6s linear infinite;
      vertical-align: middle; margin-right: 5px;
    }}
    @keyframes spin {{ to {{ transform: rotate(360deg); }} }}

    /* DONE FLASH */
    @keyframes flash-border {{
      0%,100% {{ border-color: var(--border2); }}
      50% {{ border-color: var(--amber); box-shadow: 0 0 20px rgba(255,171,0,0.3); }}
    }}
    .episode-done {{ animation: flash-border 1.5s ease 2; }}
  </style>
</head>
<body>

<!-- ── TOP BAR ── -->
<header>
  <div class="logo">
    <div class="logo-icon">⚡</div>
    <div>
      <div class="logo-text">AETHER-TaskFlow</div>
      <div class="logo-sub">OPENENV · RL ENVIRONMENT · v1.0.0</div>
    </div>
  </div>
  <div class="live-indicator"><div class="live-dot"></div>LIVE</div>
  <div class="header-links">
    <a class="hlink" href="/docs">Swagger</a>
    <a class="hlink" href="/schema">Schema</a>
    <a class="hlink" href="/state">State</a>
    <a class="hlink" href="/health">Health</a>
  </div>
</header>

<!-- ── WORKSPACE ── -->
<div class="workspace">

  <!-- COMMAND BAR -->
  <div class="cmd-bar">
    <span class="cmd-label">Difficulty</span>
    <select class="diff-select" id="difficulty">
      <option value="easy"   {"selected" if difficulty=="easy"   else ""}>🟢 EASY — 5 tasks, stable</option>
      <option value="medium" {"selected" if difficulty=="medium" else ""}>🟡 MEDIUM — 8 tasks, dynamic</option>
      <option value="hard"   {"selected" if difficulty=="hard"   else ""}>🔴 HARD — 12 tasks, scarce</option>
    </select>
    <button class="cmd-btn btn-reset" id="reset-btn" onclick="resetEnv()">
      <span id="reset-spinner" style="display:none" class="spinner"></span>⟳ RESET
    </button>
    <button class="cmd-btn btn-step"  id="step-btn"  onclick="stepEnv()"  disabled>▶ STEP</button>
    <button class="cmd-btn btn-auto"  id="auto-btn"  onclick="toggleAuto()" disabled>⏩ AUTO</button>
    <span class="status-pill" id="status-pill">STANDBY</span>
  </div>

  <!-- LEFT: Task Queue -->
  <div class="panel-left">
    <div class="panel-title">
      <span>📋 TASK QUEUE</span>
      <span id="task-badge" style="font-family:'Space Mono',monospace;font-size:0.65rem;color:var(--accent);">0 tasks</span>
    </div>
    <div class="task-scroll" id="task-list">
      <div class="task-empty">
        Reset the environment<br/>to load tasks
      </div>
    </div>
  </div>

  <!-- CENTRE: Main visualisation -->
  <div class="panel-centre">

    <!-- Resource meters -->
    <div class="resource-row">
      <div class="res-cell" id="rc-time">
        <div class="res-label">⏱ Time Remaining</div>
        <div class="res-value" id="res-time">—</div>
      </div>
      <div class="res-cell" id="rc-energy">
        <div class="res-label">⚡ Energy</div>
        <div class="res-value" id="res-energy">—</div>
      </div>
      <div class="res-cell" id="rc-budget">
        <div class="res-label">💰 Budget</div>
        <div class="res-value" id="res-budget">—</div>
      </div>
      <div class="res-cell" id="rc-value">
        <div class="res-label">🏆 Cumulative Value</div>
        <div class="res-value" id="res-value">—</div>
      </div>
    </div>

    <!-- Action picker -->
    <div class="action-panel">
      <div class="chart-title">SELECT ACTION</div>
      <div class="action-grid">
        <div class="act-btn active" onclick="selectAction('execute')" id="act-execute">
          <span class="act-icon">⚡</span>
          <div class="act-name">EXECUTE</div>
          <div class="act-sub">Full reward</div>
        </div>
        <div class="act-btn" onclick="selectAction('optimize')" id="act-optimize">
          <span class="act-icon">🔧</span>
          <div class="act-name">OPTIMIZE</div>
          <div class="act-sub">Reduce risk</div>
        </div>
        <div class="act-btn" onclick="selectAction('delegate')" id="act-delegate">
          <span class="act-icon">📤</span>
          <div class="act-name">DELEGATE</div>
          <div class="act-sub">35% reward</div>
        </div>
        <div class="act-btn" onclick="selectAction('defer')" id="act-defer">
          <span class="act-icon">⏳</span>
          <div class="act-name">DEFER</div>
          <div class="act-sub">Small penalty</div>
        </div>
      </div>
    </div>

    <!-- Outcome readout -->
    <div class="outcome-bar" id="outcome-bar">
      <span>›</span> Waiting for first action...
    </div>

    <!-- Reward history chart -->
    <div class="chart-area">
      <div class="chart-title">REWARD HISTORY — per step</div>
      <div class="reward-chart" id="reward-chart">
        <div style="color:var(--muted);font-size:0.72rem;width:100%;text-align:center;padding-bottom:10px;">No data yet</div>
      </div>
    </div>
  </div>

  <!-- RIGHT: Telemetry -->
  <div class="panel-right">

    <!-- Score ring -->
    <div class="score-hero" id="score-hero">
      <div class="score-ring">
        <svg viewBox="0 0 100 100" width="110" height="110">
          <circle class="ring-bg" cx="50" cy="50" r="45"/>
          <circle class="ring-fill" id="ring-fill" cx="50" cy="50" r="45"/>
        </svg>
        <div class="score-num" id="score-num">—</div>
      </div>
      <div class="score-label">Episode Score</div>
    </div>

    <!-- Stats -->
    <div class="stat-grid">
      <div class="stat-box cyan">
        <div class="sn" id="stat-step">0</div>
        <div class="sl">Steps</div>
      </div>
      <div class="stat-box amber">
        <div class="sn" id="stat-reward">0.00</div>
        <div class="sl">Last Reward</div>
      </div>
      <div class="stat-box green">
        <div class="sn" id="stat-done">0</div>
        <div class="sl">Completed</div>
      </div>
      <div class="stat-box red">
        <div class="sn" id="stat-fail">0</div>
        <div class="sl">Failed</div>
      </div>
    </div>

    <!-- Health -->
    <div class="health-section">
      <div class="health-label">
        <span>SYSTEM HEALTH</span>
        <span id="health-pct" style="color:var(--green);font-family:'Space Mono',monospace;">—</span>
      </div>
      <div class="health-track">
        <div class="health-fill" id="health-fill" style="width:0%"></div>
      </div>
    </div>

    <!-- Log -->
    <div class="log-panel">
      <div class="log-title">ACTION LOG</div>
      <div class="log-scroll" id="log">
        <div class="log-entry inf">› System online. Select difficulty and press RESET.</div>
      </div>
    </div>
  </div>

  <!-- STATUS BAR -->
  <div class="status-bar">
    <div class="sb-item"><div class="sb-dot ok" id="sb-api"></div>API ONLINE</div>
    <div class="sb-item"><div class="sb-dot" id="sb-ep"></div><span id="sb-ep-text">NO EPISODE</span></div>
    <div class="sb-item" style="margin-left:auto;">
      <span>Meta PyTorch OpenEnv Hackathon 2025</span>
    </div>
    <div class="sb-item">
      <a href="/docs" style="color:var(--muted);text-decoration:none;">OpenAPI ↗</a>
    </div>
  </div>
</div>

<script>
// ── STATE ──
let selectedTask   = null;
let selectedAction = 'execute';
let episodeDone    = false;
let autoTimer      = null;
let rewardHistory  = [];
let stepCount      = 0;

// ── HELPERS ──
function log(msg, cls='') {{
  const container = document.getElementById('log');
  const el = document.createElement('div');
  el.className = 'log-entry ' + cls;
  el.textContent = new Date().toLocaleTimeString('en-GB', {{hour:'2-digit',minute:'2-digit',second:'2-digit'}}) + ' › ' + msg;
  container.appendChild(el);
  container.scrollTop = container.scrollHeight;
}}

function setStatusPill(text, cls='') {{
  const p = document.getElementById('status-pill');
  p.textContent = text;
  p.className   = 'status-pill ' + cls;
}}

function selectAction(a) {{
  selectedAction = a;
  ['execute','optimize','delegate','defer'].forEach(x => {{
    document.getElementById('act-' + x).classList.toggle('active', x === a);
  }});
}}

function selectTask(id) {{
  selectedTask = id;
  document.querySelectorAll('.task-card').forEach(c => {{
    c.classList.toggle('selected', parseInt(c.dataset.id) === id);
  }});
}}

// ── RESOURCES ──
function updateResources(obs) {{
  const time   = obs.time_remaining   ?? 0;
  const energy = obs.energy_remaining ?? 0;
  const budget = obs.budget_remaining ?? 0;
  const value  = obs.cumulative_value ?? 0;
  const health = obs.system_health    ?? 1;

  document.getElementById('res-time').textContent   = time;
  document.getElementById('res-energy').textContent = energy.toFixed(1);
  document.getElementById('res-budget').textContent = budget.toFixed(1);
  document.getElementById('res-value').textContent  = value.toFixed(1);

  // fill bars — approximate max from difficulty
  const maxEnergy = 12; const maxBudget = 60;
  setResFill('rc-time',   Math.min(time/10,1)          * 100, time < 3 ? 'var(--red)' : time < 6 ? 'var(--amber)' : 'var(--green)');
  setResFill('rc-energy', Math.min(energy/maxEnergy,1) * 100, energy < 2 ? 'var(--red)' : energy < 5 ? 'var(--amber)' : 'var(--cyan)');
  setResFill('rc-budget', Math.min(budget/maxBudget,1) * 100, budget < 5 ? 'var(--red)' : budget < 15 ? 'var(--amber)' : 'var(--accent2)');
  setResFill('rc-value',  Math.min(value/50,1)         * 100, 'var(--amber)');

  // health bar
  const h   = Math.max(0, Math.min(1, health));
  const hc  = h > 0.6 ? 'var(--green)' : h > 0.3 ? 'var(--amber)' : 'var(--red)';
  document.getElementById('health-fill').style.width      = (h * 100) + '%';
  document.getElementById('health-fill').style.background = hc;
  document.getElementById('health-pct').textContent       = (h * 100).toFixed(0) + '%';
  document.getElementById('health-pct').style.color       = hc;

  // stats
  document.getElementById('stat-step').textContent = obs.step_number    ?? stepCount;
  document.getElementById('stat-done').textContent = obs.tasks_completed ?? 0;
  document.getElementById('stat-fail').textContent = obs.tasks_failed    ?? 0;

  // outcome
  if (obs.last_action_outcome) {{
    document.getElementById('outcome-bar').innerHTML =
      '<span>›</span> ' + obs.last_action_outcome;
  }}
}}

function setResFill(id, pct, color) {{
  const el = document.getElementById(id);
  el.style.setProperty('--fill', pct + '%');
  el.style.setProperty('--fill-color', color);
}}

// ── SCORE RING ──
function updateScore(score) {{
  const n = typeof score === 'number' ? score : parseFloat(score);
  if (isNaN(n)) return;
  const circumference = 283;
  const offset = circumference * (1 - Math.max(0, Math.min(1, n)));
  document.getElementById('ring-fill').style.strokeDashoffset = offset;
  const color = n >= 0.7 ? 'var(--green)' : n >= 0.4 ? 'var(--amber)' : 'var(--red)';
  document.getElementById('ring-fill').style.stroke = color;
  document.getElementById('score-num').textContent = n.toFixed(2);
  document.getElementById('score-num').style.color = color;
}}

// ── REWARD CHART ──
function pushReward(reward) {{
  rewardHistory.push(reward);
  if (rewardHistory.length > 30) rewardHistory.shift();
  renderChart();
}}

function renderChart() {{
  const chart = document.getElementById('reward-chart');
  if (rewardHistory.length === 0) {{
    chart.innerHTML = '<div style="color:var(--muted);font-size:0.72rem;width:100%;text-align:center;padding-bottom:10px;">No data yet</div>';
    return;
  }}
  const max = Math.max(...rewardHistory.map(Math.abs), 0.01);
  chart.innerHTML = '';
  rewardHistory.forEach((r, i) => {{
    const bar = document.createElement('div');
    bar.className   = 'reward-bar';
    bar.dataset.val = r.toFixed(2);
    const hPct = (Math.abs(r) / max) * 90 + 5;
    bar.style.height     = hPct + '%';
    bar.style.background = r >= 0.5 ? 'var(--green)' : r >= 0.3 ? 'var(--cyan)' : r >= 0 ? 'var(--amber)' : 'var(--red)';
    bar.style.opacity    = 0.5 + (i / rewardHistory.length) * 0.5;
    chart.appendChild(bar);
  }});
}}

// ── TASK CARDS ──
function renderTasks(tasks) {{
  const list = document.getElementById('task-list');
  document.getElementById('task-badge').textContent = (tasks?.length ?? 0) + ' tasks';

  if (!tasks || tasks.length === 0) {{
    list.innerHTML = '<div class="task-empty">All tasks resolved</div>';
    return;
  }}

  list.innerHTML = '';
  tasks.forEach(t => {{
    const urgency   = t.deadline <= 2;
    const medUrgent = t.deadline <= 4;
    const card = document.createElement('div');
    card.className = 'task-card' + (urgency ? ' urgent' : '') + (selectedTask === t.task_id ? ' selected' : '');
    card.dataset.id = t.task_id;
    card.onclick    = () => selectTask(t.task_id);

    const priCls  = t.priority  > 0.7 ? 'hi' : t.priority  > 0.4 ? 'mid' : 'lo';
    const dlCls   = urgency ? 'lo' : medUrgent ? 'mid' : 'hi';
    const riskCls = t.uncertainty > 0.6 ? 'lo' : t.uncertainty > 0.3 ? 'mid' : 'hi';

    card.innerHTML = `
      <div class="task-name">[#${{t.task_id}}] ${{t.name}}</div>
      <div class="task-chips">
        <span class="chip ${{priCls}}">PRI ${{(t.priority*100).toFixed(0)}}%</span>
        <span class="chip ${{dlCls}}">⏰ ${{t.deadline}}s</span>
        <span class="chip ${{riskCls}}">RISK ${{(t.uncertainty*100).toFixed(0)}}%</span>
        <span class="chip val">${{t.value.toFixed(1)}}v</span>
        <span class="chip">E${{t.required_energy.toFixed(1)}}</span>
        <span class="chip">B${{t.required_budget.toFixed(1)}}</span>
      </div>`;
    list.appendChild(card);
  }});

  if (selectedTask === null && tasks.length > 0) selectTask(tasks[0].task_id);
}}

// ── RESET ──
async function resetEnv() {{
  stopAuto();
  episodeDone    = false;
  rewardHistory  = [];
  stepCount      = 0;
  selectedTask   = null;
  renderChart();

  document.getElementById('reset-spinner').style.display = 'inline-block';
  document.getElementById('score-num').textContent = '—';
  document.getElementById('score-num').style.color = 'var(--text)';
  document.getElementById('ring-fill').style.strokeDashoffset = 283;
  document.getElementById('ring-fill').style.stroke = 'var(--accent)';
  document.getElementById('stat-reward').textContent = '0.00';

  setStatusPill('RESETTING…');
  const diff = document.getElementById('difficulty').value;

  try {{
    const r   = await fetch('/reset', {{
      method: 'POST',
      headers: {{'Content-Type': 'application/json'}},
      body: JSON.stringify({{difficulty: diff}})
    }});
    const data = await r.json();
    const obs  = data.observation || data;

    renderTasks(obs.tasks || []);
    updateResources(obs);

    document.getElementById('step-btn').disabled = false;
    document.getElementById('auto-btn').disabled = false;
    document.getElementById('sb-ep').classList.add('ok');
    document.getElementById('sb-ep-text').textContent = 'EPISODE ACTIVE · ' + diff.toUpperCase();
    setStatusPill('READY — ' + diff.toUpperCase(), 'ready');
    log('Episode reset · ' + diff + ' · ' + (obs.tasks||[]).length + ' tasks loaded', 'inf');
  }} catch(e) {{
    log('Reset failed: ' + e.message, 'wrn');
    setStatusPill('ERROR');
  }}
  document.getElementById('reset-spinner').style.display = 'none';
}}

// ── STEP ──
async function stepEnv() {{
  if (episodeDone) return;
  const tid = selectedTask ?? 0;

  try {{
    const r   = await fetch('/step', {{
      method: 'POST',
      headers: {{'Content-Type': 'application/json'}},
      body: JSON.stringify({{action_type: selectedAction, task_id: tid, reasoning: 'UI step'}})
    }});
    const data = await r.json();
    const obs  = data.observation || data;
    const reward = data.reward ?? obs.reward ?? 0;
    const done   = data.done   ?? obs.done   ?? false;

    stepCount++;
    pushReward(reward);
    renderTasks(obs.tasks || []);
    updateResources(obs);

    document.getElementById('stat-reward').textContent = reward.toFixed(3);
    document.getElementById('stat-reward').style.color = reward >= 0.5 ? 'var(--green)' : reward >= 0 ? 'var(--amber)' : 'var(--red)';

    const cls = reward >= 0.5 ? 'pos' : reward >= 0 ? 'wrn' : 'neg';
    log(selectedAction + '(task=' + tid + ') → reward=' + reward.toFixed(3) + (done ? ' [DONE]' : ''), cls);

    if (done) {{
      episodeDone = true;
      stopAuto();
      document.getElementById('step-btn').disabled = true;
      document.getElementById('auto-btn').disabled = true;
      document.getElementById('sb-ep-text').textContent = 'EPISODE COMPLETE';
      document.getElementById('sb-ep').classList.remove('ok');
      setStatusPill('EPISODE COMPLETE', 'done');
      document.getElementById('score-hero').classList.add('episode-done');
      await fetchFinalScore();
    }}
  }} catch(e) {{
    log('Step error: ' + e.message, 'wrn');
  }}
}}

// ── FETCH SCORE ──
async function fetchFinalScore() {{
  try {{
    const r    = await fetch('/state');
    const state = await r.json();
    const completed = state.tasks_completed ?? 0;
    const failed    = state.tasks_failed    ?? 0;
    const total     = completed + failed + (state.tasks?.length ?? 0);
    const eff       = total > 0 ? completed / total : 0;
    updateScore(eff);
    log('Final — completed ' + completed + '/' + total + ' · health ' + ((state.system_health??1)*100).toFixed(0) + '%', 'inf');
  }} catch(e) {{}}
}}

// ── AUTO PLAY ──
function toggleAuto() {{
  if (autoTimer) {{ stopAuto(); return; }}
  const btn = document.getElementById('auto-btn');
  btn.textContent = '⏹ STOP';
  btn.className   = 'cmd-btn btn-stop';
  autoTimer = setInterval(async () => {{
    if (episodeDone) {{ stopAuto(); return; }}
    await stepEnv();
  }}, 700);
}}

function stopAuto() {{
  clearInterval(autoTimer);
  autoTimer = null;
  const btn = document.getElementById('auto-btn');
  btn.textContent = '⏩ AUTO';
  btn.className   = 'cmd-btn btn-auto';
}}
</script>
</body>
</html>"""
    return HTMLResponse(content=html)


# ---------------------------------------------------------------------------
# API Routes
# ---------------------------------------------------------------------------

@app.post("/reset", response_model=ResetResponse, tags=["Environment Control"], summary="Reset the environment")
async def reset(request: ResetRequest = Body(default_factory=ResetRequest)) -> ResetResponse:
    response = await _persistent_manager.reset_environment(request.model_dump(exclude_unset=True))
    return ResetResponse(**response)


@app.post("/step", response_model=StepResponse, tags=["Environment Control"], summary="Execute an action")
async def step(payload: Dict[str, Any] = Body(default_factory=dict)) -> StepResponse:
    try:
        action_payload = _extract_action_payload(payload)
        response = await _persistent_manager.step_environment(action_payload)
    except ValidationError as exc:
        raise HTTPException(status_code=status.HTTP_422_UNPROCESSABLE_ENTITY, detail=exc.errors()) from exc
    return StepResponse(**response)


@app.get("/state", response_model=AetherTaskFlowState, tags=["State Management"], summary="Get current state")
async def get_state() -> AetherTaskFlowState:
    return _persistent_manager.env.state


@app.get("/schema", response_model=SchemaResponse, tags=["Schema"], summary="Get JSON schemas")
async def get_schemas() -> SchemaResponse:
    return SchemaResponse(
        action=AetherTaskFlowAction.model_json_schema(),
        observation=AetherTaskFlowObservation.model_json_schema(),
        state=AetherTaskFlowState.model_json_schema(),
    )

def _run_server(host: str, port: int, difficulty: str) -> None:
    if difficulty != _DIFFICULTY:
        os.environ["AETHER_DIFFICULTY"] = difficulty
    import uvicorn
    uvicorn.run(app, host=host, port=port)


def main() -> None:
    import argparse
    parser = argparse.ArgumentParser(description="AETHER-TaskFlow server")
    parser.add_argument("--host", default="0.0.0.0")
    parser.add_argument("--port", type=int, default=7860)
    parser.add_argument("--difficulty", choices=["easy", "medium", "hard"], default=_DIFFICULTY)
    args = parser.parse_args()
    _run_server(host=args.host, port=args.port, difficulty=args.difficulty)


if __name__ == "__main__":
    main()