#!/bin/bash
# Rebuild kernel (and required prior stages) after zinit config or init scripts change
# Usage:
#   scripts/rebuild-after-zinit.sh                          # minimal rebuild of initramfs only (skip boot tests; no kernel rebuild)
#   scripts/rebuild-after-zinit.sh --run-tests              # include boot tests (still no kernel rebuild by default)
#   scripts/rebuild-after-zinit.sh --with-kernel            # also rebuild kernel (re-embed updated initramfs)
#   scripts/rebuild-after-zinit.sh --refresh-container-mods # rebuild container /lib/modules if missing (kernel modules stage)
#   scripts/rebuild-after-zinit.sh -- ...                   # pass extra args to build.sh

set -euo pipefail

SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
PROJECT_ROOT="$(dirname "$SCRIPT_DIR")"
STAGES_DIR="${PROJECT_ROOT}/.build-stages"

log() { echo "[rebuild-zinit] $*"; }

run_tests=0
extra_args=()
rebuild_kernel=0
refresh_container_mods=0

# Parse flags; pass through any remaining args to build.sh after --
while [[ $# -gt 0 ]]; do
  case "$1" in
    --run-tests)
      run_tests=1
      shift
      ;;
    --with-kernel)
      rebuild_kernel=1
      shift
      ;;
    --refresh-container-mods)
      refresh_container_mods=1
      shift
      ;;
    --)
      shift
      extra_args=("$@")
      break
      ;;
    *)
      # Pass unknown args through to build.sh
      extra_args+=("$1")
      shift
      ;;
  esac
done

log "PROJECT_ROOT=${PROJECT_ROOT}"
log "STAGES_DIR=${STAGES_DIR}"

# Minimal set of stages to clear when zinit changes:
# - zinit_setup: recopy zinit YAML and init scripts into initramfs
# - validation: re-check initramfs contents
# - initramfs_create: recreate archive including updated zinit files
# - initramfs_test: re-test archive
# - kernel_build: re-embed updated initramfs into kernel
# - boot_tests: optional, depends on --run-tests
stages_to_clear=(
  # Ensure new/changed module selections in [config/modules.conf](config/modules.conf)
  # are re-resolved and copied into the initramfs:
  "modules_setup"
  "modules_copy"

  # Re-copy /init into the initramfs root from [config/init](config/init)
  "init_script"

  # Ensure zinit YAML/init script changes are reapplied:
  "zinit_setup"

  # Re-validate and recreate archive (no kernel rebuild by default):
  "validation"
  "initramfs_create"
  "initramfs_test"
)

# Optionally rebuild container modules if requested (fresh container scenario)
if [[ "$refresh_container_mods" -eq 1 ]]; then
  stages_to_clear=("kernel_modules" "${stages_to_clear[@]}")
fi

# Optionally rebuild kernel (re-embed updated initramfs)
if [[ "$rebuild_kernel" -eq 1 ]]; then
  stages_to_clear+=("kernel_build")
fi

# Remove completion markers to force incremental rebuild of those stages
for s in "${stages_to_clear[@]}"; do
  marker="${STAGES_DIR}/${s}.done"
  if [[ -f "$marker" ]]; then
    log "Removing stage marker: ${marker}"
    rm -f "$marker"
  else
    log "Marker not present (already pending): ${s}"
  fi
done

# Build
log "Starting incremental rebuild (zinit changes)"
# If we plan to rebuild the kernel, force the pipeline to run from initramfs_create
# so the cpio archive is recreated before kernel_build (ignoring prior .done markers).
build_from_args=()
if [[ "$rebuild_kernel" -eq 1 ]]; then
  build_from_args=(--rebuild-from=initramfs_create)
  log "Rebuild-from: initramfs_create (guarantee cpio is recreated before kernel_build)"
fi

if [[ "$run_tests" -eq 1 ]]; then
  log "Including boot tests"
  DEBUG=1 "${PROJECT_ROOT}/scripts/build.sh" "${build_from_args[@]}" "${extra_args[@]}"
else
  log "Skipping boot tests (use --run-tests to include)"
  DEBUG=1 "${PROJECT_ROOT}/scripts/build.sh" --skip-tests "${build_from_args[@]}" "${extra_args[@]}"
fi