#compdef worktree

__worktree_branches() {
  local -a branches
  branches=(
    ${(fo)"$(git branch --format='%(refname:short)' 2>/dev/null)"}
    ${${${${(fo)"$(git for-each-ref --format='%(refname)' refs/remotes 2>/dev/null)"}#refs/remotes/}#*/}:#HEAD}
  )
  _describe 'branch' branches
}

__worktree_active_branches() {
  local -a branches
  branches=(${(fo)"$(git worktree list --porcelain 2>/dev/null | \
    awk '/^worktree /{ wt++ } wt>1 && /^branch refs\/heads\//{ sub(/^branch refs\/heads\//, ""); print }')"})
  _describe 'branch' branches
}

_worktree() {
  local context curcontext="$curcontext" state line
  typeset -A opt_args

  _arguments -C \
    '1: :->cmd' \
    '*:: :->args'

  case $state in
    (cmd)
      local commands; commands=(
        'add:Create a worktree for a branch'
        'remove:Remove a worktree by branch'
        'rm:Remove a worktree by branch'
      )
      _describe -t commands 'worktree command' commands "$@"
      ;;
    (args)
      curcontext="${curcontext%:*:*}:worktree-cmd-$words[1]:"
      case $line[1] in
        (add)
          _arguments '1:: :__worktree_branches'
          ;;
        (rm|remove)
          _arguments '1:: :__worktree_active_branches'
          ;;
      esac
      ;;
  esac
}

_worktree "$@"