#!/usr/bin/env bash
set -euo pipefail

usage() {
  cat <<USAGE
Usage:
  $(basename "$0") <iso-root>

Remasters the Ubuntu Server install source inside <iso-root> by merging the
base + default server squashfs layers, installing appliance-required packages,
applying security updates from noble-security, and repacking the rootfs.

Environment:
  ALGA_APPLIANCE_PREINSTALL_PACKAGES  Space-separated package names to install.
  ALGA_APPLIANCE_PREINSTALL_IMAGE     Docker image used for Linux build tooling.
USAGE
}

if [[ "${1:-}" == "-h" || "${1:-}" == "--help" ]]; then
  usage
  exit 0
fi

ISO_ROOT="${1:-}"
if [[ -z "$ISO_ROOT" ]]; then
  usage >&2
  exit 1
fi

if [[ ! -d "$ISO_ROOT/casper" ]]; then
  echo "ISO root does not contain casper directory: $ISO_ROOT" >&2
  exit 1
fi

if [[ ! -f "$ISO_ROOT/casper/ubuntu-server-minimal.squashfs" ]]; then
  echo "Missing base Ubuntu server squashfs in: $ISO_ROOT/casper" >&2
  exit 1
fi

if [[ ! -f "$ISO_ROOT/casper/ubuntu-server-minimal.ubuntu-server.squashfs" ]]; then
  echo "Missing default Ubuntu server layer squashfs in: $ISO_ROOT/casper" >&2
  exit 1
fi

if ! command -v docker >/dev/null 2>&1; then
  echo "docker is required for full rootfs preinstall remastering." >&2
  exit 1
fi

if ! docker info >/dev/null 2>&1; then
  echo "docker is not running; start Docker/Colima before using --preinstall-rootfs." >&2
  exit 1
fi

PACKAGES="${ALGA_APPLIANCE_PREINSTALL_PACKAGES:-ca-certificates curl git jq net-tools nodejs openssh-server unzip ufw}"
BUILDER_IMAGE="${ALGA_APPLIANCE_PREINSTALL_IMAGE:-ubuntu:24.04}"

printf 'Preinstalling packages/security updates into Ubuntu install source:\n'
printf '  ISO root: %s\n' "$ISO_ROOT"
printf '  Builder image: %s\n' "$BUILDER_IMAGE"
printf '  Packages: %s\n' "$PACKAGES"

docker run --rm -i --platform linux/amd64 --privileged \
  -e DEBIAN_FRONTEND=noninteractive \
  -e PREINSTALL_PACKAGES="$PACKAGES" \
  -v "$ISO_ROOT:/iso-root" \
  "$BUILDER_IMAGE" \
  bash -euo pipefail <<'CONTAINER_SCRIPT'
set -x

