#!/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
    # Persist firmaware as well if $STORAGE_DIR is specified
    if [ -z "$STORAGE_DIR" ]; 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
    else
      OVMF_VARS_COPY="$STORAGE_DIR/OVMF_VARS.fd"

      # Only create firmaware file if it doesn't exist to allow persistence
      if [ ! -f "$OVMF_VARS_COPY" ]; then
        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
      fi
    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 - Smart boot order that handles both installation and post-install
  # Credits to Ryan Hughes for this elegant solution:
  # - Empty disk (bootindex=1) fails to boot -> falls back to ISO (bootindex=2)
  # - After installation, disk is bootable -> boots from disk automatically
  # This achieves David's goal without any complex monitoring!
  qemu_args+=("-drive" "file=$DISK1,format=qcow2,if=none,id=drive0")
  qemu_args+=("-device" "virtio-blk-pci,drive=drive0,bootindex=1")

  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 second boot device (fallback)
  qemu_args+=("-drive" "file=$iso_file,media=cdrom,if=none,format=raw,id=cdrom0")
  qemu_args+=("-device" "ide-cd,drive=cdrom0,bootindex=2")

  qemu_args+=("-device" "virtio-vga-gl")
  qemu_args+=("-display" "gtk,gl=on")

  # 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 "  --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"
      ;;
    *)
      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

  print_info "Boot order: Disk (if bootable) -> ISO (fallback)"
  print_info "After installation, system will automatically boot from disk"

  # 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 "$@"