#!/usr/bin/env bash
# Build a dd-able Ubuntu 24.04 server image for the tersler that boots
# straight into SSH. No root needed: it customizes the official cloud image in
# a QEMU/KVM VM (user-data.build.yaml), tests the result in a second VM
# (UEFI, bigger disk), and writes out/tersler.img.zst.
#
#   image/build-image.sh            build + test
#   image/build-image.sh --no-test  build only
#
# Settings (environment): RIG_HOSTNAME, RIG_USER, RIG_SSH_KEY, RIG_TZ,
# RIG_KEYMAP, IMG_SIZE, OUT, CACHE.
set -euo pipefail

HERE=$(cd "$(dirname "$0")" && pwd)
REPO=$(cd "$HERE/.." && pwd)
OUT=${OUT:-$HERE/out}
CACHE=${CACHE:-$HERE/cache}
WORK=$OUT/work

RIG_HOSTNAME=${RIG_HOSTNAME:-tersler}
RIG_USER=${RIG_USER:-$USER}
RIG_SSH_KEY=${RIG_SSH_KEY:-$HOME/.ssh/id_ed25519.pub}
RIG_TZ=${RIG_TZ:-Europe/Berlin}
RIG_KEYMAP=${RIG_KEYMAP:-de}
IMG_SIZE=${IMG_SIZE:-6G}          # what dd writes; the root partition grows to the whole SSD on first boot

CLOUD_BASE=https://cloud-images.ubuntu.com/noble/current
CLOUD_IMG=noble-server-cloudimg-amd64.img
NEED_MB=7000                      # peak: cloud image + build overlay + raw image + .zst
BUILD_TIMEOUT=1800
TEST_PORT=2222

log() { printf '\e[1;34m[image]\e[0m %s\n' "$*" >&2; }
die() { printf '\e[1;31m[image]\e[0m %s\n' "$*" >&2; exit 1; }

TEST=1
[[ ${1:-} == --no-test ]] && TEST=0

# ---- preflight -------------------------------------------------------------
for t in qemu-img qemu-system-x86_64 genisoimage mkpasswd zstd curl sha256sum ssh; do
    command -v "$t" >/dev/null || die "missing tool: $t"
done
[[ -r /dev/kvm && -w /dev/kvm ]] || die "no access to /dev/kvm"
[[ -r $RIG_SSH_KEY ]] || die "SSH public key not found: $RIG_SSH_KEY (set RIG_SSH_KEY)"
mkdir -p "$OUT" "$WORK" "$CACHE"
free_mb=$(df -BM --output=avail "$OUT" | tail -1 | tr -dc 0-9)
have_img=0; [[ -f $CACHE/$CLOUD_IMG ]] && have_img=600
(( free_mb + have_img >= NEED_MB )) || die "only ${free_mb} MB free at $OUT, need ~${NEED_MB} MB (set OUT=/other/disk)"
log "free space at $OUT: ${free_mb} MB"

# ---- 1. official cloud image, verified ------------------------------------
curl -fsSL -o "$CACHE/SHA256SUMS" "$CLOUD_BASE/SHA256SUMS"
want=$(awk -v f="*$CLOUD_IMG" '$2 == f {print $1}' "$CACHE/SHA256SUMS")
[[ -n $want ]] || die "$CLOUD_IMG not listed in SHA256SUMS"
if [[ ! -f $CACHE/$CLOUD_IMG ]] || [[ $(sha256sum < "$CACHE/$CLOUD_IMG" | cut -d' ' -f1) != "$want" ]]; then
    log "downloading $CLOUD_IMG"
    curl -fL --progress-bar -o "$CACHE/$CLOUD_IMG.part" "$CLOUD_BASE/$CLOUD_IMG"
    [[ $(sha256sum < "$CACHE/$CLOUD_IMG.part" | cut -d' ' -f1) == "$want" ]] || die "checksum mismatch"
    mv "$CACHE/$CLOUD_IMG.part" "$CACHE/$CLOUD_IMG"
fi
log "cloud image OK ($want)"

# ---- 2. seed ISO: user-data, meta-data, the project -----------------------
pwfile=$OUT/console-password.txt
if [[ ! -s $pwfile ]]; then
    (umask 077; head -c 30 /dev/urandom | base64 -w0 | tr -d '+/=' | cut -c1-20 > "$pwfile")
fi
pwhash=$(mkpasswd -m sha-512 "$(head -1 "$pwfile")")
sshkey=$(head -1 "$RIG_SSH_KEY")

seed=$WORK/seed
rm -rf "$seed"; mkdir -p "$seed"
sed -e "s|@HOSTNAME@|$RIG_HOSTNAME|g" -e "s|@USER@|$RIG_USER|g" -e "s|@TZ@|$RIG_TZ|g" \
    -e "s|@KEYMAP@|$RIG_KEYMAP|g" -e "s|@PASSWORD_HASH@|${pwhash//&/\\&}|g" -e "s|@SSH_KEY@|$sshkey|g" \
    "$HERE/user-data.build.yaml" > "$seed/user-data"
