#!/usr/bin/env bash
# ==============================================================================
# Hermes Portable Rescue — Custom Live ISO Builder
# ==============================================================================
#
# Builds a custom Xubuntu 24.04 LTS live ISO with Hermes Agent pre-installed,
# diagnostic tools, autorun setup, and persistent storage support.
#
# Usage:
#   sudo ./build/build-custom-iso.sh [options]
#
# Options:
#   --input-iso PATH       Path to Xubuntu 24.04 ISO (default: auto-detect)
#   --output-dir DIR       Output directory (default: ./build/output)
#   --work-dir DIR         Working directory for build (default: /tmp/hermes-iso-build)
#   --skip-download        Don't download ISO if missing
#   --keep-workdir         Don't clean up workdir after build
#   --help                 Show this message
#
# Requirements:
#   - squashfs-tools (mksquashfs, unsquashfs)
#   - xorriso
#   - syslinux-common, isolinux
#   - sudo access for mount operations
#
# ==============================================================================

set -euo pipefail

SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
PROJECT_DIR="$(cd "$SCRIPT_DIR/.." && pwd)"
BUILD_DIR="$PROJECT_DIR/build"

# ─── Defaults ──────────────────────────────────────────────────────────────────

DEFAULT_INPUT_ISO="$BUILD_DIR/xubuntu-24.04.4-desktop-amd64.iso"
OUTPUT_DIR="${OUTPUT_DIR:-$BUILD_DIR/output}"
WORKDIR="${WORKDIR:-/tmp/hermes-iso-build}"
SKIP_DOWNLOAD=false
KEEP_WORKDIR=false

# ─── Colors ────────────────────────────────────────────────────────────────────

RED='\033[0;31m'; GREEN='\033[0;32m'; YELLOW='\033[1;33m'; BLUE='\033[0;34m'; NC='\033[0m'
info()  { echo -e "${BLUE}[INFO]${NC}  $*"; }
ok()    { echo -e "${GREEN}[OK]${NC}    $*"; }
warn()  { echo -e "${YELLOW}[WARN]${NC}  $*" >&2; }
error() { echo -e "${RED}[ERROR]${NC} $*" >&2; }
die()   { error "$*"; exit 1; }

# ─── Parse Arguments ───────────────────────────────────────────────────────────

INPUT_ISO="$DEFAULT_INPUT_ISO"

while [[ $# -gt 0 ]]; do
    case "$1" in
        --input-iso)    INPUT_ISO="$2"; shift 2 ;;
        --output-dir)   OUTPUT_DIR="$2"; shift 2 ;;
        --work-dir)     WORKDIR="$2"; shift 2 ;;
        --skip-download) SKIP_DOWNLOAD=true; shift ;;
        --keep-workdir)  KEEP_WORKDIR=true; shift ;;
        -h|--help)      head -50 "$0"; exit 0 ;;
        *)              die "Unknown option: $1" ;;
    esac
done

# ─── Prerequisites ─────────────────────────────────────────────────────────────

check_command() {
    if ! command -v "$1" &>/dev/null; then
        die "Required command '$1' not found"
    fi
}

prerequisites() {
    info "Checking prerequisites..."
    check_command unsquashfs
    check_command mksquashfs
    check_command xorriso
    check_command dd
    check_command mount
    check_command rsync
    check_command chroot
    ok "All prerequisites available."
}

# ─── Phase 1: Download / Verify ISO ────────────────────────────────────────────

