"""NEURAL BOY (HTML-component build) — games dreamed by a diffusion world model.

This is an alternate front-end to ../space that renders the liked `mockup.html`
design as a single Gradio 6 custom HTML component:

    gr.HTML(html_template=..., js_on_load=..., server_functions=[load_game, step])

- html_template : the mockup console + cartridge library (static shell).
- css (global)  : the mockup styles + Press Start 2P font.
- js_on_load    : controls (keyboard + d-pad/A-B), drag-drop cartridges, frame loop.
- server_functions: load_game() seeds a session and step() advances one dreamed
  frame. Gradio passes a single arg per server call, so the js_on_load loop calls
  `await server.step({sid, action})` and paints the returned base64 PNG onto the
  console screen — no gr.Timer / gr.Image / hidden buttons. Actions are passed
  straight in, so no held-action TTL hack is needed.

Model code (`src/`) and cartridges (`assets/games/`) are reused from ../space if
not present locally, so nothing large is duplicated.
"""
import sys
from pathlib import Path

HERE = Path(__file__).resolve().parent


def _resolve(rel: str) -> Path:
    for base in (HERE, HERE.parent / "space"):
        if (base / rel).exists():
            return base / rel
    return HERE / rel


SRC = _resolve("src")
GAMES_DIR = _resolve("assets/games")
sys.path.insert(0, str(SRC))

import base64
import glob
import io
import json
import os

import gradio as gr
import numpy as np
import torch
from hydra.utils import instantiate
from omegaconf import OmegaConf
from PIL import Image

from data.episode import Episode
from models.diffusion import Denoiser, DiffusionSampler

DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu")
CTX = 4
DENOISING_STEPS = 10      # diffusion steps per frame (higher = sharper, slower)

# default action set for cartridges without meta.json (our Game Boy games)
DEFAULT_META = {
    "num_actions": 6,
    "action_names": ["noop", "left", "right", "fire", "leftfire", "rightfire"],
    "display_name": None,
}

TOKENS = ["up", "down", "left", "right", "fire"]

# Per-cartridge {game action token -> UI token}. Pong names its paddle actions
# "right"/"left" but they move the paddle up/down (ALE paddle quirk), so we bind
# the ↑/↓ keys instead of ◀/▶.
TOKEN_REMAP = {"pong": {"right": "up", "left": "down"}}

HIDDEN_GAMES = {"tobu", "pacman", "kangaroo", "crazyclimber"}    # cartridges to skip (case-insensitive substring)

if not OmegaConf.has_resolver("eval"):
    OmegaConf.register_new_resolver("eval", eval)


def tokens_of(name: str, remap: dict = None):
    toks = [t for t in TOKENS if t in name]
    if remap:
        toks = sorted((remap.get(t, t) for t in toks), key=TOKENS.index)
    return toks


def canonical(name: str, remap: dict = None) -> str:
    return "" if name == "noop" else "+".join(tokens_of(name, remap))


# ---------------- cartridges ----------------
GAME_DIRS = {
    p.name: p
    for p in sorted(GAMES_DIR.iterdir())
    if p.is_dir() and not any(h in p.name.lower() for h in HIDDEN_GAMES)
} if GAMES_DIR.exists() else {}
GAME_NAMES = list(GAME_DIRS)

_cache = {}


