wt

#!/usr/bin/env bash
set -euo pipefail

prog="${0##*/}"

usage() {
  cat <<EOF
Usage:
  $prog <command> [options]

Commands:
  mk <branch>            Create a sibling worktree for <branch>.
                         If <branch> is a local branch, check it out there.
                         If <branch> is a remote-tracking ref like origin/foo,
                         create local branch foo tracking origin/foo.
                         If <branch> does not exist, create it from HEAD.

  rm [-f] <branch>       Remove the worktree directory for <branch>.
                         Use -f to force removal of a dirty worktree.

  path <branch>          Print the absolute path to the worktree for <branch>.

  ls [--porcelain]       List worktrees attached to this repository.
                         Use --porcelain for machine-readable output.

  prune                  Clean up stale worktree metadata.

  shell-init [shell]     Print shell integration enabling 'wt cd <branch>'.
                         Supported shells: zsh, bash.

  completion [shell]     Print shell completion code.
                         Supported shells: zsh, bash.

  help                   Show this help message.

Examples:
  $prog mk dev
  $prog mk origin/feature/auth
  $prog ls
  $prog ls --porcelain
  eval "\$($prog shell-init zsh)"
  source <($prog completion zsh)
EOF
}

die() {
  printf '%s: %s\n' "$prog" "$*" >&2
  exit 1
}

note() {
  printf '%s: %s\n' "$prog" "$*" >&2
}

require_git_worktree() {
  git rev-parse --is-inside-work-tree >/dev/null 2>&1 \
    || die "not inside a Git working tree"
}

repo_root() {
  git rev-parse --show-toplevel
}

repo_name() {
  basename "$(repo_root)"
}

repo_parent() {
  dirname "$(repo_root)"
}

validate_branch_name() {
  local branch="$1"

  git check-ref-format --branch "$branch" >/dev/null 2>&1 \
    || die "invalid branch name: $branch"

  case "$branch" in
    ""|"."|".."|*..*|*~*|*^*|*:*|*\?*|*\**|*\[*|*\\*|/*)
      die "suspicious branch name: $branch"
      ;;
  esac
}

short_hash() {
  local input="$1"

  if command -v sha1sum >/dev/null 2>&1; then
    printf '%s' "$input" | sha1sum | awk '{print substr($1,1,8)}'
  elif command -v shasum >/dev/null 2>&1; then
    printf '%s' "$input" | shasum -a 1 | awk '{print substr($1,1,8)}'
  elif command -v openssl >/dev/null 2>&1; then
    printf '%s' "$input" | openssl sha1 | sed -E 's/^.*= //' | cut -c1-8
  else
    die "need one of: sha1sum, shasum, or openssl"
  fi
}

sanitize_branch_for_path() {
  local branch="$1"
  local out="$branch"

  out="${out//\//-}"
  out="$(printf '%s' "$out" | tr '[:space:]' '-')"
  out="$(printf '%s' "$out" | sed -E 's/[^A-Za-z0-9._-]+/-/g')"
  out="$(printf '%s' "$out" | sed -E 's/-+/-/g')"
  out="$(printf '%s' "$out" | sed -E 's/^[-.]+//; s/[-.]+$//')"

  [[ -n "$out" ]] || die "branch name sanitizes to empty path component: $branch"
  printf '%s\n' "$out"
}

safe_dir_component_for_branch() {
  local branch="$1"
  local sanitized hash

  sanitized="$(sanitize_branch_for_path "$branch")"

  if [[ "$sanitized" == "$branch" ]]; then
    printf '%s\n' "$sanitized"
  else
    hash="$(short_hash "$branch")"
    printf '%s-%s\n' "$sanitized" "$hash"
  fi
}

canonical_branch_for_spec() {
  local spec="$1"

  if git show-ref --verify --quiet "refs/heads/$spec"; then
    printf '%s\n' "$spec"
    return 0
  fi

  if git show-ref --verify --quiet "refs/remotes/$spec"; then
    local local_branch="${spec#*/}"
    [[ -n "$local_branch" && "$local_branch" != "$spec" ]] \
      || die "could not derive local branch name from remote ref: $spec"
    printf '%s\n' "$local_branch"
    return 0
  fi

  printf '%s\n' "$spec"
}

target_dir_for_branch_spec() {
  local spec="$1"
  local canonical safe_branch
  canonical="$(canonical_branch_for_spec "$spec")"
  safe_branch="$(safe_dir_component_for_branch "$canonical")"
  printf '%s/%s-%s\n' "$(repo_parent)" "$(repo_name)" "$safe_branch"
}

