#! /bin/bash

#######################################
# WSL (Windows Subsystem for Linux) bash completion
#
# Globals:
#   COMPREPLY
#   COMP_CWORD
#   COMP_WORDS
# Arguments:
#  None
#######################################

# Helper function to get installed WSL distributions
_wsl_get_distros() {
  local distros
  # Parse wsl --list --verbose output (locale-agnostic format)
  # Skip header line, remove leading spaces and asterisk (default marker), extract first column
  # Note: wsl.exe outputs UTF-16 LE, so we need to filter null bytes before processing
  # Convert to lowercase for easier completion without case sensitivity
  distros=$(wsl.exe --list --verbose 2>/dev/null | tr -d '\0' | sed -e 's/\r$//' | tail -n +2 | sed -e 's/^[* ]*//' | awk 'NF > 0 {print tolower($1)}')
  echo "$distros"
}

# Helper function to get available WSL distributions for installation
_wsl_get_online_distros() {
  local distros
  # Parse wsl --list --online output, skip header lines and extract NAME column
  # Skip first 4 lines (2 description lines + blank line + column headers), then get first column
  # Note: wsl.exe outputs UTF-16 LE, so we need to filter null bytes before processing
  # Convert to lowercase for easier completion without case sensitivity
  distros=$(wsl.exe --list --online 2>/dev/null | tr -d '\0' | sed -e 's/\r$//' | tail -n +5 | awk 'NF > 0 {print tolower($1)}')
  echo "$distros"
}