def get_game(name: str):
    if name not in _cache:
        d = GAME_DIRS[name]
        meta_path = d / "meta.json"
        meta = json.load(open(meta_path)) if meta_path.exists() else dict(DEFAULT_META)
        num_actions = int(meta["num_actions"])
        action_names = meta["action_names"]

        cfg = OmegaConf.load(d / "trainer.yaml")
        dcfg = instantiate(cfg.agent.denoiser)
        dcfg.inner_model.num_actions = num_actions
        denoiser = Denoiser(dcfg).to(DEVICE).eval()
        sd = torch.load(d / "agent.pt", map_location=DEVICE)
        denoiser.load_state_dict({k[len("denoiser."):]: v for k, v in sd.items() if k.startswith("denoiser.")})
        samp_cfg = instantiate(cfg.world_model_env.diffusion_sampler)
        samp_cfg.num_steps_denoising = DENOISING_STEPS
        sampler = DiffusionSampler(denoiser, samp_cfg)

        seeds = [Episode.load(p, map_location=DEVICE) for p in sorted(glob.glob(str(d / "seeds" / "*.pt")))]
        assert seeds, f"no seed episodes in {d/'seeds'}"

        noop = action_names.index("noop") if "noop" in action_names else 0
        remap = TOKEN_REMAP.get(name, {})
        keymap = {canonical(n, remap): i for i, n in enumerate(action_names)}
        keymap.setdefault("", noop)
        _cache[name] = {
            "sampler": sampler, "seeds": seeds, "noop": noop,
            "keymap": keymap, "display": (meta.get("display_name") or name.replace("-", " ")).upper(),
        }
    return _cache[name]


def pick_seed(seeds):
    """Pick a high-motion window so the dream starts somewhere interesting."""
    import random
    best, best_motion = None, -1.0
    for _ in range(30):
        ep = random.choice(seeds)
        max_start = len(ep) - CTX - 2
        if max_start < 1:
            continue
        start = random.randrange(max_start)
        clip = ep.obs[start + CTX : start + CTX + 16]
        motion = (clip[1:] - clip[:-1]).abs().mean().item() if len(clip) > 1 else 0.0
        if motion > best_motion:
            best, best_motion = (ep, start), motion
    return best


def to_dataurl(frame: torch.Tensor) -> str:
    """Encode a CHW [-1,1] frame as a tiny base64 PNG (CSS upscales it pixelated)."""
    x = ((frame.clamp(-1, 1) + 1) / 2).permute(1, 2, 0).cpu().numpy()
    im = Image.fromarray((x * 255).astype(np.uint8))
    buf = io.BytesIO()
    im.save(buf, format="PNG")
    return "data:image/png;base64," + base64.b64encode(buf.getvalue()).decode()


@torch.no_grad()
def model_step(sampler, obs_buf, act_buf, action_id):
    act_buf[:, -1] = action_id
    next_obs, _ = sampler.sample(obs_buf, act_buf)
    next_obs = next_obs.clamp(-1, 1)
    obs_buf = obs_buf.roll(-1, dims=1)
    act_buf = act_buf.roll(-1, dims=1)
    obs_buf[:, -1] = next_obs
    return obs_buf, act_buf, next_obs


# ---------------- sessions + server functions ----------------
# sid (generated client-side) -> {"game", "obs", "act"}. One frame per step() call.
SESSIONS = {}


# NOTE: Gradio server functions receive a SINGLE argument (the backend calls
# `fn(body.data)`, and the frontend forwards only the first JS argument), so each
# of these takes one dict that the JS side packs — extra positional args would be
# silently dropped and padded to None.
def load_game(payload):
    """Seed a session for a cartridge; return first frame + control keymap."""
    sid, name = payload["sid"], payload["name"]
    game = get_game(name)
    ep, start = pick_seed(game["seeds"])
    SESSIONS[sid] = {
        "game": name,
        "obs": ep.obs[start : start + CTX].unsqueeze(0).to(DEVICE),
        "act": ep.act[start : start + CTX].unsqueeze(0).to(DEVICE),
    }
    return {
        "frame": to_dataurl(SESSIONS[sid]["obs"][0, -1]),
        "keymap": game["keymap"],
        "name": game["display"],
        "noop": game["noop"],
    }


