#!/bin/bash
set -e

# Signature: ./setup_base.sh <qcow2_path> <prep_script_path> <output_golden_image> <ssh_pub_key>
QCOW2_PATH="$1"
PREP_SCRIPT_PATH="$2"
GOLDEN_IMAGE="$3"
SSH_PUB_KEY="$4"

if [ -z "$QCOW2_PATH" ] || [ -z "$GOLDEN_IMAGE" ] || [ -z "$SSH_PUB_KEY" ]; then
    echo "Usage: $0 <qcow2_path> <prep_script_path> <output_golden_image> <ssh_pub_key>"
    exit 1
fi

echo "=========================================="
echo "Setting up golden image"
echo "=========================================="
echo "Source QCOW2: $QCOW2_PATH"
echo "Output Golden: $GOLDEN_IMAGE"
echo "Prep Script: $PREP_SCRIPT_PATH"
echo "SSH Key: $SSH_PUB_KEY"

# Ensure libvirt is running
sudo systemctl is-active --quiet libvirtd || sudo systemctl start libvirtd
sleep 2

# Copy original to golden
echo "[Step 1/3] Copying base image to golden image..."
cp "$QCOW2_PATH" "$GOLDEN_IMAGE"

# Apply custom preparation script if provided
if [ -f "$PREP_SCRIPT_PATH" ]; then
    echo "[Step 2/3] Applying custom preparation script..."
    export LIBGUESTFS_BACKEND=direct

    # Run the prep script inside the image
    sudo virt-customize -a "$GOLDEN_IMAGE" \
        --run "$PREP_SCRIPT_PATH" \
        --ssh-inject root:file:"$SSH_PUB_KEY" \
        --ssh-inject rocky:file:"$SSH_PUB_KEY" \
        --root-password password:rockytesting \
        --selinux-relabel 2>&1 || {
        echo "ERROR: virt-customize failed"
        exit 1
    }
else
    echo "[Step 2/3] No custom prep script provided, applying defaults..."
    export LIBGUESTFS_BACKEND=direct

    sudo virt-customize -a "$GOLDEN_IMAGE" \
        --ssh-inject root:file:"$SSH_PUB_KEY" \
        --root-password password:rockytesting \
        --install perl,git,wget,tar,openssh-server \
        --run-command 'systemctl enable sshd' \
        --selinux-relabel 2>&1 || {
        echo "ERROR: virt-customize failed"
        exit 1
    }
fi

echo "[Step 3/3] Verifying golden image..."
qemu-img info "$GOLDEN_IMAGE" | head -5

echo "=========================================="
echo "Golden image ready: $GOLDEN_IMAGE"
echo "=========================================="