#!/usr/bin/env bash
# Minimal OSWorld guest-server install for a Debian-based FastVM guest.
#
# "Minimal" = enough to make POST /screenshot and POST /execute respond.
# No XFCE, Chromium, LibreOffice, etc. — that lives in the full install
# script. Good enough for:
#   * quickstart.py single-step smoke tests
#   * control-plane throughput / snapshot-latency benchmarks
#   * wiring up an agent harness without baking the fat desktop stack
#
# Run inside the VM as root; the bake script uploads the server payload
# to /opt/osworld-server/ before invoking this.
set -euxo pipefail
export DEBIAN_FRONTEND=noninteractive

# -------- system deps ----------------------------------------------------
apt-get update
apt-get install -y --no-install-recommends \
    xvfb x11-utils \
    python3 python3-pip python3-tk python3-xlib \
    python3-pyatspi at-spi2-core \
    scrot wmctrl xclip socat ffmpeg gnome-screenshot \
    libxtst6 libxkbcommon0 \
    curl ca-certificates

# -------- python deps ----------------------------------------------------
# Debian 12 enforces PEP 668; --break-system-packages is intentional because
# this image is single-purpose. Pinning is intentionally loose — the OSWorld
# requirements.txt pins an ancient pynput off a PR we don't need for the
# minimal server.
python3 -m pip install --break-system-packages --no-cache-dir \
    flask "pyautogui>=0.9" "Pillow>=10" requests "lxml>=4" numpy

# Provide a `python` symlink; the upstream systemd unit ExecStart points at
# /usr/bin/python, but Debian only ships /usr/bin/python3.
if [ ! -e /usr/bin/python ]; then
    ln -s /usr/bin/python3 /usr/bin/python
fi

# -------- Xvfb systemd unit ---------------------------------------------
cat >/etc/systemd/system/xvfb.service <<'EOF'
[Unit]
Description=X Virtual Framebuffer for OSWorld
After=network.target

[Service]
ExecStart=/usr/bin/Xvfb :99 -screen 0 1920x1080x24 -ac +extension RANDR +extension GLX +extension XTEST
Restart=on-failure
RestartSec=2

[Install]
WantedBy=multi-user.target
EOF

# Xlib's Xauthority lookup keys off $HOME; systemd service units start
# with a minimal env, so without an explicit HOME Python-xlib crashes with
# "$HOME not set, cannot find ~/.Xauthority" before main.py even parses.
# Touching /root/.Xauthority satisfies the file existence check; Xvfb's -ac
# means we don't actually need a cookie in it.
touch /root/.Xauthority

# -------- dual-stack Flask wrapper --------------------------------------
# Upstream main.py hardcodes `app.run(host="0.0.0.0")` (IPv4 only) under
# its `__main__` guard. FastVM guests only expose public IPv6, so an
# IPv4-only listener is unreachable from the outside. Bolt on a tiny
# wrapper that imports the same `app` and serves on `::` (dual-stack on
# Linux — IPV6_V6ONLY=0 by default, so v4 connects still work).
cat >/opt/osworld-server/serve.py <<'EOF'
import os, sys
HERE = os.path.dirname(os.path.abspath(__file__))
if HERE not in sys.path:
    sys.path.insert(0, HERE)
from main import app
if __name__ == "__main__":
    # debug=False avoids the Werkzeug reloader, which double-forks and
    # confuses systemd Type=simple (systemd tracks the parent, which exits).
    app.run(host="::", port=5000, debug=False, threaded=True)
EOF

# -------- osworld_server systemd unit ------------------------------------
# Differs from the upstream unit in desktop_env/server/osworld_server.service
# because we run as root against Xvfb :99 (no real user/session/dbus). For
# the minimal image that's fine — DBUS is only used for wallpaper changes.
cat >/etc/systemd/system/osworld_server.service <<'EOF'
[Unit]
Description=OSWorld guest server
After=xvfb.service network.target
Requires=xvfb.service
StartLimitIntervalSec=60
StartLimitBurst=4

[Service]
Type=simple
ExecStart=/usr/bin/python3 /opt/osworld-server/serve.py
WorkingDirectory=/opt/osworld-server
Environment=DISPLAY=:99
Environment=HOME=/root
Environment=XAUTHORITY=/root/.Xauthority
Environment=PYTHONUNBUFFERED=1
Restart=on-failure
RestartSec=3

[Install]
WantedBy=multi-user.target
EOF

systemctl daemon-reload
systemctl enable --now xvfb.service

# Give Xvfb a beat to come up before the server imports Xlib/pyatspi. Cheap
# and keeps the first service start from flapping through on-failure.
for i in 1 2 3 4 5; do
    if xdpyinfo -display :99 >/dev/null 2>&1; then break; fi
    sleep 1
done

systemctl enable --now osworld_server.service

# -------- verify ---------------------------------------------------------
# Poll /screenshot for up to 30s; bake script will also re-verify from
# outside the VM.
for i in $(seq 1 30); do
    if curl -sS --max-time 5 -o /tmp/osworld-boot.png http://127.0.0.1:5000/screenshot; then
        echo "OSWorld server is serving screenshots (/tmp/osworld-boot.png)."
        ls -l /tmp/osworld-boot.png
        exit 0
    fi
    sleep 1
done

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