#!/usr/bin/env bash
# Build provisioning initramfs for Raspberry Pi 4 / CM4 (aarch64).
# Produces initrd.img (gzip cpio) for TFTP boot (config.txt: initramfs initrd.img followkernel).
#
# On x86_64: tries Docker/Podman with --platform linux/arm64; if that fails (no
# arm64 emulation), downloads prebuilt static aarch64 busybox and curl. No sudo needed.
# On aarch64 (e.g. Raspberry Pi): uses local busybox and curl if installed.

set -e
SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
OUTPUT="${1:-$SCRIPT_DIR/initrd.img}"
BUILD_DIR=$(mktemp -d)
trap "rm -rf $BUILD_DIR" EXIT

echo "Build dir: $BUILD_DIR"

# Layout: /init, /provisioning-client.sh, /bin/busybox, /bin/sh, /usr/bin/curl, /lib/*.so
mkdir -p "$BUILD_DIR"/{bin,usr/bin,proc,sys,dev,dev/pts,lib,mnt}
cp "$SCRIPT_DIR/init" "$BUILD_DIR/init"
cp "$SCRIPT_DIR/provisioning-client.sh" "$BUILD_DIR/provisioning-client.sh"
cp "$SCRIPT_DIR/rescue-eeprom.sh" "$BUILD_DIR/rescue-eeprom.sh"
chmod +x "$BUILD_DIR/init" "$BUILD_DIR/provisioning-client.sh" "$BUILD_DIR/rescue-eeprom.sh"

ARCH=$(uname -m 2>/dev/null)
if [ "$ARCH" = "aarch64" ] || [ "$ARCH" = "arm64" ] || [ "$ARCH" = "armv8l" ]; then
  if ! command -v busybox >/dev/null 2>&1 || ! command -v curl >/dev/null 2>&1; then
    echo "On aarch64: install busybox and curl (apt install busybox curl) then re-run."
    exit 1
  fi
  echo "Copying busybox and curl from host (aarch64)..."
  cp "$(command -v busybox)" "$BUILD_DIR/bin/busybox"
  cp "$(command -v curl)" "$BUILD_DIR/usr/bin/curl"
  chmod +x "$BUILD_DIR/bin/busybox" "$BUILD_DIR/usr/bin/curl"
  for f in $(ldd "$(command -v curl)" 2>/dev/null | awk '/=>/{print $3}'); do
    [ -f "$f" ] && cp "$f" "$BUILD_DIR/lib/"
  done
  cp /lib/ld-linux-aarch64.so.1 "$BUILD_DIR/lib/" 2>/dev/null || true