prune_stale() {
  git worktree prune >/dev/null 2>&1 || true
}

local_branch_exists() {
  local branch="$1"
  git show-ref --verify --quiet "refs/heads/$branch"
}

remote_ref_exists() {
  local ref="$1"
  git show-ref --verify --quiet "refs/remotes/$ref"
}

branch_checked_out_elsewhere() {
  local wanted_local_branch="$1"
  local current_root
  current_root="$(repo_root)"

  while IFS= read -r wt; do
    [[ -n "$wt" ]] || continue
    [[ "$wt" == "$current_root" ]] && continue

    local head_ref
    head_ref="$(git -C "$wt" symbolic-ref -q --short HEAD 2>/dev/null || true)"
    if [[ "$head_ref" == "$wanted_local_branch" ]]; then
      printf '%s\n' "$wt"
      return 0
    fi
  done < <(git worktree list --porcelain | awk '/^worktree /{print substr($0,10)}')

  return 1
}

detect_shell() {
  if [[ -n "${1:-}" ]]; then
    printf '%s\n' "$1"
    return
  fi

  if [[ -n "${ZSH_VERSION:-}" ]]; then
    printf 'zsh\n'
  elif [[ -n "${BASH_VERSION:-}" ]]; then
    printf 'bash\n'
  else
    die "specify a shell explicitly: zsh or bash"
  fi
}

cmd_mk() {
  local spec="${1:-}"
  [[ -n "$spec" ]] || die "mk: missing branch name"

  require_git_worktree
  validate_branch_name "$spec"
  prune_stale

  local canonical target_dir
  canonical="$(canonical_branch_for_spec "$spec")"
  target_dir="$(target_dir_for_branch_spec "$spec")"

  [[ ! -e "$target_dir" ]] || die "mk: target already exists: $target_dir"

  local existing=""
  if existing="$(branch_checked_out_elsewhere "$canonical")"; then
    die "mk: branch '$canonical' already checked out at: $existing"
  fi

  note "spec:      $spec"
  note "branch:    $canonical"
  note "target:    $target_dir"

  if local_branch_exists "$canonical"; then
    git worktree add -- "$target_dir" "$canonical"
  elif remote_ref_exists "$spec"; then
    git worktree add -b "$canonical" -- "$target_dir" "$spec"
    git -C "$target_dir" branch --set-upstream-to="$spec" "$canonical" >/dev/null 2>&1 || true
  else
    git worktree add -b "$canonical" -- "$target_dir"
  fi

  printf 'Created worktree: %s\n' "$target_dir"
}

cmd_rm() {
  local force=0

  OPTIND=1
  while getopts ":f" opt; do
    case "$opt" in
      f) force=1 ;;
      \?) die "rm: unknown option -$OPTARG" ;;
    esac
  done
  shift $((OPTIND - 1))

  local spec="${1:-}"
  [[ -n "$spec" ]] || die "rm: missing branch name"

  require_git_worktree
  validate_branch_name "$spec"
  prune_stale

  local target_dir
  target_dir="$(target_dir_for_branch_spec "$spec")"

  [[ -d "$target_dir" ]] || die "rm: no such worktree: $target_dir"

  local current_root
  current_root="$(repo_root)"
  [[ "$target_dir" != "$current_root" ]] \
    || die "rm: refusing to remove the current worktree"

  if (( force )); then
    git worktree remove --force -- "$target_dir"
  else
    git worktree remove -- "$target_dir"
  fi

  printf 'Removed worktree: %s\n' "$target_dir"
}

cmd_path() {
  local spec="${1:-}"
  [[ -n "$spec" ]] || die "path: missing branch name"

  require_git_worktree
  validate_branch_name "$spec"

  local target_dir
  target_dir="$(target_dir_for_branch_spec "$spec")"

  [[ -d "$target_dir" ]] || die "path: no such worktree: $target_dir"
  printf '%s\n' "$target_dir"
}

cmd_ls() {
  require_git_worktree

  case "${1:-}" in
    "")
      git worktree list
      ;;
    --porcelain)
      git worktree list --porcelain
      ;;
    *)
      die "ls: unknown option: $1"
      ;;
  esac
}

cmd_prune() {
  require_git_worktree
  git worktree prune
  printf 'Pruned stale worktree metadata\n'
}

cmd_shell_init() {
  local shell
  shell="$(detect_shell "${1:-}")"

  case "$shell" in
    zsh|bash)
      cat <<'EOF'