function _wsl() {

  local cur=${COMP_WORDS[COMP_CWORD]}
  local prev=${COMP_WORDS[COMP_CWORD - 1]}
  local cmd=""
  
  # Find the main command if present
  local i
  for ((i = 1; i < COMP_CWORD; i++)); do
    case "${COMP_WORDS[i]}" in
      --install|--manage|--mount|--export|--import|--import-in-place|--list|-l|--set-default|-s|--set-version|--terminate|-t|--unregister|--shutdown|--update)
        cmd="${COMP_WORDS[i]}"
        break
        ;;
    esac
  done

  # Handle options based on previous word
  case "$prev" in
    --cd)
      # Suggest directories
      mapfile -t COMPREPLY < <(compgen -d -- "${cur}")
      return 0
      ;;
    --distribution|-d|--set-default|-s|--terminate|-t|--unregister|--set-version|--manage|--export)
      # Suggest distribution names from wsl --list output
      local distros
      distros=$(_wsl_get_distros)
      mapfile -t COMPREPLY < <(compgen -W "$distros" -- "${cur}")
      return 0
      ;;
    --mount|--unmount)
      # Disk path - provide file completion
      mapfile -t COMPREPLY < <(compgen -f -- "${cur}")
      return 0
      ;;
    --distribution-id)
      # Distribution GUID - no completion
      COMPREPLY=()
      return 0
      ;;
    --user|-u)
      # Username - no completion
      COMPREPLY=()
      return 0
      ;;
    --shell-type)
      mapfile -t COMPREPLY < <(compgen -W "standard login none" -- "${cur}")
      return 0
      ;;
    --exec|-e)
      # Command to execute - provide command completion
      mapfile -t COMPREPLY < <(compgen -c -- "${cur}")
      return 0
      ;;
    --set-default-version|--version)
      mapfile -t COMPREPLY < <(compgen -W "1 2" -- "${cur}")
      return 0
      ;;
    --from-file|--import-in-place)
      # File path completion
      mapfile -t COMPREPLY < <(compgen -f -- "${cur}")
      return 0
      ;;
    --location|--move)
      # Directory path completion
      mapfile -t COMPREPLY < <(compgen -d -- "${cur}")
      return 0
      ;;
    --name)
      # Distribution name - no completion
      COMPREPLY=()
      return 0
      ;;
    --type)
      # Filesystem type
      mapfile -t COMPREPLY < <(compgen -W "ext4 ext3 ext2 btrfs xfs" -- "${cur}")
      return 0
      ;;
    --partition)
      # Partition index - no completion
      COMPREPLY=()
      return 0
      ;;
    --options)
      # Mount options - no completion
      COMPREPLY=()
      return 0
      ;;
    --vhd-size|--resize)
      # Memory string - no completion
      COMPREPLY=()
      return 0
      ;;
    --set-sparse)
      mapfile -t COMPREPLY < <(compgen -W "true false" -- "${cur}")
      return 0
      ;;
    --set-default-user)
      # Username - no completion
      COMPREPLY=()
      return 0
      ;;
    --format)
      mapfile -t COMPREPLY < <(compgen -W "tar tar.gz tar.xz vhd" -- "${cur}")
      return 0
      ;;
  esac

  # Command-specific completions
  case "$cmd" in
    --install)
      if [[ "$cur" == -* ]]; then
        # Complete with install options
        mapfile -t COMPREPLY < <(compgen -W "--enable-wsl1 --fixed-vhd --from-file --legacy --location --name --no-distribution --no-launch -n --version --vhd-size --web-download" -- "${cur}")
      else
        # Check if we already have a distribution name
        local has_distro=0
        for ((i = 1; i < COMP_CWORD; i++)); do
          if [[ "${COMP_WORDS[i]}" == "--install" ]]; then
            continue
          elif [[ "${COMP_WORDS[i]}" != -* ]]; then
            has_distro=1
            break
          fi
        done
        
        if [[ $has_distro -eq 0 ]]; then
          # No distro yet, complete with online distribution names
          local distros
          distros=$(_wsl_get_online_distros)
          mapfile -t COMPREPLY < <(compgen -W "$distros" -- "${cur}")
        else
          # Already have a distro, complete with options (no hyphen prefix)
          COMPREPLY=()
        fi
      fi
      return 0
      ;;
    --manage)
      mapfile -t COMPREPLY < <(compgen -W "--move --set-sparse --set-default-user --resize" -- "${cur}")
      return 0
      ;;
    --mount)
      if [[ "$COMP_CWORD" -eq 2 ]]; then
        # First argument after --mount is the disk path
        mapfile -t COMPREPLY < <(compgen -f -- "${cur}")
      else
        mapfile -t COMPREPLY < <(compgen -W "--vhd --bare --name --type --options --partition" -- "${cur}")
      fi
      return 0
      ;;
    --shutdown)
      mapfile -t COMPREPLY < <(compgen -W "--force" -- "${cur}")
      return 0
      ;;
    --update)
      mapfile -t COMPREPLY < <(compgen -W "--pre-release" -- "${cur}")
      return 0
      ;;
    --export)
      mapfile -t COMPREPLY < <(compgen -W "--format" -- "${cur}")
      return 0
      ;;
    --import)
      mapfile -t COMPREPLY < <(compgen -W "--version --vhd" -- "${cur}")
      return 0
      ;;
    --list|-l)
      mapfile -t COMPREPLY < <(compgen -W "--all --running --quiet -q --verbose -v --online -o" -- "${cur}")
      return 0
      ;;
  esac

  # If no command selected yet, provide main options
  if [[ "$cur" == -* ]]; then
    local main_opts="--help --debug-shell --install --manage --mount --set-default-version --shutdown --status --unmount --uninstall --update --version -v --export --import --import-in-place --list -l --set-default -s --set-version --terminate -t --unregister"
    local running_opts="--exec -e --shell-type -- --cd --distribution -d --distribution-id --user -u --system"
    mapfile -t COMPREPLY < <(compgen -W "$main_opts $running_opts" -- "${cur}")
  else
    # No hyphen prefix - could be distribution name or positional argument
    COMPREPLY=()
  fi

  return 0
}

complete -F _wsl wsl
complete -F _wsl wsl.exe