else
  # x86_64: try container first; on "Exec format error" (no arm64 emulation) use static downloads
  CONTAINER_RUNTIME=""
  for cmd in docker podman; do
    if command -v "$cmd" >/dev/null 2>&1; then
      CONTAINER_RUNTIME="$cmd"
      break
    fi
  done

  CONTAINER_OK=0
  if [ -n "$CONTAINER_RUNTIME" ]; then
    echo "Trying $CONTAINER_RUNTIME (linux/arm64)..."
    CNT_NAME="cm4-initramfs-build-$$"
    # Use a container-internal dir (no bind mount); copy out with podman cp (works with rootless)
    $CONTAINER_RUNTIME run --name "$CNT_NAME" --platform linux/arm64 debian:bookworm-slim bash -c '
      apt-get update -qq && apt-get install -y -qq busybox curl
      mkdir -p /out/bin /out/usr/bin /out/lib
      cp /bin/busybox /out/bin/busybox
      cp /usr/bin/curl /out/usr/bin/curl
      chmod +x /out/bin/busybox /out/usr/bin/curl
      ldd /usr/bin/curl 2>/dev/null | awk "/=>/{print \$3}" | while read f; do [ -n "$f" ] && [ -f "$f" ] && cp "$f" /out/lib/; done
      cp /lib/ld-linux-aarch64.so.1 /out/lib/ 2>/dev/null || true
    ' 2>/dev/null || true
    # Always copy from container (avoids rootless bind-mount issues)
    $CONTAINER_RUNTIME cp "$CNT_NAME:/out/bin/busybox" "$BUILD_DIR/bin/" 2>/dev/null && \
    $CONTAINER_RUNTIME cp "$CNT_NAME:/out/usr/bin/curl" "$BUILD_DIR/usr/bin/" 2>/dev/null && \
    $CONTAINER_RUNTIME cp "$CNT_NAME:/out/lib/." "$BUILD_DIR/lib/" 2>/dev/null || true
    $CONTAINER_RUNTIME rm -f "$CNT_NAME" >/dev/null 2>&1
    if [ -f "$BUILD_DIR/bin/busybox" ] && [ -f "$BUILD_DIR/usr/bin/curl" ]; then
      CONTAINER_OK=1
      echo "Container build succeeded."
    fi
  fi

  if [ "$CONTAINER_OK" -ne 1 ]; then
    echo "Using prebuilt static aarch64 binaries (no container/emulation needed)..."
    DOWNLOAD_DIR=$(mktemp -d)
    trap "rm -rf $BUILD_DIR $DOWNLOAD_DIR" EXIT
    if command -v curl >/dev/null 2>&1; then
      GET="curl -sL"
      GET_O="curl -sL -o"
    else
      GET="wget -q -O -"
      GET_O="wget -q -O"
    fi
    # Static busybox aarch64: try busybox.net, then Alpine busybox-static package
    BB_OK=0
    $GET_O "$DOWNLOAD_DIR/busybox" "https://busybox.net/downloads/binaries/1.35.0-defconfig-multiarch-musl/busybox-armv8l" 2>/dev/null || true
    if [ -f "$DOWNLOAD_DIR/busybox" ] && [ -s "$DOWNLOAD_DIR/busybox" ]; then
      BB_OK=1
    fi
    if [ "$BB_OK" -ne 1 ]; then
      echo "Trying Alpine busybox-static aarch64..."
      $GET_O "$DOWNLOAD_DIR/bb.apk" "https://dl-cdn.alpinelinux.org/alpine/v3.19/main/aarch64/busybox-static-1.36.1-r11.apk" 2>/dev/null || true
      if [ -f "$DOWNLOAD_DIR/bb.apk" ] && [ -s "$DOWNLOAD_DIR/bb.apk" ]; then
        case "$(head -c 4 "$DOWNLOAD_DIR/bb.apk" | od -An -tx1 2>/dev/null | tr -d ' ')" in
          3c21444f|3c68746d) ;; # HTML response, skip extract
          *)
            (cd "$DOWNLOAD_DIR" && (tar xf bb.apk 2>/dev/null || tar xzf bb.apk 2>/dev/null) && [ -f data.tar.gz ] && tar xzf data.tar.gz 2>/dev/null)
            ;;
        esac
      fi
      if [ -f "$DOWNLOAD_DIR/bin/busybox" ] && [ -s "$DOWNLOAD_DIR/bin/busybox" ]; then
        cp "$DOWNLOAD_DIR/bin/busybox" "$DOWNLOAD_DIR/busybox"
        BB_OK=1
      fi
    fi
    if [ "$BB_OK" -ne 1 ]; then
      echo "Failed to download busybox (x86 host cannot run arm64 container without emulation)."
      echo ""
      echo "Option A - Enable arm64 containers (one-time, needs sudo):"
      echo "  Fedora: sudo dnf install -y qemu-user-static"
      echo "  Then re-run this script (Podman will use QEMU to run the arm64 build)."
      echo ""
      echo "Option B - Build on a Raspberry Pi 4 (aarch64):"
      echo "  scp -r $(dirname "$SCRIPT_DIR") pi@<pi-ip>:~/  && ssh pi@<pi-ip> 'cd ~/emmc-provisioning/network-boot-initramfs && sudo apt install -y busybox curl && ./build.sh'"
      echo "  Then scp pi@<pi-ip>:~/emmc-provisioning/network-boot-initramfs/initrd.img ."
      exit 1
    fi
    chmod +x "$DOWNLOAD_DIR/busybox"
    cp "$DOWNLOAD_DIR/busybox" "$BUILD_DIR/bin/busybox"
    # Static curl aarch64 glibc (Raspberry Pi OS uses glibc)
    $GET "https://github.com/stunnel/static-curl/releases/download/8.18.0/curl-linux-aarch64-glibc-8.18.0.tar.xz" -o "$DOWNLOAD_DIR/curl.tar.xz" || true
    if [ ! -f "$DOWNLOAD_DIR/curl.tar.xz" ] || [ ! -s "$DOWNLOAD_DIR/curl.tar.xz" ]; then
      echo "Failed to download static curl."
      exit 1
    fi
    (cd "$DOWNLOAD_DIR" && tar xf curl.tar.xz)
    CURL_BIN=$(find "$DOWNLOAD_DIR" -maxdepth 3 -name "curl" -type f 2>/dev/null | head -1)
    if [ -n "$CURL_BIN" ] && [ -x "$CURL_BIN" ]; then
      cp "$CURL_BIN" "$BUILD_DIR/usr/bin/curl"
      chmod +x "$BUILD_DIR/usr/bin/curl"
    else
      echo "Could not find curl binary in tarball."
      exit 1
    fi
    rm -rf "$DOWNLOAD_DIR"
    trap "rm -rf $BUILD_DIR" EXIT
  fi
fi

# Verify we have busybox (container or fallback must have left it)
if [ ! -f "$BUILD_DIR/bin/busybox" ] || [ ! -s "$BUILD_DIR/bin/busybox" ]; then
  echo "Error: busybox not found in $BUILD_DIR/bin. If the container ran, check Podman volume mount."
  exit 1
fi
chmod +x "$BUILD_DIR/bin/busybox" 2>/dev/null || true

# Busybox applets we need (sh, mount, udhcpc, etc.)
cd "$BUILD_DIR/bin"
./busybox --list 2>/dev/null | while read applet; do
  case "$applet" in
    sh|ash|mount|umount|mkdir|cat|ip|udhcpc|sleep|echo|grep|cut|awk|hostname|dd|reboot) ln -sf busybox "$applet"; ;;
  esac
done
[ -e sh ] || ln -sf busybox sh

# Build cpio (gzip)
echo "Building cpio..."
( cd "$BUILD_DIR"; find . -print0 | cpio -o -H newc -0 2>/dev/null ) | gzip -9 > "$OUTPUT"
echo "Written: $OUTPUT ($(stat -c%s "$OUTPUT" 2>/dev/null || stat -f%z "$OUTPUT" 2>/dev/null) bytes)"
echo ""
echo "Next: copy initrd.img to TFTP root (e.g. /srv/tftpboot on LXC) and in config.txt add:"
echo "  initramfs initrd.img followkernel"
echo "Then ensure the kernel line loads the initrd (followkernel does that)."
echo "Default provisioning server: http://10.20.50.1:5000 (override with kernel cmdline: provisioning_server=http://...)"