phase_download_iso() {
    if [[ -f "$INPUT_ISO" ]]; then
        local iso_size
        iso_size=$(stat -c%s "$INPUT_ISO" 2>/dev/null || echo 0)
        if [[ "$iso_size" -gt 1000000000 ]]; then
            ok "Input ISO found: $INPUT_ISO ($(du -h "$INPUT_ISO" | cut -f1))"
            # Verify SHA256
            local expected_sha="fc2e995bb05c41ea19f1dbfd91f6deea7b2aed7a83b9934d98fc9d9cac527d97"
            local actual_sha
            actual_sha=$(sha256sum "$INPUT_ISO" | awk '{print $1}')
            if [[ "$actual_sha" == "$expected_sha" ]]; then
                ok "SHA256 verified."
            else
                warn "SHA256 mismatch! Expected $expected_sha, got $actual_sha"
                warn "Continuing anyway — ISO may be incomplete or different build."
            fi
            return 0
        fi
    fi

    if $SKIP_DOWNLOAD; then
        die "ISO not found at $INPUT_ISO and --skip-download is set."
    fi

    info "Downloading Xubuntu 24.04.4 LTS Desktop ISO (~4.1 GB)..."
    wget -c "https://cdimage.ubuntu.com/xubuntu/releases/24.04/release/xubuntu-24.04.4-desktop-amd64.iso" \
        -O "$INPUT_ISO" 2>&1
    ok "Download complete."
}

# ─── Phase 2: Extract ISO ─────────────────────────────────────────────────────

phase_extract_iso() {
    info "Extracting ISO contents..."

    local iso_mount="/mnt/hermes-iso-source"
    mkdir -p "$iso_mount" "$WORKDIR/iso-contents"

    # Mount ISO (handle both .iso and .img)
    mount -o loop "$INPUT_ISO" "$iso_mount" 2>/dev/null || \
        die "Failed to mount ISO. Is it a valid ISO file?"

    # List contents for debugging
    info "ISO contents:"
    ls -la "$iso_mount/" 2>/dev/null
    echo ""

    # Determine boot structure
    if [[ -d "$iso_mount/isolinux" ]]; then
        ok "BIOS boot: isolinux detected"
    fi
    if [[ -f "$iso_mount/boot/grub/efi.img" ]] || [[ -f "$iso_mount/boot/grub/eltorito.img" ]]; then
        ok "UEFI boot: GRUB efi.img detected"
    fi

    # Check squashfs location
    local squashfs_src=""
    if [[ -f "$iso_mount/casper/filesystem.squashfs" ]]; then
        squashfs_src="$iso_mount/casper/filesystem.squashfs"
        ok "Squashfs: casper/filesystem.squashfs"
    elif [[ -f "$iso_mount/casper/filesystem.squashfs.xz" ]]; then
        squashfs_src="$iso_mount/casper/filesystem.squashfs.xz"
        ok "Squashfs: casper/filesystem.squashfs.xz (compressed)"
    else
        ls -la "$iso_mount/" "$iso_mount/casper/" 2>/dev/null
        umount "$iso_mount"
        die "Cannot locate filesystem.squashfs in the ISO."
    fi

    # Copy everything EXCEPT the squashfs (we'll replace it)
    info "Copying ISO contents (excluding squashfs)..."
    rsync -a --exclude="casper/filesystem.squashfs*" \
        "$iso_mount/" "$WORKDIR/iso-contents/" 2>&1 | tail -5

    # Copy the squashfs for extraction
    cp "$squashfs_src" "$WORKDIR/filesystem.squashfs"

    # Extract the MBR from the ISO for hybrid boot
    info "Extracting MBR boot sector..."
    dd if="$INPUT_ISO" of="$WORKDIR/iso-contents/isohdpfx.bin" bs=440 count=1 2>/dev/null
    ok "MBR extracted."

    umount "$iso_mount"
    rmdir "$iso_mount"
    ok "ISO contents extracted to $WORKDIR/iso-contents"
}

# ─── Phase 3: Unsquash Filesystem ──────────────────────────────────────────────

phase_unsquash() {
    info "Unsquashing filesystem (this may take a few minutes)..."

    mkdir -p "$WORKDIR/squashfs-root"

    if [[ "$WORKDIR/filesystem.squashfs" == *.xz ]]; then
        xz -d "$WORKDIR/filesystem.squashfs" -c > "$WORKDIR/filesystem.squashfs.unxz" 2>/dev/null
        unsquashfs -d "$WORKDIR/squashfs-root" -f "$WORKDIR/filesystem.squashfs.unxz" 2>&1 | tail -3
    else
        unsquashfs -d "$WORKDIR/squashfs-root" -f "$WORKDIR/filesystem.squashfs" 2>&1 | tail -3
    fi

    local root_size
    root_size=$(du -sh "$WORKDIR/squashfs-root" | cut -f1)
    ok "Filesystem unsquashed: $root_size"
}

