#!/bin/bash

# omarchy:summary=Install, launch, stop, inspect, or remove the Windows VM
# omarchy:args=<install|remove|launch|stop|status> [options]
# omarchy:requires-sudo=true

COMPOSE_FILE="$HOME/.config/windows/docker-compose.yml"

check_prerequisites() {
  local DISK_SIZE_GB=${1:-64}
  local REQUIRED_SPACE=$((DISK_SIZE_GB + 10))  # Add 10GB for Windows ISO and overhead

  # Check for KVM support
  if [[ ! -e /dev/kvm ]]; then
    gum style \
      --border normal \
      --padding "1 2" \
      --margin "1" \
      "❌ KVM virtualization not available!" \
      "" \
      "Please enable virtualization in BIOS or run:" \
      "  sudo modprobe kvm-intel  # for Intel CPUs" \
      "  sudo modprobe kvm-amd    # for AMD CPUs"
    exit 1
  fi

  # Check disk space
  AVAILABLE_SPACE=$(df "$HOME" | awk 'NR==2 {print int($4/1024/1024)}')
  if (( AVAILABLE_SPACE < REQUIRED_SPACE )); then
    echo "❌ Insufficient disk space!"
    echo "   Available: ${AVAILABLE_SPACE}GB"
    echo "   Required: ${REQUIRED_SPACE}GB (${DISK_SIZE_GB}GB disk + 10GB for Windows image)"
    exit 1
  fi
}

