#!/usr/bin/env bash
# Full OSWorld guest install — XFCE desktop running on Xvfb with all the
# apps whose binaries OSWorld's evaluators hard-reference. Roughly matches
# what the AWS AMI `ami-0d23263edb96951d8` ships with, modulo the
# Debian-vs-Ubuntu delta.
#
# Time budget: ~20–40 min on a c4m8 with gigabit egress. Most of that is
# apt download + LibreOffice/VSCode fetches.
#
# Prerequisites: run on a Debian 12+ FastVM guest as root. The bake script
# uploads the OSWorld server payload to /opt/osworld-server/ and the
# minimal install script first.
#
# Architecture delta vs minimal:
#   * minimal   : Xvfb :99 + `osworld_server` as its only graphical client.
#                 Screenshots come back as a 1920×1080 empty framebuffer.
#   * full      : single long-lived `osworld_session.service` that owns
#                 the whole desktop — starts Xvfb, starts a fixed-path
#                 session D-Bus, launches an XFCE session, then serves the
#                 OSWorld Flask app. Screenshots come back with a real
#                 desktop + task launcher + app windows.
#
# Using one service (not five that Require= each other) avoids cross-unit
# DBUS address plumbing and keeps teardown atomic.
set -euxo pipefail
export DEBIAN_FRONTEND=noninteractive

# Reuse the minimal install for system + python deps, Xauthority touch,
# etc. We'll override the service units below to a single-unit model.
/opt/osworld-server/install_osworld_minimal.sh

# Stop the split-unit services the minimal install set up; the full image
# replaces them with one osworld_session.service.
systemctl disable --now osworld_server.service xvfb.service || true
rm -f /etc/systemd/system/osworld_server.service /etc/systemd/system/xvfb.service
systemctl daemon-reload

# -------- build-essentials for adding third-party apt repos -------------
# We use `curl | gpg --dearmor` below to trust Google's + Microsoft's
# signing keys. gnupg isn't in the FastVM base image.
apt-get install -y --no-install-recommends gnupg apt-transport-https

# -------- desktop stack --------------------------------------------------
apt-get install -y --no-install-recommends \
    xfce4 xfce4-goodies xfce4-terminal \
    x11vnc novnc websockify \
    dbus-x11 dbus-user-session \
    fonts-dejavu fonts-liberation fonts-noto fonts-noto-cjk \
    accerciser \
    python3-pyatspi at-spi2-core gsettings-desktop-schemas dconf-cli

# -------- browsers -------------------------------------------------------
# Chromium via apt for the `chromium` binary name checks; google-chrome
# from Google's apt repo for evaluators that assert `google-chrome`.
apt-get install -y chromium chromium-driver || apt-get install -y chromium || true

if ! command -v google-chrome >/dev/null 2>&1; then
    curl -fsSL https://dl.google.com/linux/linux_signing_key.pub \
        | gpg --dearmor -o /usr/share/keyrings/google-chrome.gpg
    echo "deb [arch=amd64 signed-by=/usr/share/keyrings/google-chrome.gpg] https://dl.google.com/linux/chrome/deb/ stable main" \
        > /etc/apt/sources.list.d/google-chrome.list
    apt-get update
    apt-get install -y google-chrome-stable || true
fi

# Chrome-without-keyring workaround from upstream README — stops Chrome
# prompting for a login keyring on first launch.
mkdir -p /root/.local/share/keyrings
touch /root/.local/share/keyrings/login.keyring

# Patch .desktop launchers so Chrome/Chromium expose remote debugging on
# :9222 regardless of how they're launched (GUI menu, xdg-open, etc.).
for entry in /usr/share/applications/google-chrome.desktop /usr/share/applications/chromium.desktop; do
    [ -f "$entry" ] || continue
    sed -Ei 's#^(Exec=[^ ]+)( %U)?$#\1 --remote-debugging-port=9222 --remote-debugging-address=0.0.0.0\2#' "$entry"
done