printf 'instance-id: %s-build\nlocal-hostname: %s\n' "$RIG_HOSTNAME" "$RIG_HOSTNAME" > "$seed/meta-data"
tar -C "$REPO" --exclude=./image/out --exclude=./image/cache --exclude='*.log' \
    -czf "$seed/tersler.tar.gz" .
genisoimage -quiet -output "$WORK/seed.iso" -volid cidata -joliet -rock "$seed"

# ---- 3. customize in a VM --------------------------------------------------
build=$WORK/build.qcow2
rm -f "$build"
qemu-img create -q -f qcow2 -F qcow2 -b "$CACHE/$CLOUD_IMG" "$build" "$IMG_SIZE"
log "customizing in a VM (apt upgrade + packages, ~5-15 min); console: $WORK/build-console.log"
timeout "$BUILD_TIMEOUT" qemu-system-x86_64 -enable-kvm -cpu host -m 4096 -smp 4 \
    -display none -no-reboot \
    -drive if=virtio,file="$build",discard=unmap,detect-zeroes=unmap \
    -drive media=cdrom,file="$WORK/seed.iso",readonly=on \
    -nic user,model=virtio-net-pci \
    -serial file:"$WORK/build-console.log" \
    || die "build VM failed or timed out (see $WORK/build-console.log)"
grep -q JUNK-RIG-BUILD-OK "$WORK/build-console.log" \
    || die "build VM did not finish its runcmd (see $WORK/build-console.log)"

raw=$OUT/tersler.img
log "writing $raw"
rm -f "$raw"
qemu-img convert -O raw "$build" "$raw"     # sparse: only used blocks take disk space
rm -f "$build"

# ---- 4. test boot: UEFI, a disk bigger than the image, no seed ISO ---------
if (( TEST )); then
    ovmf_code=/usr/share/OVMF/OVMF_CODE_4M.fd
    cp /usr/share/OVMF/OVMF_VARS_4M.fd "$WORK/vars.fd"
    test=$WORK/test.qcow2
    rm -f "$test"
    qemu-img create -q -f qcow2 -F raw -b "$raw" "$test" 40G
    log "test boot (UEFI, 40 GB disk); console: $WORK/test-console.log"
    qemu-system-x86_64 -enable-kvm -cpu host -m 2048 -smp 2 -display none \
        -drive if=pflash,format=raw,readonly=on,file="$ovmf_code" \
        -drive if=pflash,format=raw,file="$WORK/vars.fd" \
        -drive if=virtio,file="$test" \
        -nic user,model=e1000e,hostfwd=tcp:127.0.0.1:$TEST_PORT-:22 \
        -serial file:"$WORK/test-console.log" &
    qpid=$!
    sshc=(ssh -p "$TEST_PORT" -o BatchMode=yes -o ConnectTimeout=5 -o StrictHostKeyChecking=no
          -o UserKnownHostsFile=/dev/null -o LogLevel=ERROR "$RIG_USER@127.0.0.1")
    ok=0
    for _ in $(seq 60); do
        if "${sshc[@]}" true 2>/dev/null; then ok=1; break; fi
        kill -0 "$qpid" 2>/dev/null || break
        sleep 5
    done
    if (( ok )); then
        "${sshc[@]}" 'cloud-init status --wait >/dev/null; set -x
            hostname; uname -r; findmnt -no SIZE /; ip -br addr
            cloud-init query instance_id; id ubuntu 2>&1 | head -1
            dpkg -l | grep -cE "^ii +(xserver-xorg-core|snapd|ubuntu-desktop)" || true
            ls ~/tersler; sudo -n true && echo sudo-ok' 2>&1 | tee "$WORK/test-report.txt"
        "${sshc[@]}" 'sudo poweroff' || true
    fi
    for _ in $(seq 30); do kill -0 "$qpid" 2>/dev/null || break; sleep 2; done
    kill "$qpid" 2>/dev/null || true
    rm -f "$test" "$WORK/vars.fd"
    (( ok )) || die "test VM: no SSH login (see $WORK/test-console.log); raw image kept at $raw"
    log "test boot OK (report: $WORK/test-report.txt)"
fi

# ---- 5. compress -----------------------------------------------------------
log "compressing"
zstd -q -T0 -10 -f "$raw" -o "$raw.zst"
rm -f "$raw"
(cd "$OUT" && sha256sum tersler.img.zst > tersler.img.zst.sha256)
log "done: $raw.zst ($(du -h "$raw.zst" | cut -f1)); console password: $pwfile"
log "write it:  zstdcat $raw.zst | sudo dd of=/dev/sdX bs=4M conv=fsync status=progress"