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

log() {
  printf '[make-void-kernel] %s\n' "$*"
}

usage() {
  cat <<'EOF'
Usage: ./make-void-kernel.sh [--out-dir <path>] [--mirror <url>] [--arch <arch>] [--kernel-package <name>] [--print-register-flags]

Download and stage a Void Linux kernel under ./runtime/void-kernel for the
experimental Void guest flow.

Defaults:
  --out-dir         ./runtime/void-kernel
  --mirror          https://repo-default.voidlinux.org
  --arch            x86_64
  --kernel-package  linux6.12

The staged output contains:
  boot/vmlinux-<version>     Firecracker-usable kernel extracted from vmlinuz
  boot/vmlinuz-<version>     Raw distro boot image from the Void package
  boot/initramfs-<version>.img  Matching initramfs generated with dracut
  boot/config-<version>      Void kernel config
  lib/modules/<version>/     Matching kernel modules tree

If --print-register-flags is passed, the script does not download anything. It
prints the banger image register flags for an existing staged Void kernel.
EOF
}

require_command() {
  local name="$1"
  command -v "$name" >/dev/null 2>&1 || {
    log "required command not found: $name"
    exit 1
  }
}

normalize_mirror() {
  local mirror="${1%/}"
  mirror="${mirror%/current}"
  mirror="${mirror%/static}"
  printf '%s\n' "$mirror"
}

find_static_binary() {
  local name="$1"
  find "$STATIC_DIR" -type f \( -name "$name" -o -name "$name.static" \) -perm -u+x | sort | head -n 1
}

find_static_keys_dir() {
  find "$STATIC_DIR" -type d -path '*/var/db/xbps/keys' | sort | head -n 1
}

find_latest_matching() {
  local dir="$1"
  local pattern="$2"
  if [[ ! -d "$dir" ]]; then
    return 1
  fi
  find "$dir" -maxdepth 1 -type f -name "$pattern" | sort | tail -n 1
}

find_latest_module_dir() {
  local root="$1"
  if [[ ! -d "$root" ]]; then
    return 1
  fi
  find "$root" -mindepth 1 -maxdepth 1 -type d | sort | tail -n 1
}

print_register_flags() {
  local kernel=""
  local initrd=""
  local modules=""

  kernel="$(find_latest_matching "$OUT_DIR/boot" 'vmlinux-*' || true)"
  initrd="$(find_latest_matching "$OUT_DIR/boot" 'initramfs-*' || true)"
  modules="$(find_latest_module_dir "$OUT_DIR/lib/modules" || true)"

  if [[ -z "$kernel" || -z "$modules" ]]; then
    log "staged Void kernel not found under $OUT_DIR"
    exit 1
  fi

  printf -- '--kernel %q ' "$kernel"
  if [[ -n "$initrd" ]]; then
    printf -- '--initrd %q ' "$initrd"
  fi
  printf -- '--modules %q\n' "$modules"
}

check_elf() {
  local path="$1"
  readelf -h "$path" >/dev/null 2>&1
}

ensure_stage_root_layout() {
  mkdir -p "$STAGE_ROOT/usr"

  if [[ ! -e "$STAGE_ROOT/bin" ]]; then
    ln -snf usr/bin "$STAGE_ROOT/bin"
  fi
  if [[ ! -e "$STAGE_ROOT/sbin" ]]; then
    ln -snf usr/bin "$STAGE_ROOT/sbin"
  fi
  if [[ ! -e "$STAGE_ROOT/usr/sbin" ]]; then
    ln -snf bin "$STAGE_ROOT/usr/sbin"
  fi
  if [[ ! -e "$STAGE_ROOT/lib" ]]; then
    ln -snf usr/lib "$STAGE_ROOT/lib"
  fi
  if [[ ! -e "$STAGE_ROOT/lib64" ]]; then
    ln -snf usr/lib "$STAGE_ROOT/lib64"
  fi
  if [[ ! -e "$STAGE_ROOT/usr/lib64" ]]; then
    ln -snf lib "$STAGE_ROOT/usr/lib64"
  fi
  if [[ -x "$STAGE_ROOT/usr/bin/udevd" ]]; then
    mkdir -p "$STAGE_ROOT/usr/lib/udev" "$STAGE_ROOT/usr/lib/systemd"
    if [[ ! -e "$STAGE_ROOT/usr/lib/udev/udevd" ]]; then
      ln -snf ../../bin/udevd "$STAGE_ROOT/usr/lib/udev/udevd"
    fi
    if [[ ! -e "$STAGE_ROOT/usr/lib/systemd/systemd-udevd" ]]; then
      ln -snf ../../bin/udevd "$STAGE_ROOT/usr/lib/systemd/systemd-udevd"
    fi
  fi
}

