#!/usr/bin/env bash
# ───────────────────────────────────────────────────────────
# build-qemu.sh — Dockerfile → Bootable QEMU Disk
#
# Pipeline:
#   1. Build the Docker image (ansible-node)
#   2. Export the container rootfs as a tarball
#   3. Extract kernel + initramfs for direct -kernel boot
#   4. Create an ext4 disk image, label ANSIBLE_ROOT
#   5. Populate with rootfs contents
#   6. Optionally convert raw → qcow2 (if qemu-img available)
#
# Output (written to ./output/):
#   ansible-node.qcow2 (or .raw)  — root filesystem disk
#   vmlinuz-virt                   — Linux kernel
#   initramfs-virt                 — initramfs
#
# Prerequisites:
#   docker, sudo, mkfs.ext4, optional: qemu-img
#
# Usage:
#   ./scripts/build-qemu.sh                    # default 2GB
#   DISK_SIZE_MB=4096 ./scripts/build-qemu.sh  # custom size
# ───────────────────────────────────────────────────────────
set -euo pipefail

SCRIPT_DIR="$(cd "$(dirname "$0")" && pwd)"
PROJECT_DIR="$(dirname "$SCRIPT_DIR")"
OUTPUT_DIR="$PROJECT_DIR/output"
IMAGE_NAME="${IMAGE_NAME:-ansible-node}"
DISK_SIZE_MB="${DISK_SIZE_MB:-2048}"

# ── Colour helpers ──────────────────────────────────────────
RED='\033[0;31m'; GREEN='\033[0;32m'; YELLOW='\033[1;33m'
BLUE='\033[0;34m'; BOLD='\033[1m'; NC='\033[0m'
info()  { echo -e "${BLUE}[*]${NC} $*"; }
ok()    { echo -e "${GREEN}[✓]${NC} $*"; }
warn()  { echo -e "${YELLOW}[!]${NC} $*"; }
err()   { echo -e "${RED}[✗]${NC} $*"; }

# ── Preflight checks ────────────────────────────────────────
for cmd in docker sudo mkfs.ext4; do
    if ! command -v "$cmd" &>/dev/null; then
        err "Missing required tool: $cmd"
        exit 1
    fi
done

# ── Banner ──────────────────────────────────────────────────
echo ""
echo -e "${BOLD}══════════════════════════════════════════════${NC}"
echo -e "${BOLD}  Ansible Control Node — QEMU Image Builder${NC}"
echo -e "${BOLD}══════════════════════════════════════════════${NC}"
echo ""
info "Image name:   ${IMAGE_NAME}"
info "Disk size:    ${DISK_SIZE_MB}MB"
info "Output dir:   ${OUTPUT_DIR}"
echo ""

mkdir -p "$OUTPUT_DIR"

# ── Step 1: Build Docker image ──────────────────────────────
info "Step 1/6: Building Docker image '${IMAGE_NAME}'..."
docker build \
    -t "$IMAGE_NAME" \
    -f "$PROJECT_DIR/Dockerfile" \
    "$PROJECT_DIR"
ok "Docker image built"

# ── Step 2: Export rootfs ───────────────────────────────────
info "Step 2/6: Exporting container rootfs..."
ROOTFS_TAR="$PROJECT_DIR/.ansible-node-rootfs.tar"
CID=$(docker create "$IMAGE_NAME")
docker export "$CID" -o "$ROOTFS_TAR"
docker rm "$CID" >/dev/null
ROOTFS_SIZE=$(du -sh "$ROOTFS_TAR" | cut -f1)
ok "Rootfs exported (${ROOTFS_SIZE})"

# ── Step 3: Extract kernel + initramfs ──────────────────────
info "Step 3/6: Extracting kernel and initramfs..."
TMP_BOOT="$(mktemp -d)"
tar -xf "$ROOTFS_TAR" -C "$TMP_BOOT" boot/ 2>/dev/null

# Find kernel/initramfs (handle different naming patterns)
KERNEL_SRC=$(find "$TMP_BOOT/boot" -name 'vmlinuz-*' 2>/dev/null | head -1)
INITRD_SRC=$(find "$TMP_BOOT/boot" -name 'initramfs-*' 2>/dev/null | head -1)

if [ -z "$KERNEL_SRC" ] || [ -z "$INITRD_SRC" ]; then
    err "Could not find kernel/initramfs in rootfs."
    err "Expected files in /boot/ from linux-virt package."
    ls -la "$TMP_BOOT/boot/" 2>/dev/null || echo "(no /boot directory)"
    rm -rf "$TMP_BOOT" "$ROOTFS_TAR"
    exit 1
fi

KERNEL_NAME=$(basename "$KERNEL_SRC")
INITRD_NAME=$(basename "$INITRD_SRC")

cp "$KERNEL_SRC" "$OUTPUT_DIR/vmlinuz-virt"
cp "$INITRD_SRC" "$OUTPUT_DIR/initramfs-virt"
rm -rf "$TMP_BOOT"

ok "Kernel:   ${KERNEL_NAME}"
ok "Initrd:   ${INITRD_NAME}"

# ── Step 4: Create raw disk image ───────────────────────────
info "Step 4/6: Creating disk image (${DISK_SIZE_MB}MB)..."
RAW_DISK="$OUTPUT_DIR/ansible-node.raw"
dd if=/dev/zero of="$RAW_DISK" bs=1M count="$DISK_SIZE_MB" status=progress 2>/dev/null
mkfs.ext4 -q -L ANSIBLE_ROOT "$RAW_DISK"
ok "ext4 filesystem created (label: ANSIBLE_ROOT)"

# ── Step 5: Mount and populate rootfs ───────────────────────
info "Step 5/6: Populating root filesystem..."
MNT="$(mktemp -d)"
sudo mount -o loop "$RAW_DISK" "$MNT"
sudo tar -xf "$ROOTFS_TAR" -C "$MNT"
sudo umount "$MNT"
rmdir "$MNT"
ok "Rootfs written to disk"

# ── Step 6: Convert to qcow2 (optional) ─────────────────────
info "Step 6/6: Finalizing..."
if command -v qemu-img &>/dev/null; then
    QCOW2_DISK="$OUTPUT_DIR/ansible-node.qcow2"
    qemu-img convert -f raw -O qcow2 "$RAW_DISK" "$QCOW2_DISK"
    rm "$RAW_DISK"
    ok "Converted to qcow2: ansible-node.qcow2"
    FINAL_DISK="$QCOW2_DISK"
    FINAL_FMT="qcow2"
else
    warn "qemu-img not found — keeping raw image"
    ok "Raw image: ansible-node.raw"
    FINAL_DISK="$RAW_DISK"
    FINAL_FMT="raw"
fi

# ── Cleanup ─────────────────────────────────────────────────
rm -f "$ROOTFS_TAR"

# ── Summary ─────────────────────────────────────────────────
echo ""
echo -e "${BOLD}══════════════════════════════════════════════${NC}"
echo -e "${GREEN}${BOLD}  Build complete!${NC}"
echo -e "${BOLD}══════════════════════════════════════════════${NC}"
echo ""
echo -e "  Disk:     ${BOLD}${FINAL_DISK}${NC} (${FINAL_FMT})"
echo -e "  Kernel:   ${BOLD}${OUTPUT_DIR}/vmlinuz-virt${NC}"
echo -e "  Initrd:   ${BOLD}${OUTPUT_DIR}/initramfs-virt${NC}"
echo ""
echo -e "  Launch:   ${BOLD}./scripts/run-qemu.sh${NC}"
echo ""
ls -lh "$OUTPUT_DIR/"
echo ""