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

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

usage() {
  cat <<'EOF'
Usage: ./scripts/make-alpine-kernel.sh [--out-dir <path>] [--release <x.y.z>] [--mirror <url>] [--arch <arch>] [--print-register-flags]

Download and stage an Alpine Linux virt kernel under ./build/manual/alpine-kernel
for the experimental Alpine guest flow.

Defaults:
  --out-dir    ./build/manual/alpine-kernel
  --release    3.23.3
  --mirror     https://dl-cdn.alpinelinux.org/alpine
  --arch       x86_64

The staged output contains:
  boot/vmlinuz-<version>        Alpine virt kernel image
  boot/initramfs-<version>.img  Matching Alpine initramfs
  boot/config-<version>         Alpine kernel config when present
  lib/modules/<version>/        Matching kernel modules from modloop-virt

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

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

check_elf() {
  local path="$1"
  readelf -h "$path" >/dev/null 2>&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"
  local dir=""
  if [[ ! -d "$root" ]]; then
    return 1
  fi
  while IFS= read -r dir; do
    if [[ -d "$dir/kernel" || -f "$dir/modules.dep" || -f "$dir/modules.dep.bin" ]]; then
      printf '%s\n' "$dir"
      return 0
    fi
  done < <(find "$root" -mindepth 1 -maxdepth 1 -type d | sort)
  return 1
}

find_tar_entry() {
  local archive="$1"
  local needle="$2"
  local entry=""

  while IFS= read -r entry; do
    case "$entry" in
      "$needle"|*/"$needle")
        printf '%s\n' "$entry"
        return 0
        ;;
    esac
  done < <(tar -tf "$archive")

  return 1
}

find_tar_config_entry() {
  local archive="$1"
  local entry=""

  while IFS= read -r entry; do
    case "$entry" in
      config-*-virt|*/config-*-virt)
        printf '%s\n' "$entry"
        return 0
        ;;
    esac
  done < <(tar -tf "$archive")

  return 1
}

resolve_release_branch() {
  local release="$1"
  printf 'v%s\n' "${release%.*}"
}

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
}

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

  kernel="$(find_latest_matching "$OUT_DIR/boot" 'vmlinux-*' || true)"
  if [[ -z "$kernel" ]]; then
    kernel="$(find_latest_matching "$OUT_DIR/boot" 'vmlinuz-*' || true)"
  fi
  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 Alpine 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"
}

cleanup() {
  if [[ "${MODLOOP_MOUNTED:-0}" == "1" ]] && [[ -n "${MODLOOP_MOUNT:-}" ]]; then
    sudo umount "$MODLOOP_MOUNT" || true
  fi
  if [[ -n "${TMP_DIR:-}" && -d "${TMP_DIR:-}" ]]; then
    rm -rf "$TMP_DIR"
  fi
}

SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
REPO_ROOT="$(cd "$SCRIPT_DIR/.." && pwd)"
MANUAL_DIR="${BANGER_MANUAL_DIR:-$REPO_ROOT/build/manual}"
OUT_DIR="$MANUAL_DIR/alpine-kernel"
RELEASE="${ALPINE_RELEASE:-3.23.3}"
MIRROR="https://dl-cdn.alpinelinux.org/alpine"
ARCH="x86_64"
PRINT_REGISTER_FLAGS=0

while [[ $# -gt 0 ]]; do
  case "$1" in
    --out-dir)
      OUT_DIR="${2:-}"
      shift 2
      ;;
    --release)
      RELEASE="${2:-}"
      shift 2
      ;;
    --mirror)
      MIRROR="${2:-}"
      shift 2
      ;;
    --arch)
      ARCH="${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

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

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

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

if command -v unsquashfs >/dev/null 2>&1; then
  USE_UNSQUASHFS=1
else
  USE_UNSQUASHFS=0
  require_command sudo
  require_command mount
  require_command umount
fi

TMP_DIR="$(mktemp -d -t banger-alpine-kernel-XXXXXX)"
EXTRACT_DIR="$TMP_DIR/extract"
MODLOOP_DIR="$TMP_DIR/modloop"
MODLOOP_MOUNT="$TMP_DIR/modloop.mount"
ARCHIVE="$TMP_DIR/alpine-netboot.tar.gz"
MODLOOP_MOUNTED=0
trap cleanup EXIT

mkdir -p "$EXTRACT_DIR" "$MODLOOP_DIR" "$MODLOOP_MOUNT"

BRANCH="$(resolve_release_branch "$RELEASE")"
RELEASE_DIR="$MIRROR/$BRANCH/releases/$ARCH"
ARCHIVE_URL="$RELEASE_DIR/alpine-netboot-$RELEASE-$ARCH.tar.gz"
SHA256_URL="$ARCHIVE_URL.sha256"