sync_host_dracut_tree() {
  if [[ ! -d /usr/lib/dracut ]]; then
    log "host dracut support files not found under /usr/lib/dracut"
    exit 1
  fi
  rm -rf "$STAGE_ROOT/usr/lib/dracut"
  mkdir -p "$STAGE_ROOT/usr/lib"
  cp -a /usr/lib/dracut "$STAGE_ROOT/usr/lib/dracut"
}

build_initramfs() {
  local kver="$1"
  local modules_dir="$2"
  local out="$3"
  local config_dir="$TMP_DIR/dracut.conf.d"
  local tmpdir="$TMP_DIR/dracut-tmp"
  local force_drivers="virtio virtio_ring virtio_mmio virtio_blk virtio_net virtio_console ext4 vsock vmw_vsock_virtio_transport"

  mkdir -p "$config_dir" "$tmpdir"
  ensure_stage_root_layout
  sync_host_dracut_tree

  log "generating initramfs for kernel $kver with host dracut against the staged Void sysroot"
  env dracutbasedir="/usr/lib/dracut" dracut \
    --force \
    --kver "$kver" \
    --sysroot "$STAGE_ROOT" \
    --kmoddir "$modules_dir" \
    --conf /dev/null \
    --confdir "$config_dir" \
    --tmpdir "$tmpdir" \
    --no-hostonly \
    --filesystems "ext4" \
    --force-drivers "$force_drivers" \
    --gzip \
    "$out"
}

extract_vmlinux() {
  local image="$1"
  local out="$2"
  local tmp="$TMP_DIR/vmlinux.extract"

  if check_elf "$image"; then
    install -m 0644 "$image" "$out"
    return 0
  fi

  try_decompress() {
    local header="$1"
    local marker="$2"
    local command="$3"
    local pos=""

    while IFS= read -r pos; do
      [[ -n "$pos" ]] || continue
      pos="${pos%%:*}"
      tail -c+"$pos" "$image" | eval "$command" >"$tmp" 2>/dev/null || true
      if check_elf "$tmp"; then
        install -m 0644 "$tmp" "$out"
        return 0
      fi
    done < <(tr "$header\n$marker" "\n$marker=" < "$image" | grep -abo "^$marker" || true)

    return 1
  }

  try_decompress '\037\213\010' "xy" "gunzip" && return 0
  try_decompress '\3757zXZ\000' "abcde" "unxz" && return 0
  try_decompress "BZh" "xy" "bunzip2" && return 0
  try_decompress '\135\000\000\000' "xxx" "unlzma" && return 0
  try_decompress '\002!L\030' "xxx" "lz4 -d" && return 0
  try_decompress '(\265/\375' "xxx" "unzstd" && return 0

  return 1
}

resolve_kernel_package_file() {
  local escaped_name=""
  escaped_name="$(printf '%s\n' "$KERNEL_PACKAGE" | sed 's/[.[\*^$()+?{|]/\\&/g')"

  curl -fsSL "$REPO_URL/" |
    grep -o "${escaped_name}-[0-9][^\" >]*\\.${ARCH}\\.xbps" |
    sort -u |
    tail -n 1
}

