#!/bin/bash
set -euo pipefail

export DEBIAN_FRONTEND=noninteractive
assets=/tmp/beam-vm-assets
desktop=${1:-false}
kiosk=${2:-false}

apt_install() {
  apt-get install -y --no-install-recommends "$@"
}

apt_repository() {
  local name=$1 key_url=$2 repository=$3 suite=$4 component=$5
  install -d -m 0755 /etc/apt/keyrings
  curl -fsSL "$key_url" -o "/etc/apt/keyrings/$name.asc"
  printf 'deb [arch=%s signed-by=/etc/apt/keyrings/%s.asc] %s %s %s\n' \
    "$arch" "$name" "$repository" "$suite" "$component" \
    > "/etc/apt/sources.list.d/$name.list"
}

desktop_mime() {
  local desktop_id=$1 mime
  shift
  for mime in "$@"; do
    runuser -u "$desktop_user" -- xdg-mime default "$desktop_id" "$mime"
  done
}

if ! command -v apt-get >/dev/null; then
  echo 'Beam VM images require Debian/Ubuntu and systemd' >&2
  exit 1
fi
apt-get update
# Runtime libraries required by native coding-agent binaries.
apt_install systemd systemd-sysv dbus ca-certificates curl openssh-server rsync git \
  sudo procps python3 python3-pip libatomic1 libstdc++6 libgcc-s1

arch=$(dpkg --print-architecture)
case "$arch" in
  amd64) ttyd_arch=x86_64 ;;
  arm64) ttyd_arch=aarch64 ;;
  *) echo "unsupported VM architecture: $arch" >&2; exit 1 ;;
esac
curl -fsSL "https://github.com/tsl0922/ttyd/releases/download/1.7.7/ttyd.$ttyd_arch" -o /usr/local/bin/ttyd
chmod 0755 /usr/local/bin/ttyd

. /etc/os-release
case "$ID" in
  ubuntu|debian) ;;
  *) echo 'Beam VM images require Debian/Ubuntu' >&2; exit 1 ;;
esac
apt_repository nodesource https://deb.nodesource.com/gpgkey/nodesource-repo.gpg.key \
  https://deb.nodesource.com/node_22.x nodistro main
apt_repository docker "https://download.docker.com/linux/$ID/gpg" \
  "https://download.docker.com/linux/$ID" "$VERSION_CODENAME" stable
apt-get update
apt_install nodejs
if ! { command -v dockerd && command -v docker && docker compose version && docker buildx version; } >/dev/null 2>&1; then
  apt_install docker-ce docker-ce-cli containerd.io docker-buildx-plugin docker-compose-plugin
fi
rm -f /usr/local/bin/docker-compose
printf '#!/bin/sh\nexec docker compose "$@"\n' > /usr/local/bin/docker-compose
chmod 0755 /usr/local/bin/docker-compose
docker-compose version >/dev/null

npm install -g --no-fund --no-audit --allow-scripts=@anthropic-ai/claude-code \
  @openai/codex@0.147.0 @anthropic-ai/claude-code@2.1.226
codex --version
claude --version
npm cache clean --force

install -d /opt/beam-vm /workspace /run/sshd /etc/systemd/system
install -m 0755 "$assets/boot" "$assets/desktop" "$assets/xstartup" "$assets/kiosk" \
  "$assets/kiosk-openbox-build" "$assets/kiosk-selkies-patch" "$assets/kiosk-mcp" /opt/beam-vm/