def step(payload):
    """Advance one dreamed frame for the session using the held action index."""
    sid = payload["sid"]
    s = SESSIONS.get(sid)
    if s is None:
        return {}
    game = get_game(s["game"])
    action = payload.get("action")
    a = int(action) if action is not None else game["noop"]
    s["obs"], s["act"], next_obs = model_step(game["sampler"], s["obs"], s["act"], a)
    return {"frame": to_dataurl(next_obs[0])}


if GAME_NAMES:
    get_game(GAME_NAMES[0])  # warm the first cartridge

GAMES_LIST = [{"id": n, "name": n.replace("-", " ").upper()} for n in GAME_NAMES]


# ---------------- UI ----------------
CSS = """
@import url('https://fonts.googleapis.com/css2?family=Press+Start+2P&display=swap');

html, body { background:#14102a; }
.gradio-container {
  max-width:100% !important;
  background: radial-gradient(ellipse at 20% 0%, #3b2d63 0%, #221a3d 45%, #14102a 100%) fixed !important;
}
.gradio-container, .gradio-container * { font-family:'Press Start 2P', monospace !important; }
footer { display:none !important; }
.gradio-container .block, .gradio-container .form, .html-container, .prose, .gap {
  background:transparent !important; border:none !important; box-shadow:none !important; padding:0 !important;
}

.nb-root { color:#f5edff; display:flex; flex-direction:column; align-items:center; gap:26px; padding:24px 16px 60px; }
.nb-root * { box-sizing:border-box; }

.nb-root header { text-align:center; }
.nb-root header h1 { font-size:30px; letter-spacing:3px; margin:0 0 12px;
  text-shadow:0 0 18px #c685ff88, 3px 3px 0 #1a1430; }
.nb-root header p { font-size:10px; color:#bda9e8; line-height:1.9; margin:0; }
.nb-root .pink { color:#ff8fc6; }

/* ---------- console ---------- */
.nb-root .console {
  width:460px; max-width:95vw;
  background:linear-gradient(160deg,#bdb4cf 0%,#9a8fb4 100%);
  border-radius:22px 22px 70px 22px; padding:26px 26px 34px;
  box-shadow:0 24px 60px #00000088, inset 0 2px 0 #ffffff55, inset 0 -6px 14px #00000033;
}
.nb-root .bezel { position:relative; background:#2e2a3e; border-radius:14px 14px 40px 14px;
  padding:20px 34px; box-shadow:inset 0 4px 12px #000000aa; }
.nb-root .powerled { position:absolute; left:14px; top:50%; transform:translateY(-50%);
  width:9px; height:9px; border-radius:50%; background:#ff3b5c; box-shadow:0 0 10px #ff3b5c; z-index:2; }
.nb-root .screen { position:relative; background:#0d1a0d; border-radius:6px; overflow:hidden; aspect-ratio:1/1; }
.nb-root .screen-img { position:absolute; inset:0; width:100%; height:100%; image-rendering:pixelated; display:none; }
.nb-root .screen .empty { position:absolute; inset:0; display:flex; align-items:center; justify-content:center;
  font-size:9px; color:#3f6b3f; text-align:center; line-height:2; }
.nb-root .scanlines { position:absolute; inset:0; pointer-events:none;
  background:repeating-linear-gradient(to bottom,#0000 0 2px,#00000022 2px 3px); }
.nb-root .gamename { text-align:center; font-size:8px; color:#9be29b; margin-top:8px; min-height:12px; letter-spacing:1px; }

.nb-root .brand { text-align:center; font-size:9px; color:#4a4060; letter-spacing:1px; margin:12px 0 16px; }
.nb-root .brand b { color:#c2447c; font-size:11px; font-weight:normal; }

/* ---------- controls ---------- */
.nb-root .controls { display:flex; align-items:center; justify-content:space-between;
  background:#00000022; border-radius:14px; padding:14px 22px; }
.nb-root .dpad { display:grid; grid-template-columns:repeat(3,32px); grid-template-rows:repeat(3,32px); }
.nb-root .pad { display:flex; align-items:center; justify-content:center; color:#cdbff0; font-size:11px;
  background:#3a3450; box-shadow:0 4px 0 #221d33; cursor:pointer; user-select:none; -webkit-user-select:none;
  touch-action:none; transition:transform .04s, box-shadow .04s, background .04s; }
.nb-root .pad.on, .nb-root .pad:active { transform:translateY(3px); box-shadow:0 1px 0 #221d33; background:#473d6b; color:#fff; }
.nb-root .pad.up { grid-area:1/2; border-radius:8px 8px 0 0; }
.nb-root .pad.left { grid-area:2/1; border-radius:8px 0 0 8px; }
.nb-root .pad.mid { grid-area:2/2; background:#3a3450; box-shadow:none; cursor:default; }
.nb-root .pad.mid::after { content:''; width:11px; height:11px; border-radius:50%; background:#221d33; }
.nb-root .pad.right { grid-area:2/3; border-radius:0 8px 8px 0; }
.nb-root .pad.down { grid-area:3/2; border-radius:0 0 8px 8px; }
.nb-root .abtns { display:flex; gap:12px; transform:rotate(-20deg); }
.nb-root .ab { width:42px; height:42px; border-radius:50%; display:flex; align-items:center; justify-content:center;
  background:#c2447c; color:#fff; font-size:11px; cursor:pointer; user-select:none; -webkit-user-select:none;
  touch-action:none; box-shadow:0 5px 0 #7c2450; transition:transform .04s, box-shadow .04s; }
.nb-root .ab.on, .nb-root .ab:active { transform:translateY(3px); box-shadow:0 1px 0 #7c2450; }

.nb-root .restart { display:block; margin:16px auto 0; font-family:inherit; font-size:9px; color:#fff;
  background:#c2447c; border:none; border-radius:10px; padding:10px 22px; cursor:pointer; box-shadow:0 4px 0 #7c2450; }
.nb-root .restart:active { transform:translateY(3px); box-shadow:0 1px 0 #7c2450; }

/* ---------- cartridge library ---------- */
.nb-root .library { width:460px; max-width:95vw; background:#ffffff0d; border:1px solid #ffffff1a;
  border-radius:18px; padding:18px; display:flex; flex-direction:column; gap:16px; }
.nb-root .lib-title { font-size:9px; color:#bda9e8; text-align:center; line-height:1.8; }
.nb-root .lib-body { display:flex; gap:18px; align-items:flex-start; }
.nb-root .rack { display:flex; flex-wrap:wrap; gap:12px; flex:1; }

.nb-root .cart { position:relative; width:96px; cursor:grab; user-select:none; -webkit-user-select:none;
  transition:transform .12s ease, filter .12s ease; }
.nb-root .cart:hover { transform:translateY(-6px); }
.nb-root .cart.dragging { opacity:.35; }
.nb-root .cart.active { filter:drop-shadow(0 0 10px #ff8fc6cc); }
.nb-root .cart.ghost { position:fixed; left:0; top:0; z-index:9999; width:96px; margin:0;
  pointer-events:none; opacity:.95; transition:none; transform:translate(-50%,-50%) rotate(-5deg);
  filter:drop-shadow(0 14px 22px #000a); }
.nb-root .cart.ghost:hover { transform:translate(-50%,-50%) rotate(-5deg); }
.nb-root .cart .body { background:linear-gradient(180deg,#4a4460,#332e47); border-radius:7px 7px 4px 4px;
  padding:7px 7px 12px; box-shadow:0 6px 14px #00000066, inset 0 1px 0 #ffffff22; }
.nb-root .cart .ridges { display:flex; gap:3px; justify-content:center; margin-bottom:6px; }
.nb-root .cart .ridges i { width:8px; height:5px; background:#00000044; border-radius:1px; }
.nb-root .cart .label { border-radius:4px; overflow:hidden; background:#0d1a0d; aspect-ratio:1/1; border:2px solid #1a1530; }
.nb-root .cart .label svg { width:100%; height:100%; display:block; image-rendering:pixelated; }
.nb-root .cart .name { text-align:center; font-size:7px; color:#cdbff0; margin-top:7px; letter-spacing:1px; }

.nb-root .slot { width:130px; flex:0 0 auto; }
.nb-root .slot-title { font-size:7px; color:#bda9e8; text-align:center; margin-bottom:8px; letter-spacing:1px; }
.nb-root .slot-well { background:#00000033; border:2px dashed #ffffff2a; border-radius:10px; min-height:150px;
  display:flex; align-items:center; justify-content:center; padding:12px; text-align:center;
  transition:border-color .12s, background .12s; }
.nb-root .slot-well.empty { font-size:7px; color:#7c6fa3; line-height:1.9; }
.nb-root .slot-well.over { border-color:#ff8fc6; background:#ff8fc61a; }
.nb-root .slot-well .cart { width:104px; cursor:pointer; }
.nb-root .eject { font-size:6px; color:#8d80b3; text-align:center; margin-top:8px; }

.nb-root .note { max-width:520px; text-align:center; font-size:8px; color:#8d80b3; line-height:2; }
.nb-root .note .pink { color:#ff8fc6; }
"""