# ─── Phase 4: Customize Filesystem (Chroot) ────────────────────────────────────

# This function writes the customization script that runs INSIDE the chroot
generate_chroot_script() {
    cat > "$WORKDIR/chroot-customize.sh" << 'CHROOTSCRIPT'
#!/bin/bash
# ==============================================================================
# Hermes Portable Rescue — Chroot Customization Script
# ==============================================================================
# Runs inside the live filesystem chroot to install everything.
# ==============================================================================

set -euo pipefail

# Redirect output to log
exec > /tmp/chroot-build.log 2>&1

echo "=== Hermes Portable Rescue — Chroot Customization ==="
echo "Date: $(date)"
echo ""

# ─── 1. Mount virtual filesystems ──────────────────────────────────────────────
mount -t proc none /proc
mount -t sysfs none /sys
mount -t devtmpfs none /dev || mount -t tmpfs none /dev
mkdir -p /dev/pts
mount -t devpts none /dev/pts

# ─── 2. Configure package sources ──────────────────────────────────────────────
echo "=== Configuring package sources ==="
apt-get update -qq 2>/dev/null || true

# ─── 3. Install essential packages ─────────────────────────────────────────────
echo "=== Installing packages ==="

PACKAGES=(
    # Python + dev
    python3-pip
    python3-venv
    python3-setuptools
    python3-wheel
    git

    # Diagnostic tools
    smartmontools
    dmidecode
    lshw
    testdisk
    ddrescue
    ntfs-3g
    stress-ng

    # Network tools
    nmap
    net-tools
    curl
    wget
    openssh-client
    iperf3
    traceroute

    # Storage tools
    parted
    rsync
    htop
    iotop
    exfatprogs
    btrfs-progs
    xfsprogs

    # Filesystem tools
    dosfstools
    e2fsprogs

    # Windows-specific tools
    hivex  # Windows registry access
)

apt-get install -y -qq "${PACKAGES[@]}" 2>&1 | tail -5
echo "Package installation complete."

# ─── 4. Install Hermes Agent ───────────────────────────────────────────────────
echo "=== Installing Hermes Agent ==="

# Create Hermes directory structure
mkdir -p /opt/hermes/{agent,tools,scripts,models,diagnostics,config}

# Install from pip
pip3 install --no-cache-dir \
    hermes-agent \
    pyyaml \
    psutil \
    requests \
    2>&1 | tail -3

# Freeze package list for provenance
pip3 freeze > /opt/hermes/pip-freeze.txt

echo "Hermes Agent installed."

# ─── 5. Configure Hermes autorun service ───────────────────────────────────────
echo "=== Configuring autorun ==="

# Systemd service that runs the Hermes rescue autorun on boot
cat > /etc/systemd/system/hermes-autorun.service << 'SERVICEEOF'
[Unit]
Description=Hermes Portable Rescue — Auto-Diagnostic Agent
After=network.target multi-user.target
ConditionPathExists=/opt/hermes/autorun.sh

[Service]
Type=oneshot
ExecStart=/opt/hermes/autorun.sh
RemainAfterExit=yes
StandardOutput=journal+console
StandardError=journal+console
Environment=HERMES_DIR=/opt/hermes
Environment=CONFIG=/opt/hermes/config.yaml
Environment=REPORT_DIR=/opt/hermes/reports
Environment=BOOT_MODE=live

[Install]
WantedBy=multi-user.target
SERVICEEOF

# Desktop autostart for GUI mode — opens a terminal running Hermes
mkdir -p /etc/xdg/autostart
cat > /etc/xdg/autostart/hermes-rescue.desktop << 'DESKTOPDF'
[Desktop Entry]
Type=Application
Name=Hermes Portable Rescue
Comment=AI-driven PC diagnostic & repair toolkit
Exec=xfce4-terminal --title="Hermes Rescue" -e "/opt/hermes/autorun.sh" --hold
Terminal=false
Categories=System;Utility;
StartupNotify=true
X-GNOME-Autostart-enabled=true
DESKTOPDF

# Also create a desktop shortcut on the user desktop for manual launch
mkdir -p /etc/skel/Desktop
cat > /etc/skel/Desktop/hermes-rescue.desktop << 'DESKTOPEOF'
[Desktop Entry]
Type=Application
Name=Hermes Rescue
Comment=AI-driven PC diagnostic & repair toolkit
Exec=xfce4-terminal --title="Hermes Rescue" -e "/opt/hermes/autorun.sh" --hold
Terminal=false
Icon=system-run
Categories=System;Utility;
StartupNotify=true
DESKTOPEOF
chmod +x /etc/skel/Desktop/hermes-rescue.desktop

# Enable the service for live boot
systemctl enable hermes-autorun.service 2>/dev/null || true

echo "Autorun service configured."

# ─── 6. Create symlinks for diagnostic scripts ─────────────────────────────────
echo "=== Setting up diagnostic scripts ==="

# The scripts will be copied from the project source via bind mount.
# Create placeholder structure here — actual files come from the build host.
mkdir -p /opt/hermes/scripts
cat > /opt/hermes/scripts/README.txt << 'README'
Diagnostic scripts are copied from the build host at ISO build time.
Place your scripts in /opt/hermes/scripts/ on the running system.
README

echo "Diagnostic scripts directory ready."

# ─── 7. Configure auto-login (lightdm for XFCE) ────────────────────────────────
echo "=== Configuring auto-login ==="

# LightDM auto-login for the live session user (xubuntu)
mkdir -p /etc/lightdm/lightdm.conf.d
cat > /etc/lightdm/lightdm.conf.d/50-hermes-autologin.conf << 'LIGHTDMEOF'
[Seat:*]
autologin-user=xubuntu
autologin-user-timeout=0
autologin-session=xubuntu
LIGHTDMEOF

echo "Auto-login configured for 'xubuntu'."

# ─── 8. Create persistent storage hook ─────────────────────────────────────────
echo "=== Setting up persistent storage hook ==="

# Script to mount persistent storage if available
mkdir -p /opt/hermes/scripts
cat > /opt/hermes/scripts/mount-persistent.sh << 'MOUNTEOF'
#!/bin/bash
# Mount persistent storage partition (NTFS/exFAT labeled HERMES_DATA)
set -e

PERSISTENT_LABEL="HERMES_DATA"
MOUNT_POINT="/mnt/hermes-data"

mkdir -p "$MOUNT_POINT"

# Find and mount by label
for dev in /dev/sd* /dev/nvme* /dev/mmcblk*; do
    if [ -b "$dev" ]; then
        label=$(blkid -o value -s LABEL "$dev" 2>/dev/null || echo "")
        if [ "$label" = "$PERSISTENT_LABEL" ]; then
            fstype=$(blkid -o value -s TYPE "$dev" 2>/dev/null || echo "")
            case "$fstype" in
                ntfs)  mount -t ntfs-3g "$dev" "$MOUNT_POINT" 2>/dev/null && break ;;
                ext4)  mount "$dev" "$MOUNT_POINT" 2>/dev/null && break ;;
                vfat|exfat) mount "$dev" "$MOUNT_POINT" 2>/dev/null && break ;;
            esac
        fi
    fi