apt-get update
apt-get install -y --no-install-recommends ca-certificates squashfs-tools python3
rm -rf /var/lib/apt/lists/*

WORK_DIR="/work/alga-rootfs-preinstall"
ROOTFS="$WORK_DIR/rootfs"
LAYERFS="$WORK_DIR/server-layer"
rm -rf "$WORK_DIR"
mkdir -p "$ROOTFS" "$LAYERFS"

unsquashfs -f -d "$ROOTFS" /iso-root/casper/ubuntu-server-minimal.squashfs
# The default Ubuntu Server layer contains overlayfs whiteout markers as
# character devices. Docker overlay filesystems can reject those outside /dev.
# Extract the layer separately while ignoring those write errors, apply the
# known whiteout, then copy normal layer contents over the base rootfs.
unsquashfs -ignore-errors -no-exit-code -f -d "$LAYERFS" /iso-root/casper/ubuntu-server-minimal.ubuntu-server.squashfs
rm -f "$ROOTFS/etc/dpkg/dpkg.cfg.d/excludes"
cp -a "$LAYERFS"/. "$ROOTFS"/

# Keep package installation from starting daemons inside the chroot.
printf '#!/bin/sh\nexit 101\n' > "$ROOTFS/usr/sbin/policy-rc.d"
chmod 0755 "$ROOTFS/usr/sbin/policy-rc.d"

# Build-time sources are intentionally limited to the release pocket plus
# security pocket. We do not enable noble-updates here because the intent is to
# pre-apply security updates, not drift to all available updates.
APT_BACKUP="$WORK_DIR/apt-backup"
mkdir -p "$APT_BACKUP"
if [[ -f "$ROOTFS/etc/apt/sources.list" ]]; then
  mv "$ROOTFS/etc/apt/sources.list" "$APT_BACKUP/sources.list"
fi
if [[ -d "$ROOTFS/etc/apt/sources.list.d" ]]; then
  mv "$ROOTFS/etc/apt/sources.list.d" "$APT_BACKUP/sources.list.d"
fi
mkdir -p "$ROOTFS/etc/apt/sources.list.d"
cat > "$ROOTFS/etc/apt/sources.list" <<'APT_SOURCES'
deb http://archive.ubuntu.com/ubuntu noble main restricted universe multiverse
deb http://security.ubuntu.com/ubuntu noble-security main restricted universe multiverse
APT_SOURCES

cp /etc/resolv.conf "$ROOTFS/etc/resolv.conf"
mount --bind /dev "$ROOTFS/dev"
mount -t devpts devpts "$ROOTFS/dev/pts"
mount -t proc proc "$ROOTFS/proc"
mount -t sysfs sysfs "$ROOTFS/sys"
cleanup_mounts() {
  set +e
  umount -lf "$ROOTFS/dev/pts" 2>/dev/null || true
  umount -lf "$ROOTFS/dev" 2>/dev/null || true
  umount -lf "$ROOTFS/proc" 2>/dev/null || true
  umount -lf "$ROOTFS/sys" 2>/dev/null || true
}
trap cleanup_mounts EXIT

chroot "$ROOTFS" apt-get update
# shellcheck disable=SC2086
chroot "$ROOTFS" apt-get install -y --no-install-recommends $PREINSTALL_PACKAGES
chroot "$ROOTFS" apt-get upgrade -y
chroot "$ROOTFS" apt-get clean

cleanup_mounts
trap - EXIT

rm -f "$ROOTFS/usr/sbin/policy-rc.d"
rm -rf "$ROOTFS/var/lib/apt/lists"/* "$ROOTFS/var/cache/apt/archives"/*.deb "$ROOTFS/tmp"/* "$ROOTFS/var/tmp"/*
# Never bake identity material generated by package postinst scripts into the appliance image.
rm -f "$ROOTFS/etc/ssh"/ssh_host_*
: > "$ROOTFS/etc/machine-id"
rm -f "$ROOTFS/var/lib/dbus/machine-id" "$ROOTFS/var/lib/systemd/random-seed"

rm -f "$ROOTFS/etc/apt/sources.list"
rm -rf "$ROOTFS/etc/apt/sources.list.d"
if [[ -f "$APT_BACKUP/sources.list" ]]; then
  mv "$APT_BACKUP/sources.list" "$ROOTFS/etc/apt/sources.list"
fi
if [[ -d "$APT_BACKUP/sources.list.d" ]]; then
  mv "$APT_BACKUP/sources.list.d" "$ROOTFS/etc/apt/sources.list.d"
else
  mkdir -p "$ROOTFS/etc/apt/sources.list.d"
fi

chroot "$ROOTFS" dpkg-query -W --showformat='${Package} ${Version}\n' \
  > /iso-root/casper/ubuntu-server-minimal.manifest
ROOTFS_SIZE="$(du -sx --block-size=1 "$ROOTFS" | awk '{print $1}')"
printf '%s\n' "$ROOTFS_SIZE" > /iso-root/casper/ubuntu-server-minimal.size

rm -f /iso-root/casper/ubuntu-server-minimal.squashfs
mksquashfs "$ROOTFS" /iso-root/casper/ubuntu-server-minimal.squashfs \
  -noappend -comp xz -b 1M

cat > /iso-root/casper/install-sources.yaml <<EOF
- default: true
  description:
    en: AlgaPSA Ubuntu Server with appliance-required packages and security updates preinstalled.
  id: alga-ubuntu-server
  locale_support: locale-only
  name:
    en: AlgaPSA Ubuntu Server
  path: ubuntu-server-minimal.squashfs
  size: ${ROOTFS_SIZE}
  type: fsimage
  variant: server
EOF

(
  cd /iso-root/casper
  sha256sum \
    ubuntu-server-minimal.squashfs \
    ubuntu-server-minimal.ubuntu-server.installer.generic-hwe.squashfs \
    ubuntu-server-minimal.ubuntu-server.installer.generic.squashfs \
    ubuntu-server-minimal.ubuntu-server.installer.squashfs \
    ubuntu-server-minimal.ubuntu-server.squashfs \
    > SHA256SUMS
)
rm -f /iso-root/casper/SHA256SUMS.gpg

python3 - <<'PY'
from pathlib import Path
manifest = Path('/iso-root/casper/ubuntu-server-minimal.manifest')
packages = {line.split()[0] for line in manifest.read_text().splitlines() if line.strip()}
required = set((Path('/tmp/required-packages.txt').read_text().split() if Path('/tmp/required-packages.txt').exists() else []))
missing = sorted(required - packages)
if missing:
    raise SystemExit(f'Missing expected preinstalled packages: {missing}')
PY
CONTAINER_SCRIPT