HTML = """
<div class="nb-root">
  <header>
    <h1>NEURAL BOY</h1>
    <p>"The Matrix is everywhere. It is all around us. Even now, in this very room."</p>
  </header>

  <div class="console">
    <div class="bezel">
      <div class="powerled"></div>
      <div class="screen">
        <img class="screen-img" id="nb-screen" alt=""/>
        <div class="empty" id="nb-empty">INSERT A<br/>CARTRIDGE</div>
        <div class="scanlines"></div>
      </div>
      <div class="gamename" id="nb-gamename">— NO CARTRIDGE —</div>
    </div>

    <div class="brand"><b>NEURAL BOY</b>&trade;</div>

    <div class="controls">
      <div class="dpad">
        <div class="pad up" data-tok="up">▲</div>
        <div class="pad left" data-tok="left">◀</div>
        <div class="pad mid"></div>
        <div class="pad right" data-tok="right">▶</div>
        <div class="pad down" data-tok="down">▼</div>
      </div>
      <div class="abtns" id="nb-abtns">
        <div class="ab" data-tok="fire">B</div>
        <div class="ab" data-tok="fire">A</div>
      </div>
    </div>

    <button class="restart" id="nb-restart">RESTART</button>
  </div>

  <div class="library">
    <div class="lib-title"><span class="pink">drag a cartridge</span> into the slot to load a game</div>
    <div class="lib-body">
      <div class="rack" id="nb-rack"></div>
      <div class="slot">
        <div class="slot-title">SLOT</div>
        <div class="slot-well empty" id="nb-slot">DROP<br/>CARTRIDGE<br/>HERE</div>
        <div class="eject" id="nb-eject"></div>
      </div>
    </div>
  </div>

  <div class="note">every frame is <span class="pink">generated live</span> by a diffusion world model from
    the last few frames and the button you're holding. no emulator, no ROM, no game code. pure matrix.</div>
</div>
"""