install_windows() {
  # Set up trap to handle Ctrl+C
  trap "echo ''; echo 'Installation cancelled by user'; exit 1" INT

  check_prerequisites

  omarchy-pkg-add freerdp openbsd-netcat gum

  mkdir -p "$HOME/.windows"
  mkdir -p "$HOME/.config/windows"
  mkdir -p "$HOME/.local/share/applications"

  cat << EOF | tee "$HOME/.local/share/applications/windows-vm.desktop" > /dev/null
[Desktop Entry]
Name=Windows
Comment=Start Windows VM via Docker and connect with RDP
Exec=uwsm app -- omarchy-windows-vm launch
Icon=windows
Terminal=false
Type=Application
Categories=System;Virtualization;
EOF

  # Get system resources
  TOTAL_RAM=$(free -h | awk 'NR==2 {print $2}')
  TOTAL_RAM_GB=$(awk 'NR==1 {printf "%d", $2/1024/1024}' /proc/meminfo)
  TOTAL_CORES=$(nproc)

  echo ""
  echo "System Resources Detected:"
  echo "  Total RAM: $TOTAL_RAM"
  echo "  Total CPU Cores: $TOTAL_CORES"
  echo ""

  RAM_OPTIONS=""
  for size in 2 4 8 16 32 64; do
    if (( size <= TOTAL_RAM_GB )); then
      RAM_OPTIONS="$RAM_OPTIONS ${size}G"
    fi
  done

  SELECTED_RAM=$(echo $RAM_OPTIONS | tr ' ' '\n' | gum choose --selected="4G" --header="How much RAM would you like to allocate to Windows VM?")

  # Check if user cancelled
  if [[ -z $SELECTED_RAM ]]; then
    echo "Installation cancelled by user"
    exit 1
  fi

  SELECTED_CORES=$(gum input --placeholder="Number of CPU cores (1-$TOTAL_CORES)" --value="2" --header="How many CPU cores would you like to allocate to Windows VM?" --char-limit=2)

  # Check if user cancelled (Ctrl+C in gum input returns empty string)
  if [[ -z $SELECTED_CORES ]]; then
    echo "Installation cancelled by user"
    exit 1
  fi

  if ! [[ $SELECTED_CORES =~ ^[0-9]+$ ]] || (( SELECTED_CORES < 1 )) || (( SELECTED_CORES > TOTAL_CORES )); then
    echo "Invalid input. Using default: 2 cores"
    SELECTED_CORES=2
  fi

  AVAILABLE_SPACE=$(df "$HOME" | awk 'NR==2 {print int($4/1024/1024)}')
  MAX_DISK_GB=$((AVAILABLE_SPACE - 10))  # Leave 10GB for Windows image

  # Check if we have enough space for minimum
  if (( MAX_DISK_GB < 32 )); then
    echo "❌ Insufficient disk space for Windows VM!"
    echo "   Available: ${AVAILABLE_SPACE}GB"
    echo "   Minimum required: 42GB (32GB disk + 10GB for Windows image)"
    exit 1
  fi

  DISK_OPTIONS=""
  for size in 32 64 128 256 512; do
    if (( size <= MAX_DISK_GB )); then
      DISK_OPTIONS="$DISK_OPTIONS ${size}G"
    fi
  done

  # Default to 64G if available, otherwise 32G
  DEFAULT_DISK="64G"
  if ! echo "$DISK_OPTIONS" | grep -q "64G"; then
    DEFAULT_DISK="32G"
  fi

  SELECTED_DISK=$(echo $DISK_OPTIONS | tr ' ' '\n' | gum choose --selected="$DEFAULT_DISK" --header="How much disk space would you like to give Windows VM? (64GB+ recommended)")

  # Check if user cancelled
  if [[ -z $SELECTED_DISK ]]; then
    echo "Installation cancelled by user"
    exit 1
  fi

  # Extract just the number for prerequisite check
  DISK_SIZE_NUM=$(echo "$SELECTED_DISK" | sed 's/G//')

  # Re-check prerequisites with selected disk size
  check_prerequisites "$DISK_SIZE_NUM"

  # Prompt for username and password
  USERNAME=$(gum input --placeholder="Username (Press enter to use default: docker)" --header="Enter Windows username:")
  if [[ -z $USERNAME ]]; then
    USERNAME="docker"
  fi

  PASSWORD=$(gum input --placeholder="Password (Press enter to use default: admin)" --password --header="Enter Windows password:")
  if [[ -z $PASSWORD ]]; then
    PASSWORD="admin"
    PASSWORD_DISPLAY="(default)"
  else
    PASSWORD_DISPLAY="(user-defined)"
  fi

  # Display configuration summary
  gum style \
    --border normal \
    --padding "1 2" \
    --margin "1" \
    --align left \
    --bold \
    "Windows VM Configuration" \
    "" \
    "RAM:       $SELECTED_RAM" \
    "CPU:       $SELECTED_CORES cores" \
    "Disk:      $SELECTED_DISK" \
    "Username:  $USERNAME" \
    "Password:  $PASSWORD_DISPLAY"

  # Ask for confirmation
  echo ""
  if ! gum confirm "Proceed with this configuration?"; then
    echo "Installation cancelled by user"
    exit 1
  fi

  mkdir -p $HOME/Windows

  # Create docker-compose.yml in user config directory
  cat << EOF | tee "$COMPOSE_FILE" > /dev/null
services:
  windows:
    image: dockurr/windows
    container_name: omarchy-windows
    environment:
      VERSION: "11"
      RAM_SIZE: "$SELECTED_RAM"
      CPU_CORES: "$SELECTED_CORES"
      DISK_SIZE: "$SELECTED_DISK"
      USERNAME: "$USERNAME"
      PASSWORD: "$PASSWORD"
      TZ: "$(timedatectl show -p Timezone --value 2>/dev/null || echo UTC)"
      ARGUMENTS: "-rtc base=localtime,clock=host,driftfix=slew"
    devices:
      - /dev/kvm
      - /dev/net/tun
    cap_add:
      - NET_ADMIN
    ports:
      - 127.0.0.1:8006:8006
      - 127.0.0.1:3389:3389/tcp
      - 127.0.0.1:3389:3389/udp
    volumes:
      - $HOME/.windows:/storage
      - $HOME/Windows:/shared
    restart: "no"
    stop_grace_period: 2m
EOF

  echo ""
  echo "Starting Windows VM installation..."
  echo "This will download a Windows 11 image (may take 10-15 minutes)."
  echo ""
  echo "Monitor installation progress at: http://127.0.0.1:8006"
  echo ""

  # Start docker-compose with user's config
  echo "Starting Windows VM with docker-compose..."
  if ! docker-compose -f "$COMPOSE_FILE" up -d 2>&1; then
    echo "❌ Failed to start Windows VM!"
    echo "   Common issues:"
    echo "   - Docker daemon not running: sudo systemctl start docker"
    echo "   - Port already in use: check if another VM is running"
    echo "   - Permission issues: make sure you're in the docker group"
    exit 1
  fi

  echo ""
  echo "Windows VM is starting up!"
  echo ""
  echo "Opening browser to monitor installation..."

  # Open browser to monitor installation
  sleep 3
  xdg-open "http://127.0.0.1:8006"

  echo ""
  echo "Installation is running in the background."
  echo "You can monitor progress at: http://127.0.0.1:8006"
  echo ""
  echo "Once finished, launch 'Windows' via Super + Space"
  echo ""
  echo "To stop the VM: omarchy-windows-vm stop"
  echo "To change resources: ~/.config/windows/docker-compose.yml"
  echo ""
}

