# Emit one "path<TAB>branch<TAB>head" row per worktree, main clone first,
# branch empty when the worktree is detached.
__worktree_list() {
  git worktree list --porcelain | awk '
    /^worktree /{ p = substr($0, 10); b = ""; h = "" }
    /^HEAD /{ h = substr($0, 6) }
    /^branch refs\/heads\//{ b = $0; sub(/^branch refs\/heads\//, "", b) }
    /^$/{ if (p != "") { print p "\t" b "\t" h; p = "" } }
    END{ if (p != "") print p "\t" b "\t" h }
  '
}

# Fill the caller's paths, names, dirs and mismatched arrays with an entry per
# worktree. names is the checked out branch, or its short hash when detached.
# dirs is @<suffix> for a worktree directory of the form <main clone>@<suffix>,
# and the directory name otherwise. mismatched is 1 when that suffix and the
# branch differ. With -m, only worktrees with such a suffix are returned.
__worktree_labels() {
  local managed_only=0
  [[ "$1" == "-m" ]] && managed_only=1
  local -a rows
  rows=("${(@f)$(__worktree_list)}")
  local main=${rows[1]%%$'\t'*}
  local prefix="${main:h}/${main:t}@"
  local row p b h name suffix dir mm
  for row in "${rows[@]}"; do
    [[ -z "$row" ]] && continue
    p=${row%%$'\t'*}
    b=${${row#*$'\t'}%%$'\t'*}
    h=${row##*$'\t'}
    if [[ -n "$b" ]]; then
      name=$b
    else
      name="(detached ${h[1,7]})"
    fi
    if [[ "${p[1,${#prefix}]}" == "$prefix" ]]; then
      suffix=${p[${#prefix}+1,-1]}
      dir="@$suffix"
      if [[ "$b" == "$suffix" ]]; then
        mm=0
      else
        mm=1
      fi
    elif (( managed_only )); then
      continue
    else
      dir=${p:t}
      mm=0
    fi
    paths+=("$p")
    names+=("$name")
    dirs+=("$dir")
    mismatched+=($mm)
  done
}

__worktree_find_in() {
  local dir=$1 pat=$2 icase=$3
  if (( $+commands[fd] )); then
    local -a fd_args=(
      --base-directory "$dir" --hidden --no-ignore --exclude .git --glob
    )
    if (( icase )); then
      fd_args+=(--ignore-case)
    else
      fd_args+=(--case-sensitive)
    fi
    if [[ "$pat" == */* ]]; then
      fd "${fd_args[@]}" --full-path "**/$pat"
    else
      fd "${fd_args[@]}" "$pat"
    fi
  else
    local ic=
    (( icase )) && ic=i
    local -a expr=(-${ic}name "$pat")
    [[ "$pat" == */* ]] && expr=(-${ic}path "*/$pat")
    (cd "$dir" && find . -name .git -prune -o "${expr[@]}" -print)
  fi
}

worktree() {
  if [[ "$1" == "" ]]; then
    echo "usage: worktree {add,cd,list,remove,sync} <branch> [<branch>...]
       worktree find [-i] <glob>"
    return 1
  elif [[ "$1" == "-h" ]] || [[ "$1" == "--help" ]]; then
    echo "usage: worktree {add,cd,list,remove,sync} <branch> [<branch>...]
       worktree find [-i] <glob>

Manage git worktrees by branch name."
    return 0
  fi

  local cmd=$1
  local branch=$2

  if ! git rev-parse --git-dir &>/dev/null; then
    print -P "%F{red}error:%f not a git repository"
    return 1
  fi

  if [[ "$cmd" != "cd" ]] && [[ "$cmd" != "find" ]] && \
     [[ "$(git rev-parse --is-inside-work-tree 2>/dev/null)" == "true" ]] && \
     [[ "$(git rev-parse --absolute-git-dir)" != "$(git rev-parse --show-toplevel)/.git" ]]
  then
    print -P "%F{red}error:%f must be run from the main clone, not a worktree"
    return 1
  fi

  case $cmd in
    add)
      if [[ -z "$branch" ]]; then
        print -P "%F{red}error:%f missing <branch> argument"
        return 1
      fi
      local root=$(git rev-parse --show-toplevel)
      local wt_dest=${root:h}/${root:t}@${branch}
      if ! git show-ref --verify --quiet refs/heads/"$branch" && \
         ! git show-ref --verify --quiet refs/remotes/origin/"$branch"; then
        git branch "$branch" origin/main || return 1
      fi
      git worktree add "$wt_dest" "$branch" || return 1
      if [ -f $root/.enter ] && [ -f $root/.exit ]; then
        cp $root/.enter $wt_dest/.enter
        cp $root/.exit $wt_dest/.exit
        _autoenv_authorized $wt_dest/.enter yes
        _autoenv_authorized $wt_dest/.exit yes
      fi
      if [ -f $root/local.bazelrc ]; then
        cp $root/local.bazelrc $wt_dest/local.bazelrc
      fi
      ;;

    cd)
      if [[ -z "$branch" ]]; then
        print -P "%F{red}error:%f missing <branch> argument"
        return 1
      fi
      local wt_path
      wt_path=$(git worktree list --porcelain | awk -v b="$branch" \
        '/^worktree /{ sub(/^worktree /, ""); p=$0 } /^branch refs\/heads\//{ sub(/^branch refs\/heads\//, ""); if($0==b) print p }')
      if [[ -z "$wt_path" ]]; then
        print -P "%F{red}error:%f no worktree for branch: $branch"
        return 1
      fi
      cd "$wt_path"
      ;;

    find)
      local icase=0
      local -a globs
      local arg
      for arg in "${@:2}"; do
        case "$arg" in
          -i|--ignore-case) icase=1 ;;
          *) globs+=("$arg") ;;
        esac
      done
      if (( ${#globs} == 0 )); then
        print -P "%F{red}error:%f missing <glob> argument"
        return 1
      fi
      if (( ${#globs} > 1 )); then
        print -P "%F{red}error:%f unexpected argument: ${globs[2]}"
        return 1
      fi
      local pattern=${globs[1]}
      # Matching is on the tail of a path, so leading and trailing slashes
      # carry no meaning.
      local pat=$pattern
      while [[ "$pat" == /* ]]; do pat=${pat#/}; done
      while [[ "$pat" == */ ]]; do pat=${pat%/}; done
      if [[ -z "$pat" ]]; then
        print -P "%F{red}error:%f invalid <glob> argument: $pattern"
        return 1
      fi

      local -a paths names dirs mismatched
      __worktree_labels

      # Search every worktree at once, each into its own file so that the
      # results don't interleave.
      local REPORTTIME=-1
      local tmp
      tmp=$(mktemp -d "${TMPDIR:-/tmp}/worktree-find.XXXXXX") || return 1
      local -a results
      local i
      {
        (
          for (( i = 1; i <= ${#paths}; i++ )); do
            __worktree_find_in "${paths[i]}" "$pat" "$icase" | sort >"$tmp/$i" &
          done
          wait
        )
        for (( i = 1; i <= ${#paths}; i++ )); do
          results+=("$(<$tmp/$i)")
        done
      } always {
        rm -rf "$tmp"
      }

      # Report every worktree with matches under the same heading list prints,
      # followed by its matches as absolute paths.
      local c_dir c_warn c_off
      if [[ -t 1 ]]; then
        c_dir=$(print -Pn '%F{8}')
        c_warn=$(print -Pn '%F{yellow}')
        c_off=$(print -Pn '%f')
      fi
      local found=0 line root c
      for (( i = 1; i <= ${#paths}; i++ )); do
        [[ -n "${results[i]}" ]] || continue
        (( found )) && print
        found=1
        root=${paths[i]}
        [[ "$root" == "$HOME" || "$root" == "$HOME"/* ]] && root="~${root#$HOME}"
        if (( mismatched[i] )); then
          c=$c_warn
        else
          c=$c_dir
        fi
        print -r -- "${names[i]}  ${c}${dirs[i]}${c_off}"
        for line in "${(@f)${results[i]}}"; do
          print -r -- "  $root/${line#./}"
        done
      done

      if (( ! found )); then
        print -u2 -- "no matches: $pattern"
        return 1
      fi
      ;;

    list|ls)
      local -a paths names dirs mismatched
      __worktree_labels -m

      (( ${#names} )) || return 0

      local maxw=0 name
      for name in "${names[@]}"; do
        (( ${#name} > maxw )) && maxw=${#name}
      done

      local c_dir c_warn c_off
      if [[ -t 1 ]]; then
        c_dir=$(print -Pn '%F{8}')
        c_warn=$(print -Pn '%F{yellow}')
        c_off=$(print -Pn '%f')
      fi
      local i c
      for (( i = 1; i <= ${#names}; i++ )); do
        if (( mismatched[i] )); then
          c=$c_warn
        else
          c=$c_dir
        fi
        print -r -- "  ${(r:$maxw:)names[i]}  ${c}${dirs[i]}${c_off}"
      done
      ;;

    rm|remove)
      local force=0
      local branches=()
      local arg
      for arg in "${@:2}"; do
        case "$arg" in
          -f|--force) force=1 ;;
          *) branches+=("$arg") ;;
        esac
      done
      if (( ${#branches[@]} == 0 )); then
        print -P "%F{red}error:%f missing <branch> argument"
        return 1
      fi
      local root=$(git rev-parse --show-toplevel)
      local wt_list
      wt_list=$(git worktree list --porcelain)
      local failed=0
      local wt_path
      local err
      for branch in "${branches[@]}"; do
        wt_path=$(echo "$wt_list" | awk -v b="$branch" \
          '/^worktree /{ sub(/^worktree /, ""); p=$0 } /^branch refs\/heads\//{ sub(/^branch refs\/heads\//, ""); if($0==b) print p }')
        if [[ -z "$wt_path" ]]; then
          print -P "%F{yellow}warning:%f no worktree for branch: $branch"
          continue
        fi
        if [[ -f "$wt_path/BUILD" ]] || [[ -f "$wt_path/WORKSPACE" ]]; then
          if (( $+commands[bazel] )); then
            bazel --output_base="$wt_path" clean --expunge
          elif [[ -x "$wt_path/bazelw" ]]; then
            "$wt_path/bazelw" clean --expunge
          fi
        fi
        local -a rm_args=()
        (( force )) && rm_args+=(--force)
        if err=$(git worktree remove "${rm_args[@]}" "$wt_path" 2>&1); then
          :
        elif [[ "$err" == *"working trees containing submodules"* ]]; then
          if (( ! force )) && \
             [[ -n "$(git -C "$wt_path" status --porcelain --ignore-submodules=all 2>/dev/null)" ]]; then
            print -P "%F{red}error:%f worktree has uncommitted changes (use --force): $wt_path"
            failed=1
            continue
          fi
          if ! rm -rf "$wt_path" || ! git worktree prune; then
            failed=1
            continue
          fi
        else
          print -- "$err" >&2
          failed=1
          continue
        fi
        local parent=${wt_path:h}
        while [[ "$parent" != "${root:h}" ]] && [[ -d "$parent" ]]; do
          rmdir "$parent" 2>/dev/null || break
          parent=${parent:h}
        done
      done
      return $failed
      ;;

    sync)
      local root=$(git rev-parse --show-toplevel)
      local prefix="${root:h}/${root:t}@"
      local wt_list
      wt_list=$(git worktree list --porcelain)
      local branches=("${@:2}")
      if (( ${#branches[@]} == 0 )); then
        branches=("${(@f)$(echo "$wt_list" | awk -v prefix="$prefix" '
          /^worktree /{ sub(/^worktree /, ""); p=$0 }
          /^branch refs\/heads\//{ if(index(p, prefix)==1) { sub(/^branch refs\/heads\//, ""); print } }
        ')}")
      fi
      local failed=0
      local wt_path
      for branch in "${branches[@]}"; do
        [[ -z "$branch" ]] && continue
        wt_path=$(echo "$wt_list" | awk -v b="$branch" \
          '/^worktree /{ sub(/^worktree /, ""); p=$0 } /^branch refs\/heads\//{ sub(/^branch refs\/heads\//, ""); if($0==b) print p }')
        if [[ -z "$wt_path" ]]; then
          print -P "%F{yellow}warning:%f no worktree for branch: $branch"
          failed=1
          continue
        fi
        if [[ "$wt_path" == "$root" ]]; then
          continue
        fi
        for f in .enter .exit; do
          [[ -f "$root/$f" ]] || continue
          if [[ -L "$wt_path/$f" ]]; then
            rm "$wt_path/$f"
          elif [[ -f "$wt_path/$f" ]] && cmp -s "$root/$f" "$wt_path/$f"; then
            continue
          fi
          cp "$root/$f" "$wt_path/$f"
          _autoenv_authorized "$wt_path/$f" yes
          print -P "%F{green}updated:%f $wt_path/$f"
        done
        if [[ -f "$root/local.bazelrc" ]]; then
          if [[ -L "$wt_path/local.bazelrc" ]]; then
            rm "$wt_path/local.bazelrc"
          elif [[ -f "$wt_path/local.bazelrc" ]] && cmp -s "$root/local.bazelrc" "$wt_path/local.bazelrc"; then
            continue
          fi
          cp "$root/local.bazelrc" "$wt_path/local.bazelrc"
          print -P "%F{green}updated:%f $wt_path/local.bazelrc"
        fi
      done
      return $failed
      ;;

    *)
      print -P "%F{red}error:%f unknown command: $cmd"
      return 1
      ;;
  esac
}

# A <glob> is matched against every worktree, never the current directory.
alias worktree='noglob worktree'