#!/bin/bash

set -euo pipefail

# Particle-OS Boot Test Script
# This script creates a bootable disk image and tests it with QEMU

SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
PROJECT_ROOT="$(dirname "$SCRIPT_DIR")"
BUILD_DIR="$PROJECT_ROOT/build"
TEST_DIR="$PROJECT_ROOT/test"

# Colors for output
RED='\033[0;31m'
GREEN='\033[0;32m'
YELLOW='\033[1;33m'
BLUE='\033[0;34m'
NC='\033[0m' # No Color

log_info() {
    echo -e "${BLUE}[INFO]${NC} $1"
}

log_success() {
    echo -e "${GREEN}[SUCCESS]${NC} $1"
}

log_warning() {
    echo -e "${YELLOW}[WARNING]${NC} $1"
}

log_error() {
    echo -e "${RED}[ERROR]${NC} $1"
}

# Check prerequisites
check_prerequisites() {
    log_info "Checking prerequisites..."
    
    local missing_tools=()
    
    for tool in podman qemu-system-x86_64 truncate; do
        if ! command -v "$tool" >/dev/null 2>&1; then
            missing_tools+=("$tool")
        fi
    done
    
    # Check for system tools in /usr/sbin
    for tool in parted grub-install partprobe; do
        if [ ! -x "/usr/sbin/$tool" ]; then
            missing_tools+=("$tool")
        fi
    done
    
    if [ ${#missing_tools[@]} -gt 0 ]; then
        log_error "Missing required tools: ${missing_tools[*]}"
        exit 1
    fi
    
    # Check if Phase 2 image exists
    if ! podman image exists particle-os:phase2; then
        log_error "Phase 2 image not found. Run 'just build-phase2' first."
        exit 1
    fi
    
    log_success "All prerequisites satisfied"
}

# Create test directories
setup_test_environment() {
    log_info "Setting up test environment..."
    
    mkdir -p "$BUILD_DIR" "$TEST_DIR"
    
    log_success "Test environment ready"
}

# Create bootable disk image
create_bootable_image() {
    local image_size="2G"
    local image_path="$TEST_DIR/particle-os-test.img"
    
    log_info "Creating bootable disk image ($image_size)..."
    
    # Create raw disk image
    truncate -s "$image_size" "$image_path"
    
    # Partition the disk
    echo "yes" | /usr/sbin/parted "$image_path" mklabel gpt
    /usr/sbin/parted "$image_path" mkpart primary fat32 1MiB 512MiB
    /usr/sbin/parted "$image_path" mkpart primary ext4 512MiB 100%
    /usr/sbin/parted "$image_path" set 1 boot on
    /usr/sbin/parted "$image_path" set 1 esp on
    
    log_success "Disk image created: $image_path"
    echo "$image_path"
}

# Mount and prepare the image
prepare_image() {
    local image_path="$1"
    
    log_info "Preparing bootable image..."
    
    # Create loop device
    log_info "Creating loop device for: $image_path"
    local loop_dev=$(sudo losetup --find --show "$image_path")
    if [ -z "$loop_dev" ]; then
        log_error "Failed to create loop device"
        exit 1
    fi
    log_info "Using loop device: $loop_dev"
    
    # Debug: show what devices exist
    log_info "Available loop devices:"
    ls -la /dev/loop* 2>/dev/null || echo "No loop devices found"
    
    # Wait a moment for device nodes to appear
    sleep 1
    
    # Verify partition devices exist
    if [ ! -b "${loop_dev}p1" ] || [ ! -b "${loop_dev}p2" ]; then
        log_error "Partition devices not found. Creating them manually..."
        sudo partprobe "$loop_dev"
        sleep 2
        
        # Check again
        if [ ! -b "${loop_dev}p1" ] || [ ! -b "${loop_dev}p2" ]; then
            log_error "Partition devices still not found after partprobe"
            sudo losetup -d "$loop_dev" 2>/dev/null || true
            exit 1
        fi
    fi
    
    # Create filesystems
    sudo mkfs.fat -F32 "${loop_dev}p1"
    sudo mkfs.ext4 "${loop_dev}p2"
    
    # Mount partitions
    local mount_root="$TEST_DIR/mount"
    local mount_boot="$mount_root/boot"
    local mount_rootfs="$mount_root/rootfs"
    
    mkdir -p "$mount_boot" "$mount_rootfs"
    
    sudo mount "${loop_dev}p1" "$mount_boot"
    sudo mount "${loop_dev}p2" "$mount_rootfs"
    
    # Extract Phase 2 container to rootfs
    log_info "Extracting Phase 2 container..."
    podman create --name temp-phase2 particle-os:phase2
    podman export temp-phase2 | sudo tar -x -C "$mount_rootfs"
    podman rm temp-phase2
    
    # Set up bootloader
    log_info "Setting up bootloader..."
    sudo /usr/sbin/grub-install --target=x86_64-efi --efi-directory="$mount_boot" --boot-directory="$mount_boot" --removable "$loop_dev"
    
    # Create GRUB config
    sudo tee "$mount_boot/grub/grub.cfg" > /dev/null <<EOF
set timeout=5
set default=0

menuentry "Particle-OS Phase 2" {
    search --set=root --file /vmlinuz-6.12.38+deb13-amd64
    linux /vmlinuz-6.12.38+deb13-amd64 root=/dev/sda2 rw console=ttyS0
    initrd /initrd.img-6.12.38+deb13-amd64
}
EOF
    
    # Cleanup mounts
    sudo umount "$mount_boot" "$mount_rootfs"
    sudo losetup -d "$loop_dev"
    
    log_success "Image preparation complete"
}

# Test boot with QEMU
test_boot() {
    local image_path="$1"
    
    log_info "Testing boot with QEMU..."
    
    qemu-system-x86_64 \
        -m 2048 \
        -smp 2 \
        -drive file="$image_path",format=raw,if=virtio \
        -enable-kvm \
        -serial stdio \
        -nographic \
        -boot d \
        -display none \
        -monitor none
}

# Main execution
main() {
    log_info "Starting Particle-OS boot test..."
    
    check_prerequisites
    setup_test_environment
    
    local image_path=$(create_bootable_image)
    prepare_image "$image_path"
    
    log_success "Boot test setup complete!"
    log_info "Image created at: $image_path"
    log_info "You can now test booting with:"
    log_info "  qemu-system-x86_64 -m 2048 -smp 2 -drive file=$image_path,format=raw,if=virtio -enable-kvm -serial stdio -nographic"
    
    read -p "Test boot now? (y/N): " -n 1 -r
    echo
    if [[ $REPLY =~ ^[Yy]$ ]]; then
        test_boot "$image_path"
    fi
}

# Run main function
main "$@"