#!/bin/bash
#
# prepare-boot-tools.sh — Download/extract boot files for PXE boot tools
#
# Downloads Clonezilla Live and Memtest86+ for PXE booting,
# and extracts Blancco Drive Eraser from its ISO.
#
# Usage:
#   ./prepare-boot-tools.sh [/path/to/blancco.iso]
#
# Output directories:
#   boot-tools/clonezilla/   — vmlinuz, initrd.img, filesystem.squashfs
#   boot-tools/blancco/      — extracted boot files or ISO for memdisk
#   boot-tools/memtest/      — memtest.efi

set -euo pipefail

SCRIPT_DIR="$(cd "$(dirname "$0")" && pwd)"
OUT_DIR="$SCRIPT_DIR/boot-tools"
BLANCCO_ISO="${1:-}"

# Auto-detect Blancco ISO in project directory
if [ -z "$BLANCCO_ISO" ]; then
    BLANCCO_ISO=$(find "$SCRIPT_DIR" -maxdepth 1 -name '*DriveEraser*.iso' -o -name '*blancco*.iso' 2>/dev/null | head -1)
fi

mkdir -p "$OUT_DIR"/{clonezilla,blancco,memtest}

echo "============================================"
echo "PXE Boot Tools Preparation"
echo "============================================"

# --- Clonezilla Live ---
echo ""
echo "[1/3] Clonezilla Live"

CLONEZILLA_VERSION="3.2.1-6"
CLONEZILLA_FILE="clonezilla-live-${CLONEZILLA_VERSION}-amd64.zip"
CLONEZILLA_URL="https://sourceforge.net/projects/clonezilla/files/clonezilla_live_stable/${CLONEZILLA_VERSION}/${CLONEZILLA_FILE}/download"

if [ -f "$OUT_DIR/clonezilla/vmlinuz" ] && [ -f "$OUT_DIR/clonezilla/filesystem.squashfs" ]; then
    echo "    Already prepared, skipping. Delete boot-tools/clonezilla/ to re-download."
else
    echo "    Downloading Clonezilla Live ${CLONEZILLA_VERSION}..."
    TMPDIR=$(mktemp -d)

    wget -q --show-progress -O "$TMPDIR/$CLONEZILLA_FILE" "$CLONEZILLA_URL" || {
        echo "    ERROR: Download failed. Trying alternative URL..."
        # Fallback: try OSDN mirror
        wget -q --show-progress -O "$TMPDIR/$CLONEZILLA_FILE" \
            "https://free.nchc.org.tw/clonezilla-live/stable/${CLONEZILLA_FILE}" || {
            echo "    ERROR: Could not download Clonezilla. Download manually and place in boot-tools/clonezilla/"
            echo "    Need: vmlinuz, initrd.img, filesystem.squashfs from the live ZIP"
        }
    }

    if [ -f "$TMPDIR/$CLONEZILLA_FILE" ]; then
        echo "    Extracting PXE boot files..."
        unzip -o -j "$TMPDIR/$CLONEZILLA_FILE" "live/vmlinuz" -d "$OUT_DIR/clonezilla/"
        unzip -o -j "$TMPDIR/$CLONEZILLA_FILE" "live/initrd.img" -d "$OUT_DIR/clonezilla/"
        unzip -o -j "$TMPDIR/$CLONEZILLA_FILE" "live/filesystem.squashfs" -d "$OUT_DIR/clonezilla/"
        rm -rf "$TMPDIR"
        echo "    Done."
    fi
fi

ls -lh "$OUT_DIR/clonezilla/" 2>/dev/null | grep -E 'vmlinuz|initrd|squashfs' | sed 's/^/    /'

# --- Blancco Drive Eraser ---
echo ""
echo "[2/3] Blancco Drive Eraser"

