#!/bin/bash

set -e

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

# Function to print colored messages
print_error() { echo -e "${RED}[ERROR]${NC} $1" >&2; }
print_success() { echo -e "${GREEN}[SUCCESS]${NC} $1"; }
print_warning() { echo -e "${YELLOW}[WARNING]${NC} $1"; }
print_info() { echo -e "[INFO] $1"; }

# Function to find OVMF firmware files on Arch Linux
find_ovmf_firmware() {
    local firmware_type="$1" # CODE or VARS
    
    # Arch Linux OVMF paths (from edk2-ovmf package)
    local ovmf_paths=(
        "/usr/share/edk2/x64/OVMF_${firmware_type}.4m.fd"
        "/usr/share/edk2-ovmf/x64/OVMF_${firmware_type}.fd"
        "/usr/share/ovmf/x64/OVMF_${firmware_type}.4m.fd"
    )
    
    for path in "${ovmf_paths[@]}"; do
        if [ -f "$path" ]; then
            echo "$path"
            return 0
        fi
    done
    
    return 1
}

# Function to check and install Arch Linux dependencies
check_requirements() {
    local errors=0
    local packages_to_install=()
    
    # Check for QEMU
    if ! pacman -Qi qemu-full &>/dev/null && ! pacman -Qi qemu-system-x86 &>/dev/null; then
        print_warning "QEMU is not installed"
        packages_to_install+=("qemu-full")
    else
        print_success "QEMU is installed"
    fi
    
    # Check for OVMF firmware
    if ! pacman -Qi edk2-ovmf &>/dev/null; then
        print_warning "OVMF UEFI firmware is not installed"
        packages_to_install+=("edk2-ovmf")
    else
        print_success "OVMF firmware is installed"
    fi
    
    # Install missing packages
    if [ ${#packages_to_install[@]} -gt 0 ]; then
        print_info "Installing missing packages: ${packages_to_install[*]}"
        sudo pacman -S --noconfirm --needed "${packages_to_install[@]}"
        if [ $? -ne 0 ]; then
            print_error "Failed to install required packages"
            ((errors++))
        fi
    fi
    
    # Check for KVM support
    if [ ! -e /dev/kvm ]; then
        print_warning "KVM is not available. VM will run without hardware acceleration (slower)"
        print_info "To enable KVM:"
        print_info "  1. Ensure virtualization is enabled in BIOS"
        print_info "  2. Load KVM modules: sudo modprobe kvm kvm_intel (or kvm_amd)"
        KVM_AVAILABLE=false
    else
        if [ ! -r /dev/kvm ] || [ ! -w /dev/kvm ]; then
            print_warning "KVM permissions issue. Add user to kvm group: sudo usermod -a -G kvm $USER"
            KVM_AVAILABLE=false
        else
            KVM_AVAILABLE=true
            print_success "KVM hardware acceleration is available"
        fi
    fi
    
    # Find OVMF firmware files
    OVMF_CODE=$(find_ovmf_firmware "CODE")
    if [ -z "$OVMF_CODE" ]; then
        print_error "Could not find OVMF CODE firmware"
        print_info "Install with: sudo pacman -S edk2-ovmf"
        ((errors++))
    else
        print_success "Found OVMF CODE at: $OVMF_CODE"
    fi
    
    OVMF_VARS=$(find_ovmf_firmware "VARS")
    if [ -z "$OVMF_VARS" ]; then
        # On Arch, VARS might not exist separately, use CODE
        OVMF_VARS="$OVMF_CODE"
        print_info "Using OVMF CODE for VARS"
    else
        print_success "Found OVMF VARS at: $OVMF_VARS"
    fi
    
    return $errors
}

# Function to setup storage
setup_storage() {
    local storage_dir="${STORAGE_DIR:-/tmp}"
    
    if [ ! -d "$storage_dir" ]; then
        print_error "Storage directory does not exist: $storage_dir"
        exit 1
    fi
    
    if [ ! -f "$storage_dir/test.qcow2" ]; then
        print_info "Creating virtual disk at $storage_dir/test.qcow2..."
        qemu-img create -f qcow2 "$storage_dir/test.qcow2" 20G
    else
        print_info "Using existing disk at $storage_dir/test.qcow2"
    fi
    
    if [ ! -f "$storage_dir/test_large.qcow2" ]; then
        print_info "Creating second virtual disk at $storage_dir/test_large.qcow2..."
        qemu-img create -f qcow2 "$storage_dir/test_large.qcow2" 20G
    else
        print_info "Using existing disk at $storage_dir/test_large.qcow2"
    fi
    
    DISK1="$storage_dir/test.qcow2"
    DISK2="$storage_dir/test_large.qcow2"
}

# Function to build QEMU command
build_qemu_command() {
    local iso_file="$1"
    local qemu_args=()
    
    # Basic configuration
    qemu_args+=("qemu-system-x86_64")
    
    # CPU and acceleration
    if [ "$KVM_AVAILABLE" = true ]; then
        qemu_args+=("-cpu" "host")
        qemu_args+=("-enable-kvm")
        qemu_args+=("-machine" "q35,accel=kvm")
    else
        qemu_args+=("-cpu" "qemu64")
        qemu_args+=("-machine" "q35")
    fi
    
    # Memory
    local total_mem=$(free -m | awk '/^Mem:/{print $2}')
    local vm_mem="${MEMORY:-8192}"
    if [ "$total_mem" -lt 10240 ] && [ "$vm_mem" -gt 4096 ]; then
        vm_mem=4096
        print_warning "Limited system memory. Using ${vm_mem}MB for VM"
    fi
    qemu_args+=("-m" "$vm_mem")
    
    # UEFI firmware (unless BIOS mode requested)
    if [ -n "$OVMF_CODE" ] && [ "$BOOT_MODE" != "bios" ]; then
        # Create writable VARS file from the correct template
        OVMF_VARS_COPY="/tmp/OVMF_VARS_$$.fd"
        if [ -f "$OVMF_VARS" ]; then
            cp "$OVMF_VARS" "$OVMF_VARS_COPY"
        else
            # Fallback: create empty VARS file if template not found
            dd if=/dev/zero of="$OVMF_VARS_COPY" bs=1M count=4 2>/dev/null
        fi
        
        qemu_args+=("-drive" "if=pflash,format=raw,readonly=on,file=$OVMF_CODE")
        qemu_args+=("-drive" "if=pflash,format=raw,file=$OVMF_VARS_COPY")
    fi
    
    # Storage devices - simplified for testing
    qemu_args+=("-drive" "file=$DISK1,format=qcow2,if=none,id=drive0")
    qemu_args+=("-device" "virtio-blk-pci,drive=drive0,bootindex=2")
    
    qemu_args+=("-drive" "file=$DISK2,format=qcow2,if=none,id=drive1")
    qemu_args+=("-device" "virtio-blk-pci,drive=drive1,bootindex=3")
    
    # CD-ROM with ISO - set as first boot device for UEFI
    qemu_args+=("-drive" "file=$iso_file,media=cdrom,if=none,format=raw,id=cdrom0")
    qemu_args+=("-device" "ide-cd,drive=cdrom0,bootindex=1")
    
    # Display configuration
    if [ "$HEADLESS" = true ]; then
        qemu_args+=("-nographic")
        qemu_args+=("-vnc" ":1")
        print_info "Headless mode. VNC server on :1 (port 5901)"
    else
        # Use virtio-vga for both UEFI and BIOS modes
        qemu_args+=("-device" "virtio-vga")
        # Enable OpenGL acceleration in GTK display
        qemu_args+=("-display" "gtk,gl=on")
    fi
    
    # Input devices
    qemu_args+=("-usb" "-device" "usb-tablet")
    
    # Network
    qemu_args+=("-netdev" "user,id=net0")
    qemu_args+=("-device" "virtio-net-pci,netdev=net0")
    
    echo "${qemu_args[@]}"
}

# Function to show usage
show_usage() {
    echo "Omarchy ISO Boot Tool - Arch Linux"
    echo ""
    echo "Usage: $0 <image.iso> [options]"
    echo ""
    echo "Options:"
    echo "  --bios            Force BIOS boot mode (default: UEFI)"
    echo "  --storage DIR     Directory for VM storage (default: /tmp)"
    echo "  --memory MB       Memory allocation in MB (default: 8192)"
    echo "  --no-kvm          Disable KVM acceleration"
    echo "  --headless        Run without display (VNC on :1)"
    echo "  --help            Show this help message"
    echo ""
    echo "Examples:"
    echo "  $0 omarchy.iso"
    echo "  $0 omarchy.iso --memory 4096"
    echo "  $0 omarchy.iso --bios --storage ~/vms"
    echo ""
}

# Main function
main() {
    print_info "Omarchy ISO Boot Tool for Arch Linux"
    echo "======================================"
    
    # Parse arguments
    if [ $# -eq 0 ] || [ "$1" = "--help" ] || [ "$1" = "-h" ]; then
        show_usage
        exit 0
    fi
    
    ISO_FILE="$1"
    shift
    
    # Check if ISO exists
    if [ ! -f "$ISO_FILE" ]; then
        print_error "ISO file not found: $ISO_FILE"
        exit 1
    fi
    
    print_success "ISO file: $ISO_FILE ($(du -h "$ISO_FILE" | cut -f1))"
    
    # Parse additional options
    while [ $# -gt 0 ]; do
        case "$1" in
            --bios)
                BOOT_MODE="bios"
                print_info "Using BIOS boot mode"
                ;;
            --storage)
                STORAGE_DIR="$2"
                shift
                ;;
            --memory)
                MEMORY="$2"
                shift
                ;;
            --no-kvm)
                KVM_AVAILABLE=false
                print_info "KVM acceleration disabled by user"
                ;;
            --headless)
                HEADLESS=true
                print_info "Headless mode enabled"
                ;;
            *)
                print_warning "Unknown option: $1"
                ;;
        esac
        shift
    done
    
    # Check system requirements
    print_info "Checking system requirements..."
    if ! check_requirements; then
        print_error "System requirements not met"
        exit 1
    fi
    
    # Setup storage
    print_info "Setting up storage..."
    setup_storage
    
    # Build QEMU command
    QEMU_CMD=$(build_qemu_command "$ISO_FILE")
    
    # Show command for debugging
    print_info "QEMU command:"
    echo "  $QEMU_CMD" | fold -s -w 80 | sed '2,$s/^/  /'
    echo ""
    
    print_success "Starting VM... (Press Ctrl-A X to exit)"
    echo "======================================"
    
    # Execute QEMU
    eval $QEMU_CMD
    
    # Cleanup
    if [ -f "$OVMF_VARS_COPY" ]; then
        rm -f "$OVMF_VARS_COPY"
    fi
}

# Cleanup on exit
cleanup() {
    if [ -f "$OVMF_VARS_COPY" ]; then
        rm -f "$OVMF_VARS_COPY"
    fi
}

trap cleanup EXIT

# Run main function
main "$@"