log "downloading Alpine netboot bundle from $ARCHIVE_URL"
curl -fsSL "$ARCHIVE_URL" -o "$ARCHIVE"
expected_sha="$(curl -fsSL "$SHA256_URL" | awk '{print $1}')"
actual_sha="$(sha256sum "$ARCHIVE" | awk '{print $1}')"
if [[ -z "$expected_sha" ]]; then
  log "failed to read SHA256 from $SHA256_URL"
  exit 1
fi
if [[ "$expected_sha" != "$actual_sha" ]]; then
  log "sha256 mismatch for $ARCHIVE_URL"
  log "expected: $expected_sha"
  log "actual:   $actual_sha"
  exit 1
fi

VMLINUX_ENTRY="$(find_tar_entry "$ARCHIVE" 'vmlinuz-virt' || true)"
INITRD_ENTRY="$(find_tar_entry "$ARCHIVE" 'initramfs-virt' || true)"
MODLOOP_ENTRY="$(find_tar_entry "$ARCHIVE" 'modloop-virt' || true)"
CONFIG_ENTRY="$(find_tar_config_entry "$ARCHIVE" || true)"

if [[ -z "$VMLINUX_ENTRY" || -z "$INITRD_ENTRY" || -z "$MODLOOP_ENTRY" ]]; then
  log "Alpine netboot bundle is missing expected virt boot artifacts"
  exit 1
fi

log "extracting Alpine virt boot artifacts"
tar_args=("$VMLINUX_ENTRY" "$INITRD_ENTRY" "$MODLOOP_ENTRY")
if [[ -n "$CONFIG_ENTRY" ]]; then
  tar_args+=("$CONFIG_ENTRY")
fi
tar -xf "$ARCHIVE" -C "$EXTRACT_DIR" "${tar_args[@]}"

VMLINUX_SRC="$EXTRACT_DIR/$VMLINUX_ENTRY"
INITRD_SRC="$EXTRACT_DIR/$INITRD_ENTRY"
MODLOOP_SRC="$EXTRACT_DIR/$MODLOOP_ENTRY"
CONFIG_SRC=""
if [[ -n "$CONFIG_ENTRY" ]]; then
  CONFIG_SRC="$EXTRACT_DIR/$CONFIG_ENTRY"
fi

if [[ "$USE_UNSQUASHFS" == "1" ]]; then
  log "extracting kernel modules with unsquashfs"
  unsquashfs -f -d "$MODLOOP_DIR" "$MODLOOP_SRC" >/dev/null
else
  log "extracting kernel modules with a read-only loop mount"
  sudo mount -o loop,ro "$MODLOOP_SRC" "$MODLOOP_MOUNT"
  MODLOOP_MOUNTED=1
  cp -a "$MODLOOP_MOUNT/." "$MODLOOP_DIR/"
  sudo umount "$MODLOOP_MOUNT"
  MODLOOP_MOUNTED=0
fi

MODULES_ROOT=""
if [[ -d "$MODLOOP_DIR/modules" ]]; then
  MODULES_ROOT="$MODLOOP_DIR/modules"
elif [[ -d "$MODLOOP_DIR/lib/modules" ]]; then
  MODULES_ROOT="$MODLOOP_DIR/lib/modules"
fi
if [[ -z "$MODULES_ROOT" ]]; then
  log "extracted modloop is missing a modules directory"
  exit 1
fi

MODULES_SRC="$(find_latest_module_dir "$MODULES_ROOT" || true)"
if [[ -z "$MODULES_SRC" ]]; then
  log "failed to locate a kernel modules tree inside modloop-virt"
  exit 1
fi

KERNEL_VERSION="$(basename "$MODULES_SRC")"
mkdir -p "$OUT_DIR/boot" "$OUT_DIR/lib/modules"
install -m 0644 "$VMLINUX_SRC" "$OUT_DIR/boot/vmlinuz-$KERNEL_VERSION"
install -m 0644 "$INITRD_SRC" "$OUT_DIR/boot/initramfs-$KERNEL_VERSION.img"
if [[ -n "$CONFIG_SRC" && -f "$CONFIG_SRC" ]]; then
  install -m 0644 "$CONFIG_SRC" "$OUT_DIR/boot/config-$KERNEL_VERSION"
fi
cp -a "$MODULES_SRC" "$OUT_DIR/lib/modules/"

log "extracting Firecracker kernel from vmlinuz-$KERNEL_VERSION"
if ! extract_vmlinux "$VMLINUX_SRC" "$OUT_DIR/boot/vmlinux-$KERNEL_VERSION"; then
  log "failed to extract an uncompressed vmlinux from $VMLINUX_SRC"
  log "raw kernel image type: $(file -b "$VMLINUX_SRC")"
  exit 1
fi

log "staged Alpine kernel artifacts in $OUT_DIR"
log "kernel version: $KERNEL_VERSION"