if [ -n "$BLANCCO_ISO" ] && [ -f "$BLANCCO_ISO" ]; then
    echo "    Extracting from: $BLANCCO_ISO"
    echo "    Using 7z to extract (no root required)..."

    # Blancco is Arch Linux-based. We need:
    #   arch/boot/x86_64/vmlinuz-bde-linux
    #   arch/boot/x86_64/initramfs-bde-linux.img
    #   arch/boot/intel-ucode.img
    #   arch/boot/amd-ucode.img
    #   arch/boot/config.img
    #   arch/x86_64/airootfs.sfs
    TMPDIR=$(mktemp -d)
    7z x -o"$TMPDIR" "$BLANCCO_ISO" \
        "arch/boot/x86_64/vmlinuz-bde-linux" \
        "arch/boot/x86_64/initramfs-bde-linux.img" \
        "arch/boot/intel-ucode.img" \
        "arch/boot/amd-ucode.img" \
        "arch/boot/config.img" \
        "arch/x86_64/airootfs.sfs" \
        -r 2>/dev/null || {
        echo "    7z extraction failed. Install p7zip-full: apt install p7zip-full"
    }

    # Flatten into blancco/ directory for HTTP serving
    if [ -f "$TMPDIR/arch/boot/x86_64/vmlinuz-bde-linux" ]; then
        cp "$TMPDIR/arch/boot/x86_64/vmlinuz-bde-linux" "$OUT_DIR/blancco/"
        cp "$TMPDIR/arch/boot/x86_64/initramfs-bde-linux.img" "$OUT_DIR/blancco/"
        cp "$TMPDIR/arch/boot/intel-ucode.img" "$OUT_DIR/blancco/"
        cp "$TMPDIR/arch/boot/amd-ucode.img" "$OUT_DIR/blancco/"
        cp "$TMPDIR/arch/boot/config.img" "$OUT_DIR/blancco/"
        # airootfs.sfs needs to be in arch/x86_64/ path relative to HTTP root
        mkdir -p "$OUT_DIR/blancco/arch/x86_64"
        cp "$TMPDIR/arch/x86_64/airootfs.sfs" "$OUT_DIR/blancco/arch/x86_64/"
        echo "    Extracted Blancco boot files."
    else
        echo "    Could not extract boot files from ISO."
    fi
    rm -rf "$TMPDIR"
else
    echo "    No Blancco ISO found. Provide path as argument or place in project directory."
    echo "    Usage: $0 /path/to/DriveEraser.iso"
fi

ls -lh "$OUT_DIR/blancco/" 2>/dev/null | grep -v '^total' | sed 's/^/    /'

# --- Memtest86+ ---
echo ""
echo "[3/3] Memtest86+"

MEMTEST_VERSION="7.20"
MEMTEST_URL="https://memtest.org/download/${MEMTEST_VERSION}/mt86plus_${MEMTEST_VERSION}.binaries.zip"

if [ -f "$OUT_DIR/memtest/memtest.efi" ]; then
    echo "    Already prepared, skipping."
else
    echo "    Downloading Memtest86+ v${MEMTEST_VERSION}..."
    TMPDIR=$(mktemp -d)

    wget -q --show-progress -O "$TMPDIR/memtest.zip" "$MEMTEST_URL" || {
        echo "    ERROR: Download failed. Download manually from https://memtest.org"
        TMPDIR=""
    }

    if [ -n "$TMPDIR" ] && [ -f "$TMPDIR/memtest.zip" ]; then
        echo "    Extracting EFI binary..."
        unzip -o -j "$TMPDIR/memtest.zip" "memtest64.efi" -d "$OUT_DIR/memtest/" 2>/dev/null || \
        unzip -o -j "$TMPDIR/memtest.zip" "mt86plus_${MEMTEST_VERSION}.x64.efi" -d "$OUT_DIR/memtest/" 2>/dev/null || \
        unzip -o "$TMPDIR/memtest.zip" -d "$TMPDIR/extract/"

        # Find the EFI file regardless of exact name
        EFI_FILE=$(find "$TMPDIR" "$OUT_DIR/memtest" -name '*.efi' -name '*64*' 2>/dev/null | head -1)
        if [ -n "$EFI_FILE" ] && [ ! -f "$OUT_DIR/memtest/memtest.efi" ]; then
            cp "$EFI_FILE" "$OUT_DIR/memtest/memtest.efi"
        fi
        rm -rf "$TMPDIR"
        echo "    Done."
    fi
fi

ls -lh "$OUT_DIR/memtest/" 2>/dev/null | grep -v '^total' | sed 's/^/    /'

# --- Summary ---
echo ""
echo "============================================"
echo "Boot tools prepared in: $OUT_DIR/"
echo "============================================"
echo ""

for tool in clonezilla blancco memtest; do
    COUNT=$(find "$OUT_DIR/$tool" -type f 2>/dev/null | wc -l)
    SIZE=$(du -sh "$OUT_DIR/$tool" 2>/dev/null | cut -f1)
    printf "  %-15s %s (%d files)\n" "$tool" "$SIZE" "$COUNT"
done

echo ""
echo "These files need to be copied to the PXE server's web root:"
echo "  /var/www/html/clonezilla/"
echo "  /var/www/html/blancco/"
echo "  /var/www/html/memtest/"
echo ""
echo "The build-usb.sh script will include them automatically,"
echo "or copy them manually to the server."
echo ""