remove_windows() {
  if ! gum confirm --default=false "Remove Windows VM and delete all associated data?"; then
    echo "Removal cancelled by user"
    exit 1
  fi

  echo "Removing Windows VM..."

  docker-compose -f "$COMPOSE_FILE" down 2>/dev/null || true

  docker rmi dockurr/windows 2>/dev/null || echo "Image already removed or not found"

  rm "$HOME/.local/share/applications/windows-vm.desktop"
  rm -rf "$HOME/.config/windows"
  rm -rf "$HOME/.windows"

  echo ""
  echo "Windows VM removal completed!"
}

launch_windows() {
  KEEP_ALIVE=false
  if [[ $1 = "--keep-alive" ]] || [[ $1 = "-k" ]]; then
    KEEP_ALIVE=true
  fi

  # Check if config exists
  if [[ ! -f $COMPOSE_FILE ]]; then
    echo "Windows VM not configured. Please run: omarchy-windows-vm install"
    exit 1
  fi

  # Extract credentials from compose file
  WIN_USER=$(grep "USERNAME:" "$COMPOSE_FILE" | sed 's/.*USERNAME: "\(.*\)"/\1/')
  WIN_PASS=$(grep "PASSWORD:" "$COMPOSE_FILE" | sed 's/.*PASSWORD: "\(.*\)"/\1/')

  # Use defaults if not found
  [[ -z $WIN_USER ]] && WIN_USER="docker"
  [[ -z $WIN_PASS ]] && WIN_PASS="admin"

  # Check if container is already running
  CONTAINER_STATUS=$(docker inspect --format='{{.State.Status}}' omarchy-windows 2>/dev/null)

  if [[ $CONTAINER_STATUS != "running" ]]; then
    echo "Starting Windows VM..."

    if ! docker-compose -f "$COMPOSE_FILE" up -d 2>&1; then
      echo "❌ Failed to start Windows VM!"
      echo "   Try checking: omarchy-windows-vm status"
      echo "   View logs: docker logs omarchy-windows"
      omarchy-notification-send -u critical "Windows VM" "Failed to start Windows VM"
      exit 1
    fi
  fi

  # docker logs keeps output across stop/start and is only cleared when the
  # container is removed, so an unanchored grep matches "started successfully"
  # from an earlier boot and returns immediately. Anchor the scan to the current
  # container start, and run it even when the container was already up: the
  # image restarts the guest in place whenever Windows reboots.
  windows_started() {
    local started_at
    started_at=$(docker inspect --format='{{.State.StartedAt}}' omarchy-windows 2>/dev/null)
    # An empty --since would drop the filter and match a stale boot again
    [[ -n $started_at ]] || return 1
    docker logs --since "$started_at" omarchy-windows 2>&1 | grep -qi "windows started successfully"
  }

  echo "Waiting for Windows VM to start..."
  WAIT_COUNT=0
  until windows_started; do
    sleep 2
    WAIT_COUNT=$((WAIT_COUNT + 1))
    if (( WAIT_COUNT > 60 )); then  # 2 minutes timeout
      echo ""
      echo "❌ Timeout: Windows VM failed to start within 2 minutes"
      echo "   Check logs: docker logs omarchy-windows"
      exit 1
    fi
  done

  # Build the connection info
  if [[ $KEEP_ALIVE = "true" ]]; then
    LIFECYCLE="VM will keep running after RDP closes
To stop: omarchy-windows-vm stop"
  else
    LIFECYCLE="VM will auto-stop when RDP closes"
  fi

  gum style \
    --border normal \
    --padding "1 2" \
    --margin "1" \
    --align center \
    "Connecting to Windows VM" \
    "" \
    "$LIFECYCLE"

  # FreeRDP 3 tries Kerberos before NTLM for NLA, and krb5 ships /etc/krb5.conf
  # as the MIT sample with default_realm = ATHENA.MIT.EDU. Every connect then
  # goes looking for MIT's KDC: with internet up it fails fast, but off the
  # network each attempt blocks ~23s and no RDP window is ever drawn. The VM
  # authenticates against a local Windows account, so point FreeRDP at a
  # realm-less config and let it fall straight through to NTLM.
  KRB5_CONF="$(dirname "$COMPOSE_FILE")/krb5.conf"
  if [[ ! -f $KRB5_CONF ]]; then
    printf '[libdefaults]\n  dns_lookup_kdc = false\n  dns_lookup_realm = false\n' >"$KRB5_CONF"
  fi
  export KRB5_CONFIG="$KRB5_CONF"

  # Detect display scale from Hyprland
  HYPR_SCALE=$(hyprctl monitors -j | jq -r '.[] | select (.focused == true) | .scale')
  SCALE_PERCENT=$(echo "$HYPR_SCALE" | awk '{print int($1 * 100)}')

  RDP_SCALE=""
  if (( SCALE_PERCENT >= 170 )); then
    RDP_SCALE="/scale:180"
  elif (( SCALE_PERCENT >= 130 )); then
    RDP_SCALE="/scale:140"
  fi
  # If scale is less than 130%, don't set any scale (use default 100)

  # Connect with RDP in fullscreen (auto-detects resolution)
  xfreerdp3 /u:"$WIN_USER" /p:"$WIN_PASS" /v:127.0.0.1:3389 -grab-keyboard /sound /microphone /clipboard /cert:ignore /title:"Windows VM - Omarchy" /dynamic-resolution /gfx:AVC444 /floatbar:sticky:off,default:visible,show:fullscreen $RDP_SCALE

  # After RDP closes, stop the container unless --keep-alive was specified
  if [[ $KEEP_ALIVE = "false" ]]; then
    echo ""
    echo "RDP session closed. Stopping Windows VM..."
    docker-compose -f "$COMPOSE_FILE" down
    echo "Windows VM stopped."
  else
    echo ""
    echo "RDP session closed. Windows VM is still running."
    echo "To stop it: omarchy-windows-vm stop"
  fi
}

