#compdef copy_files

_copy_files() {
  local -a dirs
  local -a cp_hosts_default
  local -a used_hosts avail_hosts
  local _dotfiles="${DOTFILES:-$HOME/.dotfiles}"
  local i host_mode=0 host_arg_index=0

  cp_hosts_default=( $(_copy_files_ssh_hosts) )

  [[ -d "$_dotfiles" ]] || return 1
  dirs=(${(f)"$(print -rl -- $_dotfiles/*(/N:t))"})

  for (( i = 2; i < CURRENT; i++ )); do
    case "${words[i]}" in
      -h|--hosts)
        host_mode=1
        host_arg_index=$i
        ;;
      --)
        (( host_mode )) && host_mode=0
        ;;
    esac
  done

  if (( host_mode )); then
    used_hosts=()

    for (( i = host_arg_index + 1; i < CURRENT; i++ )); do
      [[ "${words[i]}" == "--" ]] && break
      used_hosts+=("${words[i]}")
    done

    avail_hosts=( "${(@)cp_hosts_default:|used_hosts}" )
    compadd -S '' -- "${avail_hosts[@]}" --
    return
  fi

  _describe -t dotfiles 'dotfiles subdirectory' dirs
  compadd -S '' -- -h --hosts
}

_copy_files_ssh_hosts() {
  local cfg="${HOME}/.ssh/config"
  [[ -r "$cfg" ]] || return 0

  awk '
    BEGIN {
      in_block = 0
      has_hostname = 0
      aliases = ""
    }

    /^[[:space:]]*#/ { next }
    /^[[:space:]]*$/ { next }

    tolower($1) == "host" {
      if (in_block && has_hostname && aliases != "") {
        print aliases
      }

      in_block = 1
      has_hostname = 0
      aliases = ""

      for (i = 2; i <= NF; i++) {
        if ($i ~ /[*?\[]/) {
          continue
        }
        aliases = aliases $i "\n"
      }
      next
    }

    in_block && tolower($1) == "hostname" {
      has_hostname = 1
      next
    }

    END {
      if (in_block && has_hostname && aliases != "") {
        print aliases
      }
    }
  ' "$cfg" | awk '!seen[$0]++'
}

_copy_files "$@"