cleanup() {
  if [[ -n "${TMP_DIR:-}" && -d "${TMP_DIR:-}" ]]; then
    rm -rf "$TMP_DIR"
  fi
}

SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
DEFAULT_RUNTIME_DIR="$SCRIPT_DIR"
if [[ -d "$SCRIPT_DIR/runtime" ]]; then
  DEFAULT_RUNTIME_DIR="$SCRIPT_DIR/runtime"
fi
RUNTIME_DIR="${BANGER_RUNTIME_DIR:-$DEFAULT_RUNTIME_DIR}"
OUT_DIR="$RUNTIME_DIR/void-kernel"
MIRROR="https://repo-default.voidlinux.org"
ARCH="x86_64"
KERNEL_PACKAGE="linux6.12"
PRINT_REGISTER_FLAGS=0

while [[ $# -gt 0 ]]; do
  case "$1" in
    --out-dir)
      OUT_DIR="${2:-}"
      shift 2
      ;;
    --mirror)
      MIRROR="${2:-}"
      shift 2
      ;;
    --arch)
      ARCH="${2:-}"
      shift 2
      ;;
    --kernel-package)
      KERNEL_PACKAGE="${2:-}"
      shift 2
      ;;
    --print-register-flags)
      PRINT_REGISTER_FLAGS=1
      shift
      ;;
    -h|--help)
      usage
      exit 0
      ;;
    *)
      log "unknown option: $1"
      usage
      exit 1
      ;;
  esac
done

MIRROR="$(normalize_mirror "$MIRROR")"
REPO_URL="$MIRROR/current"
STATIC_ARCHIVE_URL="$MIRROR/static/xbps-static-latest.x86_64-musl.tar.xz"

if [[ "$PRINT_REGISTER_FLAGS" == "1" ]]; then
  print_register_flags
  exit 0
fi

if [[ "$ARCH" != "x86_64" ]]; then
  log "unsupported arch: $ARCH"
  log "this experimental downloader currently supports only x86_64"
  exit 1
fi
if [[ ! -d "$RUNTIME_DIR" ]]; then
  log "runtime bundle not found: $RUNTIME_DIR"
  exit 1
fi
if [[ -e "$OUT_DIR" ]]; then
  log "output directory already exists: $OUT_DIR"
  log "remove it first if you want to re-stage a different Void kernel"
  exit 1
fi

require_command curl
require_command tar
require_command cp
require_command find
require_command grep
require_command cut
require_command readelf
require_command file
require_command install
require_command tail
require_command xz
require_command gzip
require_command bzip2
require_command dracut

TMP_DIR="$(mktemp -d -t banger-void-kernel-XXXXXX)"
STATIC_DIR="$TMP_DIR/static"
STAGE_ROOT="$TMP_DIR/root"
STAGE_OUT="$TMP_DIR/out"
STATIC_ARCHIVE="$TMP_DIR/xbps-static.tar.xz"
trap cleanup EXIT

mkdir -p "$STATIC_DIR" "$STAGE_ROOT/var/db/xbps/keys" "$STAGE_OUT/boot" "$STAGE_OUT/lib/modules"

log "downloading static XBPS from $STATIC_ARCHIVE_URL"
curl -fsSL "$STATIC_ARCHIVE_URL" -o "$STATIC_ARCHIVE"
tar -xf "$STATIC_ARCHIVE" -C "$STATIC_DIR"

XBPS_INSTALL="$(find_static_binary xbps-install)"
STATIC_KEYS_DIR="$(find_static_keys_dir)"
if [[ -z "$XBPS_INSTALL" || ! -x "$XBPS_INSTALL" ]]; then
  log "failed to locate xbps-install in the static archive"
  exit 1