wt() {
  local cmd="${1:-}"
  case "$cmd" in
    cd)
      shift
      local target
      target="$(command wt path "$@")" || return
      builtin cd -- "$target"
      ;;
    *)
      command wt "$@"
      ;;
  esac
}
EOF
      ;;
    *)
      die "shell-init: unsupported shell: $shell (supported: zsh, bash)"
      ;;
  esac
}

cmd_completion() {
  local shell
  shell="$(detect_shell "${1:-}")"

  case "$shell" in
    zsh)
      cat <<'EOF'
#compdef wt

_wt() {
  local -a commands shells
  commands=(
    'mk:create a sibling worktree for a branch'
    'rm:remove the worktree for a branch'
    'path:print the path to a worktree'
    'ls:list worktrees'
    'prune:clean stale worktree metadata'
    'shell-init:print shell integration for wt cd'
    'completion:print shell completion code'
    'help:show help'
  )
  shells=(
    'zsh:zsh integration'
    'bash:bash integration'
  )

  _wt_worktree_branches() {
    local -a items
    items=(${(f)"$(
      command wt ls --porcelain 2>/dev/null | awk '
        /^branch refs\/heads\// {
          sub(/^branch refs\/heads\//, "", $0)
          print $0
        }
      '
    )"})
    _describe -t branches 'worktree branch' items
  }

  _wt_mk_branches() {
    local -a items
    items=(${(f)"$(
      {
        git for-each-ref --format='%(refname:short)' refs/heads 2>/dev/null
        git for-each-ref --format='%(refname:short)' refs/remotes 2>/dev/null | grep -vE '/HEAD$'
      } | awk '!seen[$0]++'
    )"})
    _describe -t branches 'branch or remote branch' items
  }

  if (( CURRENT == 2 )); then
    _describe -t commands 'wt command' commands
    return
  fi

  case "$words[2]" in
    mk)
      _wt_mk_branches
      ;;
    rm|path|cd)
      _wt_worktree_branches
      ;;
    ls)
      _arguments '--porcelain[machine-readable output]'
      ;;
    shell-init|completion)
      if (( CURRENT == 3 )); then
        _describe -t shells 'shell' shells
      fi
      ;;
  esac
}

compdef _wt wt
EOF
      ;;
    bash)
      cat <<'EOF'
_wt_completion() {
  local cur prev cmd
  COMPREPLY=()
  cur="${COMP_WORDS[COMP_CWORD]}"
  prev="${COMP_WORDS[COMP_CWORD-1]}"
  cmd="${COMP_WORDS[1]}"

  local commands="mk rm path ls prune shell-init completion help"
  local shells="zsh bash"

  _wt_worktree_branches() {
    command wt ls --porcelain 2>/dev/null | awk '
      /^branch refs\/heads\// {
        sub(/^branch refs\/heads\//, "", $0)
        print $0
      }
    '
  }

  _wt_mk_branches() {
    {
      git for-each-ref --format="%(refname:short)" refs/heads 2>/dev/null
      git for-each-ref --format="%(refname:short)" refs/remotes 2>/dev/null | grep -vE "/HEAD$"
    } | awk '!seen[$0]++'
  }

  if [[ $COMP_CWORD -eq 1 ]]; then
    COMPREPLY=( $(compgen -W "$commands" -- "$cur") )
    return 0
  fi

  case "$cmd" in
    mk)
      COMPREPLY=( $(compgen -W "$(_wt_mk_branches)" -- "$cur") )
      ;;
    rm|path|cd)
      COMPREPLY=( $(compgen -W "$(_wt_worktree_branches)" -- "$cur") )
      ;;
    ls)
      COMPREPLY=( $(compgen -W "--porcelain" -- "$cur") )
      ;;
    shell-init|completion)
      COMPREPLY=( $(compgen -W "$shells" -- "$cur") )
      ;;
  esac

  return 0
}

complete -F _wt_completion wt
EOF
      ;;
    *)
      die "completion: unsupported shell: $shell (supported: zsh, bash)"
      ;;
  esac
}

main() {
  local cmd="${1:-help}"
  shift || true

  case "$cmd" in
    mk)         cmd_mk "$@" ;;
    rm)         cmd_rm "$@" ;;
    path)       cmd_path "$@" ;;
    ls)         cmd_ls "$@" ;;
    prune)      cmd_prune "$@" ;;
    shell-init) cmd_shell_init "$@" ;;
    completion) cmd_completion "$@" ;;
    help|-h|--help) usage ;;
    *) usage >&2; die "unknown command: $cmd" ;;
  esac
}

main "$@"