#!/bin/sh
set -eu

usage() {
  cat <<'EOF'
Usage: scripts/build-kexec-image [options]

Build a Guix kexec installer tarball from this guix-tribes checkout.

Options:
  --channels=PATH       Guix channels file for time-machine
  --output=PATH         Output tarball path, or - for stdout
  --gzip-level=N        gzip level for final tarball (default: 4)
  -h, --help            Show this help
EOF
}

script_dir=$(CDPATH= cd -- "$(dirname "$0")" && pwd)
root_dir=$(CDPATH= cd -- "$script_dir/.." && pwd)
channels=
output=
gzip_level="${NBDE_KEXEC_GZIP_LEVEL:-4}"

for arg in "$@"; do
  case "$arg" in
    --channels=*)
      channels=${arg#*=}
      ;;
    --output=*)
      output=${arg#*=}
      ;;
    --gzip-level=*)
      gzip_level=${arg#*=}
      ;;
    -h|--help)
      usage
      exit 0
      ;;
    *)
      echo "unknown argument: $arg" >&2
      usage >&2
      exit 2
      ;;
  esac
done

[ -n "$channels" ] || {
  echo "missing required --channels=PATH" >&2
  usage >&2
  exit 2
}

[ -n "$output" ] || {
  echo "missing required --output=PATH" >&2
  usage >&2
  exit 2
}

load_guix_env() {
  for profile in \
    "$HOME/.config/guix/current/etc/profile" \
    "$HOME/.guix-profile/etc/profile" \
    /run/current-system/profile/etc/profile
  do
    [ -r "$profile" ] || continue
    # shellcheck disable=SC1090
    . "$profile"
  done
}

require_tool() {
  command -v "$1" >/dev/null 2>&1 || {
    echo "$1 command not found" >&2
    exit 1
  }
}

require_file() {
  path=$1
  label=$2
  [ -e "$path" ] || {
    echo "$label not found: $path" >&2
    exit 1
  }
}