# -------- office / productivity apps -------------------------------------
apt-get install -y --no-install-recommends \
    libreoffice libreoffice-gnome \
    gimp \
    vlc \
    thunderbird \
    || true    # tolerate individual package failures; we check binaries below.

# VSCode from Microsoft's repo.
if ! command -v code >/dev/null 2>&1; then
    curl -fsSL https://packages.microsoft.com/keys/microsoft.asc \
        | gpg --dearmor -o /usr/share/keyrings/packages.microsoft.gpg
    echo "deb [arch=amd64,arm64,armhf signed-by=/usr/share/keyrings/packages.microsoft.gpg] https://packages.microsoft.com/repos/code stable main" \
        > /etc/apt/sources.list.d/vscode.list
    apt-get update
    apt-get install -y code || true
fi

# -------- VLC HTTP interface --------------------------------------------
# The upstream AWS image enables the Lua HTTP interface via GUI prefs; we
# bake it as a config file so a fresh VM answers on :8080 immediately.
mkdir -p /root/.config/vlc
cat >/root/.config/vlc/vlcrc <<'EOF'
[core]
extraintf=http
[lua]
http-port=8080
http-host=::
http-password=password
EOF

# -------- dconf a11y defaults -------------------------------------------
# Toolkit accessibility is needed for Thunderbird/Firefox to expose a11y
# trees. System-level dconf default so every new session picks it up.
mkdir -p /etc/dconf/db/site.d
cat >/etc/dconf/db/site.d/00-osworld <<'EOF'
[org/gnome/desktop/interface]
toolkit-accessibility=true
EOF
mkdir -p /etc/dconf/profile
cat >/etc/dconf/profile/user <<'EOF'
user-db:user
system-db:site
EOF
dconf update || true

# -------- osworld session launcher --------------------------------------
# This script is pid1 of the osworld_session unit and owns the full
# desktop — Xvfb + D-Bus + XFCE + osworld_server — so all four share one
# env (same DISPLAY, same DBUS_SESSION_BUS_ADDRESS) without any
# cross-unit plumbing.
#
# Why not separate units: systemd's Environment= can't cleanly propagate
# a runtime-generated DBUS socket path between units. Fixed-path bus
# works, but you still need an explicit wait-for-socket + keepalive dance
# that isn't any simpler than a single wrapper script.
cat >/opt/osworld-server/start_session.sh <<'EOF'
#!/usr/bin/env bash
set -eu

export HOME=/root
export XDG_RUNTIME_DIR=/run/user/0
mkdir -p "$XDG_RUNTIME_DIR"
chmod 700 "$XDG_RUNTIME_DIR"

# Unique DBUS socket path. Using a path under XDG_RUNTIME_DIR keeps dbus
# from logging warnings about socket permissions.
export DBUS_SESSION_BUS_ADDRESS="unix:path=$XDG_RUNTIME_DIR/bus"
rm -f "$XDG_RUNTIME_DIR/bus"

# Xvfb :99 with RANDR so screen-resolution queries work and XTEST so
# PyAutoGUI can synthesize input events.
/usr/bin/Xvfb :99 -screen 0 1920x1080x24 -ac \
    +extension RANDR +extension GLX +extension XTEST &
XVFB_PID=$!
export DISPLAY=:99
export XAUTHORITY=/root/.Xauthority

# Wait for Xvfb's socket before starting anything that talks X.
for _ in 1 2 3 4 5 6 7 8 9 10; do
    if xdpyinfo >/dev/null 2>&1; then break; fi
    sleep 0.5
done

# Session D-Bus daemon at the fixed path. --nofork keeps it as a child so
# its exit takes us down too (systemd sees osworld_session die, restarts).
/usr/bin/dbus-daemon --session \
    --address="$DBUS_SESSION_BUS_ADDRESS" \
    --nofork --nopidfile --syslog-only &
DBUS_PID=$!
for _ in 1 2 3 4 5; do
    [ -S "$XDG_RUNTIME_DIR/bus" ] && break
    sleep 0.2
done