stop_windows() {
  if [[ ! -f $COMPOSE_FILE ]]; then
    echo "Windows VM not configured."
    exit 1
  fi

  echo "Stopping Windows VM..."
  docker-compose -f "$COMPOSE_FILE" down
  echo "Windows VM stopped."
}

status_windows() {
  if [[ ! -f $COMPOSE_FILE ]]; then
    echo "Windows VM not configured."
    echo "To set up: omarchy-windows-vm install"
    exit 1
  fi

  CONTAINER_STATUS=$(docker inspect --format='{{.State.Status}}' omarchy-windows 2>/dev/null)

  if [[ -z $CONTAINER_STATUS ]]; then
    echo "Windows VM container not found."
    echo "To start: omarchy-windows-vm launch"
  elif [[ $CONTAINER_STATUS = "running" ]]; then
    gum style \
      --border normal \
      --padding "1 2" \
      --margin "1" \
      --align left \
      "Windows VM Status: RUNNING" \
      "" \
      "Web interface: http://127.0.0.1:8006" \
      "RDP available: port 3389" \
      "" \
      "To connect: omarchy-windows-vm launch" \
      "To stop:    omarchy-windows-vm stop"
  else
    echo "Windows VM is stopped (status: $CONTAINER_STATUS)"
    echo "To start: omarchy-windows-vm launch"
  fi
}

show_usage() {
  echo "Usage: omarchy-windows-vm [command] [options]"
  echo ""
  echo "Commands:"
  echo "  install              Install and configure Windows VM"
  echo "  remove               Remove Windows VM and optionally its data"
  echo "  launch [options]     Start Windows VM (if needed) and connect via RDP"
  echo "                       Options:"
  echo "                         --keep-alive, -k   Keep VM running after RDP closes"
  echo "  stop                 Stop the running Windows VM"
  echo "  status               Show current VM status"
  echo "  help                 Show this help message"
  echo ""
  echo "Examples:"
  echo "  omarchy-windows-vm install           # Set up Windows VM for first time"
  echo "  omarchy-windows-vm launch            # Connect to VM (auto-stop on exit)"
  echo "  omarchy-windows-vm launch -k         # Connect to VM (keep running)"
  echo "  omarchy-windows-vm stop              # Shut down the VM"
}

# Main command dispatcher
case "$1" in
  install)
    install_windows
    ;;
  remove)
    remove_windows
    ;;
  launch|start)
    launch_windows "$2"
    ;;
  stop|down)
    stop_windows
    ;;
  status)
    status_windows
    ;;
  help|--help|-h|"")
    show_usage
    ;;
  *)
    echo "Unknown command: $1" >&2
    echo "" >&2
    show_usage >&2
    exit 1
    ;;
esac