done

if mountpoint -q "$MOUNT_POINT"; then
    echo "Persistent storage mounted at $MOUNT_POINT"
    # Symlink persistent config
    if [ -f "$MOUNT_POINT/hermes/config.yaml" ]; then
        ln -sf "$MOUNT_POINT/hermes/config.yaml" /opt/hermes/config.yaml
    fi
else
    echo "No persistent storage found. Running in temporary mode."
fi
MOUNTEOF
chmod +x /opt/hermes/scripts/mount-persistent.sh

echo "Persistent storage hook created."

# ─── 9. Clean up package caches ─────────────────────────────────────────────────
echo "=== Cleaning up ==="
apt-get clean -qq 2>/dev/null || true
apt-get autoremove --purge -qq 2>/dev/null || true
rm -rf /var/lib/apt/lists/* 2>/dev/null || true
rm -rf /root/.cache 2>/dev/null || true
rm -rf /tmp/* 2>/dev/null || true

# Clean Python caches
find /usr -type d -name '__pycache__' -exec rm -rf {} + 2>/dev/null || true
find /usr -type f -name '*.pyc' -delete 2>/dev/null || true
find /opt/hermes -type d -name '__pycache__' -exec rm -rf {} + 2>/dev/null || true
find /opt/hermes -type f -name '*.pyc' -delete 2>/dev/null || true

echo ""
echo "=== Chroot customization complete! ==="
echo "Packages installed: $(dpkg-query -W -f='${Package}\n' 2>/dev/null | wc -l)"
echo "Hermes config at: /opt/hermes/"
echo "Autorun service: hermes-autorun.service"
echo ""

# Unmount virtual filesystems
umount /dev/pts 2>/dev/null || true
umount /sys 2>/dev/null || true
umount /proc 2>/dev/null || true

exit 0
CHROOTSCRIPT

    chmod +x "$WORKDIR/chroot-customize.sh"
}

phase_customize() {
    info "Preparing chroot environment..."

    SQUASHFS_ROOT="$WORKDIR/squashfs-root"

    # Bind mount necessary filesystems
    mount --bind /dev "$SQUASHFS_ROOT/dev"
    mount --bind /proc "$SQUASHFS_ROOT/proc"
    mount --bind /sys "$SQUASHFS_ROOT/sys"
    mount -t devpts none "$SQUASHFS_ROOT/dev/pts" 2>/dev/null || true

    # Network connectivity inside chroot
    cp /etc/resolv.conf "$SQUASHFS_ROOT/etc/resolv.conf"

    # Copy the project's diagnostic scripts and config into the chroot
    mkdir -p "$SQUASHFS_ROOT/tmp/hermes-project"
    if [[ -d "$PROJECT_DIR/src/diagnostics" ]]; then
        cp -r "$PROJECT_DIR/src/diagnostics/"* "$SQUASHFS_ROOT/tmp/hermes-project/" 2>/dev/null || true
    fi
    if [[ -f "$PROJECT_DIR/src/hermes/config.yaml" ]]; then
        cp "$PROJECT_DIR/src/hermes/config.yaml" "$SQUASHFS_ROOT/tmp/hermes-project/" 2>/dev/null || true
    fi
    if [[ -f "$PROJECT_DIR/src/hermes/autorun.sh" ]]; then
        cp "$PROJECT_DIR/src/hermes/autorun.sh" "$SQUASHFS_ROOT/tmp/hermes-project/" 2>/dev/null || true
    fi
    if [[ -f "$PROJECT_DIR/src/hermes/launch.sh" ]]; then
        cp "$PROJECT_DIR/src/hermes/launch.sh" "$SQUASHFS_ROOT/tmp/hermes-project/" 2>/dev/null || true
    fi
    if [[ -d "$PROJECT_DIR/src/hermes/scripts" ]]; then
        cp -r "$PROJECT_DIR/src/hermes/scripts/"* "$SQUASHFS_ROOT/tmp/hermes-project/scripts/" 2>/dev/null || true
    fi

    # Copy the chroot script in
    cp "$WORKDIR/chroot-customize.sh" "$SQUASHFS_ROOT/tmp/chroot-customize.sh"
    chmod +x "$SQUASHFS_ROOT/tmp/chroot-customize.sh"

    info "Running chroot customization (this may take 10-20 minutes)..."
    echo ""

    if chroot "$SQUASHFS_ROOT" /bin/bash /tmp/chroot-customize.sh; then
        ok "Chroot customization completed successfully."
    else
        warn "Chroot script exited with code $? — checking log..."
    fi

    # Check the build log
    if [[ -f "$SQUASHFS_ROOT/tmp/chroot-build.log" ]]; then
        echo ""
        info "Chroot build log (last 30 lines):"
        tail -30 "$SQUASHFS_ROOT/tmp/chroot-build.log"
    fi

    # Copy project files from temp to /opt/hermes
    info "Installing project files..."
    if [[ -f "$SQUASHFS_ROOT/tmp/hermes-project/config.yaml" ]]; then
        cp "$SQUASHFS_ROOT/tmp/hermes-project/config.yaml" "$SQUASHFS_ROOT/opt/hermes/config.yaml"
    fi
    if [[ -f "$SQUASHFS_ROOT/tmp/hermes-project/autorun.sh" ]]; then
        cp "$SQUASHFS_ROOT/tmp/hermes-project/autorun.sh" "$SQUASHFS_ROOT/opt/hermes/autorun.sh"
        chmod +x "$SQUASHFS_ROOT/opt/hermes/autorun.sh"
    fi
    if [[ -f "$SQUASHFS_ROOT/tmp/hermes-project/launch.sh" ]]; then
        cp "$SQUASHFS_ROOT/tmp/hermes-project/launch.sh" "$SQUASHFS_ROOT/opt/hermes/launch.sh"
        chmod +x "$SQUASHFS_ROOT/opt/hermes/launch.sh"
    fi
    if [[ -d "$SQUASHFS_ROOT/tmp/hermes-project/scripts" ]]; then
        cp -r "$SQUASHFS_ROOT/tmp/hermes-project/scripts/"* "$SQUASHFS_ROOT/opt/hermes/scripts/" 2>/dev/null || true
    fi
    if [[ -d "$SQUASHFS_ROOT/tmp/hermes-project/diagnostics" ]]; then
        cp -r "$SQUASHFS_ROOT/tmp/hermes-project/diagnostics/"* "$SQUASHFS_ROOT/opt/hermes/scripts/" 2>/dev/null || true
    fi
    if [[ -d "$SQUASHFS_ROOT/tmp/hermes-project/__init__.py" ]]; then
        cp "$SQUASHFS_ROOT/tmp/hermes-project/__init__.py" "$SQUASHFS_ROOT/opt/hermes/scripts/" 2>/dev/null || true
    fi
    if [[ -f "$SQUASHFS_ROOT/tmp/hermes-project/__main__.py" ]]; then
        cp "$SQUASHFS_ROOT/tmp/hermes-project/__main__.py" "$SQUASHFS_ROOT/opt/hermes/scripts/" 2>/dev/null || true
    fi

    # Clean up
    rm -rf "$SQUASHFS_ROOT/tmp/hermes-project" "$SQUASHFS_ROOT/tmp/chroot-customize.sh" 2>/dev/null || true
}

# ─── Phase 5: Clean Chroot ─────────────────────────────────────────────────────

phase_clean_chroot() {
    info "Cleaning chroot..."

    SQUASHFS_ROOT="$WORKDIR/squashfs-root"

    # Remove resolv.conf
    rm -f "$SQUASHFS_ROOT/etc/resolv.conf"

    # Unmount virtual filesystems
    umount "$SQUASHFS_ROOT/dev/pts" 2>/dev/null || true
    umount "$SQUASHFS_ROOT/sys" 2>/dev/null || true
    umount "$SQUASHFS_ROOT/proc" 2>/dev/null || true
    umount "$SQUASHFS_ROOT/dev" 2>/dev/null || true

    ok "Chroot cleaned."
}

# ─── Phase 6: Re-pack Squashfs ─────────────────────────────────────────────────

phase_repack_squashfs() {
    info "Re-packing squashfs filesystem (this takes a while)..."

    local target_squashfs="$WORKDIR/iso-contents/casper/filesystem.squashfs"
    mkdir -p "$(dirname "$target_squashfs")"

    # Rebuild with high compression
    mksquashfs "$WORKDIR/squashfs-root" "$target_squashfs" \
        -comp xz -b 1M -Xdict-size 100% \
        -noappend \
        2>&1 | tail -5

    local squash_size
    squash_size=$(du -h "$target_squashfs" | cut -f1)
    ok "Squashfs re-packed: $squash_size"

    # Update filesystem.size for casper
    info "Updating filesystem.size..."
    local fs_size
    fs_size=$(du -sx --block-size=1 "$WORKDIR/squashfs-root" 2>/dev/null | cut -f1)
    echo "$fs_size" > "$WORKDIR/iso-contents/casper/filesystem.size"
    ok "filesystem.size updated: $fs_size bytes"
}

# ─── Phase 7: Generate ISO ─────────────────────────────────────────────────────

phase_generate_iso() {
    info "Generating bootable ISO..."

    mkdir -p "$OUTPUT_DIR"
    local output_iso="$OUTPUT_DIR/hermes-rescue.iso"
    local iso_contents="$WORKDIR/iso-contents"

    # Build the ISO with hybrid MBR (BIOS + UEFI)
    cd "$iso_contents"

    # Determine which efi.img to use
    local efi_img=""
    if [[ -f "boot/grub/efi.img" ]]; then
        efi_img="boot/grub/efi.img"
    elif [[ -f "boot/grub/eltorito.img" ]]; then
        efi_img="boot/grub/eltorito.img"
    else
        # Search for any efi boot image
        efi_img=$(find . -name "efi.img" -o -name "eltorito.img" 2>/dev/null | head -1)
        if [[ -z "$efi_img" ]]; then
            warn "No EFI boot image found — ISO may not boot on UEFI systems."
        fi
    fi

    if [[ -n "$efi_img" ]]; then
        info "Using EFI image: $efi_img"
        xorriso -as mkisofs -r -V "HERMES_RESCUE" \
            -J -joliet-long \
            -isohybrid-mbr "$iso_contents/isohdpfx.bin" \
            -b isolinux/isolinux.bin \
            -c isolinux/boot.cat \
            -no-emul-boot -boot-load-size 4 -boot-info-table \
            -eltorito-alt-boot -e "$efi_img" -no-emul-boot \
            -isohybrid-gpt-basdat \
            -isohybrid-apm-hfsplus \
            -o "$output_iso" \
            "$iso_contents" 2>&1 | tail -10
    else
        # BIOS-only fallback
        warn "Building BIOS-only ISO (no UEFI support)."
        xorriso -as mkisofs -r -V "HERMES_RESCUE" \
            -J -joliet-long \
            -isohybrid-mbr "$iso_contents/isohdpfx.bin" \
            -b isolinux/isolinux.bin \
            -c isolinux/boot.cat \
            -no-emul-boot -boot-load-size 4 -boot-info-table \
            -o "$output_iso" \
            "$iso_contents" 2>&1 | tail -10
    fi

    local iso_size
    iso_size=$(du -h "$output_iso" | cut -f1)
    ok "ISO generated: $output_iso ($iso_size)"

    cd "$PROJECT_DIR"
}

# ─── Phase 8: Create Persistent Storage Overlay ────────────────────────────────

phase_create_persistence() {
    info "Creating persistent storage files..."

    # Create a casper-rw file for persistence (1GB by default)
    local persist_size="${PERSIST_SIZE:-1024}"
    local persist_file="$OUTPUT_DIR/casper-rw"

    if [[ -f "$persist_file" ]]; then
        warn "casper-rw already exists — skipping."
        return 0
    fi

    dd if=/dev/zero of="$persist_file" bs=1M count="$persist_size" status=progress 2>/dev/null
    mkfs.ext4 -F -L casper-rw "$persist_file" 2>/dev/null
    ok "Persistent storage created: $persist_file (${persist_size}MB)"

    # Create a writeable partition image for Windows-accessible data
    local data_img="$OUTPUT_DIR/hermes-data.ext4"
    if [[ ! -f "$data_img" ]]; then
        dd if=/dev/zero of="$data_img" bs=1M count=512 status=progress 2>/dev/null
        mkfs.ext4 -F -L HERMES_DATA "$data_img" 2>/dev/null
        ok "Data partition image: $data_img (512MB)"
    fi
}

# ─── Phase 9: Cleanup ──────────────────────────────────────────────────────────

phase_cleanup() {
    if $KEEP_WORKDIR; then
        info "Keeping workdir at $WORKDIR"
        return 0
    fi

    info "Cleaning up build artifacts..."
    rm -rf "$WORKDIR/squashfs-root" 2>/dev/null || true
    rm -f "$WORKDIR/filesystem.squashfs" "$WORKDIR/filesystem.squashfs.unxz" 2>/dev/null || true
    rm -rf "$WORKDIR/iso-contents" 2>/dev/null || true
    rm -f "$WORKDIR/chroot-customize.sh" 2>/dev/null || true
    ok "Cleanup complete."
}

# ─── Main ──────────────────────────────────────────────────────────────────────

main() {
    echo ""
    echo "╔════════════════════════════════════════════════════════╗"
    echo "║      Hermes Portable Rescue — Live ISO Builder        ║"
    echo "╚════════════════════════════════════════════════════════╝"
    echo ""
    echo "  Input ISO:  $INPUT_ISO"
    echo "  Output:     $OUTPUT_DIR/hermes-rescue.iso"
    echo "  Workdir:    $WORKDIR"
    echo "  Project:    $PROJECT_DIR"
    echo ""

    # Phase count
    TOTAL=8
    CURRENT=0

    CURRENT=$((CURRENT + 1)); echo ""; echo "[${CURRENT}/${TOTAL}] Prerequisites"; prerequisites
    CURRENT=$((CURRENT + 1)); echo ""; echo "[${CURRENT}/${TOTAL}] Download/Verify ISO"; phase_download_iso
    CURRENT=$((CURRENT + 1)); echo ""; echo "[${CURRENT}/${TOTAL}] Extract ISO"; phase_extract_iso
    CURRENT=$((CURRENT + 1)); echo ""; echo "[${CURRENT}/${TOTAL}] Unsquash filesystem"; phase_unsquash
    CURRENT=$((CURRENT + 1)); echo ""; echo "[${CURRENT}/${TOTAL}] Generate chroot script"; generate_chroot_script
    CURRENT=$((CURRENT + 1)); echo ""; echo "[${CURRENT}/${TOTAL}] Customize (chroot)"; phase_customize
    CURRENT=$((CURRENT + 1)); echo ""; echo "[${CURRENT}/${TOTAL}] Clean chroot"; phase_clean_chroot
    CURRENT=$((CURRENT + 1)); echo ""; echo "[${CURRENT}/${TOTAL}] Re-pack squashfs"; phase_repack_squashfs
    CURRENT=$((CURRENT + 1)); echo ""; echo "[${CURRENT}/${TOTAL}] Generate ISO"; phase_generate_iso
    CURRENT=$((CURRENT + 1)); echo ""; echo "[${CURRENT}/${TOTAL}] Create persistence"; phase_create_persistence
    CURRENT=$((CURRENT + 1)); echo ""; echo "[${CURRENT}/${TOTAL}] Cleanup"; phase_cleanup

    echo ""
    echo "╔════════════════════════════════════════════════════════╗"
    echo "║             BUILD COMPLETE                            ║"
    echo "╚════════════════════════════════════════════════════════╝"
    echo ""
    echo "  ISO:        $OUTPUT_DIR/hermes-rescue.iso"
    echo "  Persist:    $OUTPUT_DIR/casper-rw"
    echo "  Data img:   $OUTPUT_DIR/hermes-data.ext4"
    echo ""
    echo "Write to USB:"
    echo "  # Copy ISO to a Ventoy USB, or write directly:"
    echo "  sudo dd if=$OUTPUT_DIR/hermes-rescue.iso of=/dev/sdX bs=4M status=progress"
    echo ""
    echo "For persistence, copy casper-rw to the USB root alongside the ISO,"
    echo "or the live system will auto-detect a partition labeled 'HERMES_DATA'."
    echo ""
}

main "$@"