require_store_dir() {
  path=$1
  label=$2
  case "$path" in
    /gnu/store/*) ;;
    *)
      echo "$label did not produce a /gnu/store path: ${path:-<empty>}" >&2
      exit 1
      ;;
  esac
  [ -d "$path" ] || {
    echo "$label store path does not exist: $path" >&2
    exit 1
  }
}

log() {
  echo "$*" >&2
}

load_guix_env

require_tool guix
require_tool guile
require_tool mksquashfs
require_file "$channels" "channels file"
require_file "$script_dir/kexec-run" "kexec runner"
require_file "$root_dir/examples/build-host-kexec-installer.scm" "system file"

tmp=$(mktemp -d /tmp/guix-kexec-image.XXXXXX)
system_build_stdout="$tmp/system-build.stdout"
static_kexec_stdout="$tmp/static-kexec.stdout"
closure_paths="$tmp/closure-paths"

cleanup() {
  chmod -R u+w "$tmp" >/dev/null 2>&1 || true
  rm -rf "$tmp"
}

trap cleanup EXIT INT TERM

log "building kexec installer system"
if ! guix time-machine -C "$channels" -- system build -L "$root_dir" \
  "$root_dir/examples/build-host-kexec-installer.scm" >"$system_build_stdout"; then
  [ ! -s "$system_build_stdout" ] || cat "$system_build_stdout" >&2
  echo "failed to build kexec installer system" >&2
  exit 1
fi
system=$(tail -n 1 "$system_build_stdout")
require_store_dir "$system" "kexec installer system"
require_file "$system/parameters" "kexec installer system parameters"
require_file "$system/boot" "kexec installer system boot directory"

log "building static kexec-tools"
if ! guix time-machine -C "$channels" -- build -e '(begin
  (use-modules (gnu packages linux)
               (guix build-system gnu))
  (static-package kexec-tools))' >"$static_kexec_stdout"; then
  [ ! -s "$static_kexec_stdout" ] || cat "$static_kexec_stdout" >&2
  echo "failed to build static kexec-tools" >&2
  exit 1
fi
static_kexec=$(tail -n 1 "$static_kexec_stdout")
require_store_dir "$static_kexec" "static kexec-tools"
require_file "$static_kexec/sbin/kexec" "static kexec binary"

guile -s /dev/stdin -- "$system" >"$tmp/metadata" <<'GUILE'
(use-modules (ice-9 match)
             (srfi srfi-1)
             (srfi srfi-13))

(define (field name fields)
  (match (assoc name fields)
    ((_ value) value)
    (_ (error "missing boot parameter field" name))))

(let* ((system (last (command-line)))
       (sexp   (call-with-input-file
                 (string-append system "/parameters")
                 read)))
     (match sexp
    (('boot-parameters fields ...)
     (let* ((kernel     (field 'kernel fields))
            (initrd     (field 'initrd fields))
            (root       (field 'root-device fields))
            (args       (field 'kernel-arguments fields))
            (boot-link  (string-append system "/boot"))
            (boot-path  (canonicalize-path boot-link))
            (boot-args
             (append
              (cond
               ((equal? root "tmpfs")
                '("rootfstype=tmpfs"))
               ((and (string? root) (not (string=? root "none")))
                (list (string-append "root=" root)))
               (else
                '()))
              (list (string-append "gnu.system=" system)
                    (string-append "gnu.load=" boot-path))
              args)))
       (display kernel)
       (newline)
       (display initrd)
       (newline)
       (display boot-path)
       (newline)
       (display (string-join boot-args " "))
       (newline)))
    (_
     (error "unrecognized boot parameters file" sexp))))
GUILE

kernel=$(sed -n '1p' "$tmp/metadata")
initrd=$(sed -n '2p' "$tmp/metadata")
boot_path=$(sed -n '3p' "$tmp/metadata")
cmdline=$(sed -n '4p' "$tmp/metadata")
boot_refs=$(guix gc --references "$boot_path")

store_item_root() {
  case "$1" in
    /gnu/store/*)
      printf '%s\n' "$1" | sed 's#^\(/gnu/store/[^/]*\).*#\1#'
      ;;
    *)
      echo "not a store item: $1" >&2
      exit 1
      ;;
  esac
}

initrd_store=$(store_item_root "$initrd")
parameters_store=$(store_item_root "$(readlink -f "$system/parameters")")
system_refs="$tmp/system-refs"

guix gc --references "$system" \
  | while IFS= read -r ref; do
      case "$ref" in
        "$initrd_store"|"$parameters_store")
          log "excluding already-copied boot artifact from embedded store: $ref"
          ;;
        *)
          printf '%s\n' "$ref"
          ;;
      esac
    done >"$system_refs"

mkdir -p "$tmp/kexec"
cp "$kernel" "$tmp/kexec/bzImage"
cp "$initrd" "$tmp/kexec/initrd"
chmod 644 "$tmp/kexec/initrd"
cp "$static_kexec/sbin/kexec" "$tmp/kexec/kexec-static"
chmod 755 "$tmp/kexec/kexec-static"
cp "$script_dir/kexec-run" "$tmp/kexec/run"
chmod 755 "$tmp/kexec/run"
printf '%s\n' "$cmdline" >"$tmp/kexec/cmdline"

{
  printf '%s\n' "$system" "$boot_path"
  printf '%s\n' $boot_refs
  cat "$system_refs"
  guix gc --requisites "$boot_path" $boot_refs $(cat "$system_refs")
} | sort -u >"$closure_paths"

squashfs_root="$tmp/squashfs-root"
mkdir -p "$squashfs_root"
while IFS= read -r path; do
  [ -n "$path" ] || continue
  cp -a "$path" "$squashfs_root/$(basename "$path")"
done <"$closure_paths"

log "packing Guix store closure into initrd"
mksquashfs "$squashfs_root" "$tmp/kexec/gnu-store.squashfs" \
  -comp gzip -Xcompression-level 9 -no-xattrs -noappend >&2

squashfs_cpio="$tmp/squashfs-cpio"
mkdir -p "$squashfs_cpio"
cp "$tmp/kexec/gnu-store.squashfs" "$squashfs_cpio/gnu-store.squashfs"
( cd "$squashfs_cpio" && printf 'gnu-store.squashfs' | cpio -o -H newc | gzip -9 ) \
  >> "$tmp/kexec/initrd"
rm "$tmp/kexec/gnu-store.squashfs"

log "writing kexec installer tarball"
case "$output" in
  -)
    GZIP="-$gzip_level" tar -C "$tmp" -czf - kexec
    ;;
  *)
    mkdir -p "$(dirname "$output")"
    GZIP="-$gzip_level" tar -C "$tmp" -czf "$output" kexec
    ;;
esac