install -m 0644 "$assets/kiosk-openbox.xml" "$assets/kiosk_config.py" /opt/beam-vm/
install -m 0644 "$assets"/*.service /etc/systemd/system/

# These are managed by Beam at boot, not by image build-time enablement.
# A missing unit aborts a batch disable on Noble. Disable each independently.
for unit in ssh.service ssh.socket getty@tty1.service serial-getty@hvc0.service \
  docker.service docker.socket containerd.service systemd-resolved.service \
  systemd-networkd.service systemd-networkd.socket; do
  systemctl disable "$unit" 2>/dev/null || true
done
rm -f /etc/ssh/ssh_host_* /etc/machine-id /var/lib/dbus/machine-id
touch /etc/machine-id
ln -s /etc/machine-id /var/lib/dbus/machine-id
cat > /etc/ssh/sshd_config.d/beam-vm.conf <<'EOF'
Port 2222
HostKey /etc/ssh/ssh_host_ed25519_key
PermitRootLogin prohibit-password
PasswordAuthentication no
KbdInteractiveAuthentication no
AuthorizedKeysFile /run/beam-vm/authorized_keys
PidFile /run/sshd/sshd.pid
EOF

if [ "$desktop" = true ] || [ "$kiosk" = true ]; then
  case "${ID}:${VERSION_CODENAME:-}" in
    ubuntu:jammy|ubuntu:noble) ;;
    *) echo 'Graphical VMs require Ubuntu 22.04 (jammy) or 24.04 (noble)' >&2; exit 1 ;;
  esac
  display_packages=(
    openbox libnss3-tools dbus-x11 xterm x11-utils x11-xserver-utils
    xdotool scrot xclip ffmpeg fonts-ubuntu xfonts-base adwaita-icon-theme
  )
  # Native packages bundle the web client and encoding extensions. Keep both
  # the version and digest pinned: image identity must not depend on "latest".
  version=2.1.0
  if [ "$VERSION_CODENAME" = noble ]; then
    package="selkies-${version}-ubuntu24.04-${arch}.deb"
    case "$arch" in
      amd64) checksum=cedba7a9d8bb4d07fd2b2139a518dc2c03452e0f0254c5bb088426623e2027f3 ;;
      arm64) checksum=7ed77ef1ddbb17e265adb01e4190a086df93f3af2bccbe326b9dcdb268279868 ;;
    esac
    curl -fsSL "https://github.com/selkies-project/selkies/releases/download/${version}/${package}" \
      -o /tmp/selkies.deb
    printf '%s  %s\n' "$checksum" /tmp/selkies.deb | sha256sum --check -
    apt_install /tmp/selkies.deb
    rm /tmp/selkies.deb
  else
    # Jammy has no native release package. Extract the portable distribution
    # at build time so starting a VM never requires FUSE or a download.
    case "$arch" in
      amd64)
        portable_arch=x86_64
        checksum=a0147192e9d20e98c7b899411d62194a4fb4e5f19447a4e08868d61cde09c43b
        ;;
      arm64)
        portable_arch=aarch64
        checksum=5573a85c1a538168d6cf04f7cc21c27ac5c9e40dab496b737d911ecf4fda8629
        ;;
    esac
    curl -fsSL "https://github.com/selkies-project/selkies/releases/download/${version}/selkies-${version}-${portable_arch}.AppImage" \
      -o /tmp/selkies.AppImage
    printf '%s  %s\n' "$checksum" /tmp/selkies.AppImage | sha256sum --check -
    chmod 0755 /tmp/selkies.AppImage
    (
      cd /opt
      /tmp/selkies.AppImage --appimage-extract >/dev/null
      mv squashfs-root selkies
    )
    ln -s /opt/selkies/AppRun /usr/local/bin/selkies
    rm /tmp/selkies.AppImage
  fi
  apt_install "${display_packages[@]}" xvfb xauth wmctrl xsettingsd pulseaudio \
    libegl1 libgbm1 libgl1-mesa-dri
  if [ "$kiosk" = true ]; then
    apt_install python3-venv
    python3 -m venv /opt/beam-vm/mcp
    mcp_version=$(PYTHONPATH=/opt/beam-vm python3 -c 'from kiosk_config import DesktopRuntime; print(DesktopRuntime().mcp_version)')
    /opt/beam-vm/mcp/bin/pip install --no-cache-dir "fastmcp==$mcp_version"
    /opt/beam-vm/kiosk-openbox-build
  fi
  /opt/beam-vm/kiosk-selkies-patch
  printf 'selkies-%s\n' "$version" > /opt/beam-vm/kiosk-backend
fi

if [ "$desktop" = true ] && [ "$kiosk" != true ]; then
  desktop_user=$(PYTHONPATH=/opt/beam-vm python3 -c 'from kiosk_config import DesktopRuntime; print(DesktopRuntime().desktop_user)')
  if ! id "$desktop_user" >/dev/null 2>&1; then
    useradd --create-home --user-group --shell /bin/bash "$desktop_user"
  fi
  desktop_home=$(getent passwd "$desktop_user" | cut -d: -f6)
  printf '%s\n' "$desktop_user" > /opt/beam-vm/desktop-user
  test -u /usr/bin/sudo
  printf '%s ALL=(ALL:ALL) NOPASSWD: ALL\n' "$desktop_user" > /etc/sudoers.d/beam-desktop
  chmod 0440 /etc/sudoers.d/beam-desktop
  visudo -cf /etc/sudoers.d/beam-desktop
  test "$(stat -c %u /usr/bin/sudo)" = 0
  chown -R "$desktop_user:$desktop_user" /workspace
  if [ "$arch" = amd64 ]; then
    apt_repository google-chrome https://dl.google.com/linux/linux_signing_key.pub \
      https://dl.google.com/linux/chrome/deb/ stable main
    browser=google-chrome-stable
  else
    # Google Chrome has no Linux arm64 package; retain Firefox on that architecture.
    apt_repository mozilla https://packages.mozilla.org/apt/repo-signing-key.gpg \
      https://packages.mozilla.org/apt mozilla main
    printf 'Package: *\nPin: origin packages.mozilla.org\nPin-Priority: 1000\n' > /etc/apt/preferences.d/mozilla
    browser=firefox
  fi
  apt-get update
  apt_install "$browser" ubuntu-session ubuntu-settings gnome-shell gnome-terminal \
    gnome-control-center gnome-software nautilus gvfs fonts-ubuntu yaru-theme-gtk yaru-theme-icon \
    gnome-shell-extension-ubuntu-dock gnome-shell-extension-appindicator \
    gnome-shell-extension-desktop-icons-ng librsvg2-common xdg-user-dirs \
    xdg-utils shared-mime-info at-spi2-core file-roller unzip
  # Native defaults remain overridable in each user's desktop settings.
  pixel_ratio=$(PYTHONPATH=/opt/beam-vm python3 -c 'from kiosk_config import DesktopRuntime; print(DesktopRuntime().pixel_ratio)')
  sed "s/@PIXEL_RATIO@/$pixel_ratio/g" "$assets/desktop.gschema.override" \
    > /usr/share/glib-2.0/schemas/99_beam.gschema.override
  glib-compile-schemas --strict /usr/share/glib-2.0/schemas
  # GNOME records this acknowledgment itself after its first non-GDM session.
  # Browser access authenticates this remote session; it has no local login screen.
  runuser -u "$desktop_user" -- mkdir -p "$desktop_home/.local/share/gnome-shell"
  runuser -u "$desktop_user" -- touch "$desktop_home/.local/share/gnome-shell/lock-warning-shown"
  install -d /etc/polkit-1/rules.d
  sed "s/@DESKTOP_USER@/$desktop_user/g" "$assets/packagekit.rules" \
    > /etc/polkit-1/rules.d/49-beam-packagekit.rules
  chmod 0644 /etc/polkit-1/rules.d/49-beam-packagekit.rules
  desktop_mime org.gnome.Software.desktop application/vnd.debian.binary-package application/x-deb
  if [ "$arch" = amd64 ]; then
    install -m 0755 "$assets/browser" /usr/local/bin/google-chrome-stable
    ln -sf google-chrome-stable /usr/local/bin/google-chrome
  fi
  install -d -m 0755 "$desktop_home/Desktop" /usr/local/share/applications
  install -m 0755 "$assets/browser.desktop" "$assets/terminal.desktop" "$desktop_home/Desktop"/
  runuser -u "$desktop_user" -- xdg-user-dirs-update
  desktop_mime org.gnome.Nautilus.desktop inode/directory
  if [ "$browser" = firefox ]; then
    sed -i 's/Name=Chrome/Name=Firefox/; s/Icon=google-chrome/Icon=firefox/; s/^Exec=.*/Exec=firefox %U/' \
      "$desktop_home/Desktop"/browser.desktop
  fi
  install -m 0644 "$desktop_home/Desktop/browser.desktop" /usr/local/share/applications/beam-browser.desktop
  desktop_mime beam-browser.desktop text/html application/xhtml+xml x-scheme-handler/http x-scheme-handler/https
  chown -R "$desktop_user:$desktop_user" "$desktop_home/Desktop" "$desktop_home/.config"
fi

if [ "$desktop" = true ] || [ "$kiosk" = true ]; then
  install -d /etc/ca-certificates/update.d
  install -m 0755 "$assets/browser-trust" /opt/beam-vm/browser-trust
  install -m 0755 "$assets/browser-trust-update" /etc/ca-certificates/update.d/beam-browser
  /opt/beam-vm/browser-trust --initialize
fi

# A desktop's package opener needs these indexes to resolve downloaded packages.
if [ "$desktop" != true ] || [ "$kiosk" = true ]; then
  rm -rf /var/lib/apt/lists/*
fi
rm -rf "$assets"