JS = r"""
const root = element;
// gr.HTML renders the template *inside* `element`, so the actual .nb-root is a
// child. The ghost cartridge must live inside .nb-root for its scoped CSS
// (position:fixed, pointer-events:none, cart styling) to apply.
const stage = root.querySelector('.nb-root') || root;
const GAMES = Array.isArray(props.value) ? props.value : [];
const sid = (window.crypto && crypto.randomUUID) ? crypto.randomUUID()
                                                  : ('s' + Math.random().toString(36).slice(2));

const screen   = root.querySelector('#nb-screen');
const empty    = root.querySelector('#nb-empty');
const gamename = root.querySelector('#nb-gamename');
const rack     = root.querySelector('#nb-rack');
const slot     = root.querySelector('#nb-slot');
const abtns    = root.querySelector('#nb-abtns');
const restart  = root.querySelector('#nb-restart');

const tokenFor = {ArrowUp:'up', w:'up', ArrowDown:'down', s:'down',
                  ArrowLeft:'left', a:'left', ArrowRight:'right', d:'right', ' ':'fire', z:'fire'};
const ORDER = ['up','down','left','right','fire'];
const heldKeys = new Set();   // keyboard keys held
const heldPads = new Set();   // control *elements* held via pointer (.pad / .ab)
let current = null, keymap = {}, running = false, busy = false;

// Light the *specific* element(s) actually held. Pointer lights the exact button
// pressed; keyboard lights one representative element per token — so the two fire
// buttons (A/B both use data-tok="fire") never light up together.
function refreshVisual(){
  root.querySelectorAll('[data-tok]').forEach(el => el.classList.remove('on'));
  heldPads.forEach(el => el.classList.add('on'));
  heldKeys.forEach(k => {
    const el = root.querySelector('[data-tok="' + tokenFor[k] + '"]:not(.mid)');
    if (el) el.classList.add('on');
  });
}
function lookupIdx(toks){
  const key = ORDER.filter(t => toks.has(t)).join('+');
  return (key in keymap) ? keymap[key] : null;
}
function activeIdx(){
  const a = new Set();
  heldKeys.forEach(k => a.add(tokenFor[k]));
  heldPads.forEach(el => a.add(el.dataset.tok));
  let idx = lookupIdx(a);
  if (idx !== null) return idx;
  // Some games (e.g. KungFuMaster) have no standalone FIRE — punch only exists
  // as RIGHTFIRE/LEFTFIRE/DOWNFIRE. If the held combo isn't a real action but
  // FIRE is down, fall back to FIRE + a default direction so the button does
  // something instead of a silent noop.
  if (a.has('fire')){
    for (const d of ['right','left','down','up']){
      const t = new Set(a); t.add(d);
      idx = lookupIdx(t);
      if (idx !== null) return idx;
    }
  }
  return ('' in keymap) ? keymap[''] : 0;
}

/* ---------- per-game art (inline SVG placeholder box-art) ---------- */
function art(game){
  game = (game || '').toLowerCase();
  if (game === 'breakout'){
    const rows = ['#e2443f','#e08b3a','#e0c93a','#5fae47','#4f8fe0','#9a5fe0'];
    let bricks = '';
    rows.forEach((c,r) => { for (let x=0;x<8;x++){ bricks += `<rect x="${4+x*7.4}" y="${10+r*5}" width="6.6" height="4" fill="${c}"/>`; } });
    return `<svg viewBox="0 0 64 64"><rect width="64" height="64" fill="#0d0f1a"/>${bricks}<rect x="26" y="56" width="14" height="3" rx="1" fill="#e6e6e6"/><rect x="40" y="42" width="3" height="3" fill="#e6e6e6"/></svg>`;
  }
  if (game === 'pong'){
    return `<svg viewBox="0 0 64 64"><rect width="64" height="64" fill="#060606"/>${[...Array(8)].map((_,i)=>`<rect x="31" y="${4+i*8}" width="2" height="4" fill="#3a3a3a"/>`).join('')}<rect x="4" y="24" width="3" height="16" fill="#fff"/><rect x="57" y="14" width="3" height="16" fill="#fff"/><rect x="40" y="34" width="3" height="3" fill="#fff"/></svg>`;
  }
  if (game === 'boxing'){
    return `<svg viewBox="0 0 64 64"><rect width="64" height="64" fill="#5a8f3a"/><rect x="10" y="16" width="44" height="36" fill="none" stroke="#d98a2b" stroke-width="2"/><rect x="27" y="26" width="6" height="14" fill="#f4f4f4"/><rect x="31" y="26" width="6" height="14" fill="#1a1a1a"/></svg>`;
  }
  if (game === 'kangaroo'){
    return `<svg viewBox="0 0 64 64"><rect width="64" height="64" fill="#16203f"/><rect y="50" width="64" height="14" fill="#2a2350"/><g fill="#d98a3a"><rect x="26" y="24" width="14" height="17" rx="4"/><rect x="36" y="16" width="9" height="9" rx="3"/><rect x="43" y="11" width="3" height="6" rx="1"/><polygon points="26,36 8,46 13,49 28,42"/><rect x="29" y="39" width="5" height="13" rx="2"/><rect x="27" y="50" width="12" height="4" rx="2"/></g><rect x="30" y="29" width="7" height="8" rx="2" fill="#f0b870"/><rect x="38" y="19" width="2" height="2" fill="#16203f"/></svg>`;
  }
  if (game === 'kungfumaster'){
    return `<svg viewBox="0 0 64 64"><rect width="64" height="64" fill="#7a1f2a"/><rect y="52" width="64" height="12" fill="#3a1016"/><g fill="#f4f1e8"><rect x="20" y="14" width="8" height="8" rx="3"/><rect x="20" y="22" width="9" height="14" rx="2"/><rect x="12" y="24" width="9" height="4" rx="2"/><rect x="20" y="34" width="6" height="14" rx="2"/><polygon points="26,32 46,28 46,33 28,40"/></g><rect x="42" y="26" width="7" height="8" rx="1" fill="#e0c060"/><rect x="18" y="20" width="12" height="2" fill="#c2447c"/></svg>`;
  }
  if (game === 'crazyclimber'){
    let win = '';
    for (let r=0;r<9;r++) for (let c=0;c<4;c++){ const lit=(r*3+c*5)%4===0; win += `<rect x="${18+c*8}" y="${8+r*6}" width="5" height="4" fill="${lit?'#e8d24a':'#26303f'}"/>`; }
    return `<svg viewBox="0 0 64 64"><rect width="64" height="64" fill="#16263f"/><rect x="14" y="4" width="36" height="60" fill="#5a5a6e"/>${win}<rect x="30" y="31" width="5" height="7" rx="1" fill="#e2443f"/><rect x="27" y="30" width="4" height="3" rx="1" fill="#e2443f"/><rect x="34" y="30" width="4" height="3" rx="1" fill="#e2443f"/></svg>`;
  }
  return `<svg viewBox="0 0 64 64"><rect width="64" height="64" fill="#1a1530"/><text x="32" y="36" fill="#ff8fc6" font-size="9" text-anchor="middle" font-family="monospace">?</text></svg>`;
}

function makeCart(g){
  const el = document.createElement('div');
  el.className = 'cart'; el.dataset.id = g.id;
  el.innerHTML = '<div class="body"><div class="ridges"><i></i><i></i><i></i><i></i><i></i></div>'
    + '<div class="label">' + art(g.id) + '</div></div><div class="name">' + g.name + '</div>';
  return el;
}

function seat(id){
  const g = GAMES.find(x => x.id === id); if (!g) return;
  slot.classList.remove('empty','over'); slot.innerHTML = '';
  slot.appendChild(makeCart(g));
  [...rack.children].forEach(ch => ch.classList.toggle('active', ch.dataset.id === id));
}

async function insert(id){
  if (!id) return;
  let res;
  try { res = await server.load_game({sid: sid, name: id}); }
  catch (e) { gamename.textContent = 'LOAD FAILED'; console.error('load_game', id, e); return; }
  if (!res) return;
  current = id; keymap = res.keymap || {}; running = true;
  if (res.frame){ screen.src = res.frame; screen.style.display = 'block'; }
  if (empty) empty.style.display = 'none';
  gamename.textContent = res.name || id.toUpperCase();
  const hasFire = Object.keys(keymap).some(k => k.split('+').indexOf('fire') >= 0);
  if (abtns) abtns.style.visibility = hasFire ? 'visible' : 'hidden';
  seat(id);
}

GAMES.forEach(g => rack.appendChild(makeCart(g)));

/* Cartridge select: native HTML5 drag-drop is flaky inside the component, so we
   drive it with pointer events. A plain click loads instantly; a real drag spawns
   a "ghost" cartridge that follows the cursor and loads when dropped on the slot. */
let pickId = null, pickEl = null, ghost = null, downX = 0, downY = 0, dragging = false;

function overSlot(x, y){
  const el = document.elementFromPoint(x, y);
  return !!(el && el.closest && el.closest('.slot'));
}

rack.addEventListener('pointerdown', e => {
  const cart = e.target.closest('.cart'); if (!cart) return;
  e.preventDefault();
  pickId = cart.dataset.id; pickEl = cart; downX = e.clientX; downY = e.clientY; dragging = false;
});
document.addEventListener('pointermove', e => {
  if (!pickId) return;
  if (!dragging && Math.hypot(e.clientX - downX, e.clientY - downY) < 6) return;
  if (!dragging){
    dragging = true;
    ghost = pickEl.cloneNode(true); ghost.className = 'cart ghost'; stage.appendChild(ghost);
    pickEl.classList.add('dragging');
  }
  ghost.style.left = e.clientX + 'px';
  ghost.style.top  = e.clientY + 'px';
  slot.classList.toggle('over', overSlot(e.clientX, e.clientY));
});
document.addEventListener('pointerup', e => {
  if (!pickId) return;
  const id = pickId, wasDragging = dragging;
  if (ghost){ ghost.remove(); ghost = null; }
  if (pickEl) pickEl.classList.remove('dragging');
  slot.classList.remove('over');
  pickId = null; pickEl = null; dragging = false;
  if (!wasDragging || overSlot(e.clientX, e.clientY)) insert(id);
});
if (restart) restart.addEventListener('click', () => { if (current) insert(current); });

document.addEventListener('keydown', e => { if (e.key in tokenFor){ e.preventDefault(); heldKeys.add(e.key); refreshVisual(); } });
document.addEventListener('keyup',   e => { if (e.key in tokenFor){ heldKeys.delete(e.key); refreshVisual(); } });

root.addEventListener('pointerdown', e => {
  const el = e.target.closest('[data-tok]'); if (!el || el.classList.contains('mid')) return;
  e.preventDefault(); heldPads.add(el); refreshVisual();
});
function releasePads(){ if (heldPads.size){ heldPads.clear(); refreshVisual(); } }
document.addEventListener('pointerup', releasePads);
document.addEventListener('pointercancel', releasePads);

async function loop(){
  if (running && current && !busy){
    busy = true;
    try { const res = await server.step({sid: sid, action: activeIdx()}); if (res && res.frame) screen.src = res.frame; }
    catch (e) {}
    busy = false;
  }
  setTimeout(loop, (running && current) ? 0 : 200);
}
loop();

if (GAMES.length) insert(GAMES[0].id);
"""

with gr.Blocks(title="NEURAL BOY") as demo:
    gr.HTML(
        value=GAMES_LIST,
        html_template=HTML,
        js_on_load=JS,
        server_functions=[load_game, step],
        apply_default_css=False,
    )

demo.queue(default_concurrency_limit=8)
# Honor Gradio's env vars so this is correct on HF Spaces (which proxies port
# 7860) and still reachable when run on a remote droplet. Override locally with
# e.g. GRADIO_SERVER_PORT=7861 python app.py.
demo.launch(
    css=CSS,
    theme=gr.themes.Base(),
    server_name=os.environ.get("GRADIO_SERVER_NAME", "0.0.0.0"),
    server_port=int(os.environ.get("GRADIO_SERVER_PORT", 7860)),
)