fi
if [[ -z "$STATIC_KEYS_DIR" || ! -d "$STATIC_KEYS_DIR" ]]; then
  log "failed to locate Void repository keys in the static archive"
  exit 1
fi

cp -a "$STATIC_KEYS_DIR/." "$STAGE_ROOT/var/db/xbps/keys/"

KERNEL_PACKAGE_FILE="$(resolve_kernel_package_file)"
if [[ -z "$KERNEL_PACKAGE_FILE" ]]; then
  log "failed to resolve a package file for $KERNEL_PACKAGE in $REPO_URL"
  exit 1
fi

log "staging $KERNEL_PACKAGE_FILE into a temporary root"
env XBPS_ARCH="$ARCH" "$XBPS_INSTALL" -S -y -U -r "$STAGE_ROOT" -R "$REPO_URL" linux-base "$KERNEL_PACKAGE" dracut eudev >/dev/null

VMLINUX_RAW="$(find_latest_matching "$STAGE_ROOT/boot" 'vmlinuz-*' || true)"
KERNEL_CONFIG="$(find_latest_matching "$STAGE_ROOT/boot" 'config-*' || true)"
MODULES_DIR="$(find_latest_module_dir "$STAGE_ROOT/usr/lib/modules" || true)"
KERNEL_VERSION="$(basename "$MODULES_DIR")"
INITRAMFS_NAME="initramfs-${KERNEL_VERSION}.img"
INITRAMFS_RAW="$STAGE_OUT/boot/$INITRAMFS_NAME"

if [[ -z "$VMLINUX_RAW" || -z "$KERNEL_CONFIG" || -z "$MODULES_DIR" ]]; then
  log "staged Void kernel is missing expected boot artifacts"
  exit 1
fi
if [[ ! -x "$STAGE_ROOT/usr/bin/udevd" ]]; then
  log "staged Void sysroot is missing /usr/bin/udevd after package install"
  exit 1
fi

VMLINUX_BASE="$(basename "$VMLINUX_RAW")"
VMLINUX_OUT="$STAGE_OUT/boot/vmlinux-${VMLINUX_BASE#vmlinuz-}"
install -m 0644 "$VMLINUX_RAW" "$STAGE_OUT/boot/$VMLINUX_BASE"
install -m 0644 "$KERNEL_CONFIG" "$STAGE_OUT/boot/$(basename "$KERNEL_CONFIG")"
build_initramfs "$KERNEL_VERSION" "$MODULES_DIR" "$INITRAMFS_RAW"
cp -a "$MODULES_DIR" "$STAGE_OUT/lib/modules/"

log "extracting Firecracker kernel from $(basename "$VMLINUX_RAW")"
if ! extract_vmlinux "$VMLINUX_RAW" "$VMLINUX_OUT"; then
  log "failed to extract an uncompressed vmlinux from $VMLINUX_RAW"
  log "raw kernel image type: $(file -b "$VMLINUX_RAW")"
  exit 1
fi

cat >"$STAGE_OUT/metadata.json" <<EOF
{
  "package": "$KERNEL_PACKAGE_FILE",
  "kernel_path": "$OUT_DIR/boot/$(basename "$VMLINUX_OUT")",
  "raw_kernel_path": "$OUT_DIR/boot/$VMLINUX_BASE",
  "config_path": "$OUT_DIR/boot/$(basename "$KERNEL_CONFIG")",
  "initrd_path": "$OUT_DIR/boot/$INITRAMFS_NAME",
  "modules_dir": "$OUT_DIR/lib/modules/$(basename "$MODULES_DIR")"
}
EOF

mv "$STAGE_OUT" "$OUT_DIR"

log "staged Void kernel artifacts in $OUT_DIR"
log "kernel image: $OUT_DIR/boot/$(basename "$VMLINUX_OUT")"
log "initrd image: $OUT_DIR/boot/$INITRAMFS_NAME"
log "modules dir: $OUT_DIR/lib/modules/$(basename "$MODULES_DIR")"