# XFCE desktop session against our Xvfb + dbus.
/usr/bin/startxfce4 &
XFCE_PID=$!

# Give XFCE a moment to paint its panel before the OSWorld server starts
# answering screenshots — nicer first-screenshot for demo screenshots.
sleep 3

# OSWorld Flask app — this is the foreground process systemd tracks.
# When main.py exits, the whole unit restarts.
exec /usr/bin/python3 /opt/osworld-server/serve.py
EOF
chmod +x /opt/osworld-server/start_session.sh

# -------- osworld_session systemd unit ----------------------------------
# One unit, one session, all services inside. Restart=always so a
# long-lived benchmark that crashes an app doesn't take the server down
# for the entire episode.
cat >/etc/systemd/system/osworld_session.service <<'EOF'
[Unit]
Description=OSWorld desktop session (Xvfb + D-Bus + XFCE + server)
After=network.target
StartLimitIntervalSec=60
StartLimitBurst=4

[Service]
Type=simple
ExecStart=/opt/osworld-server/start_session.sh
Environment=PYTHONUNBUFFERED=1
Restart=on-failure
RestartSec=3
# Kill the whole process group on stop so dbus/XFCE don't become orphans.
KillMode=control-group
TimeoutStopSec=10

[Install]
WantedBy=multi-user.target
EOF

# -------- noVNC on 5910 -------------------------------------------------
# x11vnc/novnc run *outside* the session unit because the VNC viewer
# needs to survive server restarts (otherwise every crashloop blanks the
# VNC tab the operator is watching).
cat >/etc/systemd/system/x11vnc.service <<'EOF'
[Unit]
Description=x11vnc against Xvfb :99
After=osworld_session.service
# Don't Require= the session — we want VNC to come back even if the
# session hiccups. x11vnc reconnects when the X server returns.

[Service]
ExecStart=/usr/bin/x11vnc -display :99 -rfbport 5900 -shared -forever -nopw -bg -o /var/log/x11vnc.log -q
Type=forking
Restart=on-failure
RestartSec=2

[Install]
WantedBy=multi-user.target
EOF

cat >/etc/systemd/system/novnc.service <<'EOF'
[Unit]
Description=noVNC (websockify) on :5910
After=x11vnc.service

[Service]
ExecStart=/usr/bin/websockify --web=/usr/share/novnc 5910 localhost:5900
Restart=on-failure
RestartSec=2

[Install]
WantedBy=multi-user.target
EOF

systemctl daemon-reload
systemctl enable osworld_session.service x11vnc.service novnc.service

# Start the session first and wait for Xvfb to bind :99 before firing the
# VNC services — x11vnc/novnc error out immediately if :99 isn't live,
# and systemd's StartLimitBurst=4 default means retrying under load
# trips the start-loop limiter.
systemctl start osworld_session.service
for _ in $(seq 1 30); do
    if DISPLAY=:99 xdpyinfo >/dev/null 2>&1; then break; fi
    sleep 1
done
systemctl start x11vnc.service novnc.service

# -------- verify --------------------------------------------------------
# Poll /screenshot for up to 90s — full image is slower to first-ready
# because XFCE paints a real desktop (wallpaper, panel, menu).
for i in $(seq 1 90); do
    if curl -sS --max-time 5 -o /tmp/osworld-boot.png http://127.0.0.1:5000/screenshot; then
        size=$(stat -c%s /tmp/osworld-boot.png)
        echo "OSWorld server is serving screenshots (${size} bytes)."
        # A real XFCE screenshot is typically >100 KB; a blank Xvfb
        # framebuffer is ~10 KB. This is a weak heuristic but handy.
        if [ "$size" -lt 20000 ]; then
            echo "WARNING: screenshot is small (${size} B) — XFCE may not have started yet."
        fi
        exit 0
    fi
    sleep 1
done

echo "ERROR: osworld_session did not come up in 90s" >&2
systemctl status osworld_session.service --no-pager || true
journalctl -u osworld_session.service --no-pager -n 80 || true
exit 1
