wt

#!/bin/zsh
# wt — lightweight git worktree helper
#
# Convention: worktrees live at  <base>/<branch>/<repo-name>
#   e.g.  .../capstone/veloce-trace/feature-x/veloce-trace
# <base> defaults to two levels above the current repo (the grouping dir that
# already holds main/<repo>), so it matches the layout you set up by hand.
# Override per-call with --base DIR, or globally with the WT_BASE env var.
#
# Commands:
#   wt new <branch> [-c] [-e] [--base DIR] [--link-env] [--setup CMD]
#                                            create (or check out) a worktree
#   wt open <branch> [-c] [-e]               open an existing worktree
#   wt ls                                    list worktrees
#   wt rm <branch> [-f] [--rm-branch]        remove a worktree
#   wt help
#
# Flags:
#   -c / --claude   open a new iTerm2 tab in the worktree and start `claude`
#   -e / --code     open the worktree in a VSCode window (`code`)
#   -f / --force    force-remove even if the worktree has uncommitted changes
#   --rm-branch     also delete the local branch after removing the worktree
#   --base DIR      base directory to place the worktree under
#   --link-env      symlink the main worktree's root .env into the new worktree
#   --setup CMD     run CMD inside the new worktree after creating it
self=$0
emulate -L zsh
setopt pipe_fail

die() { print -u2 "wt: $1"; exit 1 }

# Resolve the current repo's top level and derive name + default base.
repo_top=$(git rev-parse --show-toplevel 2>/dev/null) || die "not inside a git repository"
repo_name=${repo_top:t}                       # basename, e.g. veloce-trace
default_base=${WT_BASE:-${repo_top:h:h}}       # two levels up = the grouping dir

# Echo the on-disk path of the worktree currently checked out to <branch>, if any.
worktree_path_for_branch() {
  local target=$1 line wtp=""
  while IFS= read -r line; do
    case $line in
      "worktree "*)                 wtp=${line#worktree } ;;
      "branch refs/heads/$target")  print -r -- "$wtp"; return 0 ;;
    esac
  done < <(git worktree list --porcelain)
  return 1
}

# Open a path in a fresh iTerm2 tab and launch claude there.
open_iterm_claude() {
  osascript - "$1" <<'OSA'
on run argv
  set wtPath to item 1 of argv
  tell application "iTerm2"
    activate
    if (count of windows) = 0 then
      create window with default profile
    else
      tell current window to create tab with default profile
    end if
    tell current session of current window to write text "cd " & quoted form of wtPath & " && claude"
  end tell
end run
OSA
}

# Open a path in VSCode.
open_code() {
  command -v code >/dev/null 2>&1 || { print -u2 "wt: 'code' CLI not found; skipping VSCode"; return }
  code -- "$1"
}

# Echo the path of the repo's main (primary) worktree — the first entry of
# `git worktree list --porcelain`, which git always reports as the main checkout.
main_worktree_path() {
  local line
  while IFS= read -r line; do
    case $line in
      "worktree "*) print -r -- "${line#worktree }"; return 0 ;;
    esac
  done < <(git worktree list --porcelain)
  return 1
}

# Symlink the main worktree's root .env into the new worktree (no copy, so the
# secrets aren't duplicated to disk). Skips silently if there's no .env, and
# never clobbers an existing .env in the new worktree.
link_env() {
  local dest=$1 main
  main=$(main_worktree_path) || { print -u2 "wt: could not locate main worktree; skipping --link-env"; return }
  if [[ ! -e "$main/.env" ]]; then
    print "wt: no .env in main worktree ($main); skipping --link-env"
    return
  fi
  if [[ -e "$dest/.env" || -L "$dest/.env" ]]; then
    print "wt: .env already exists in worktree; leaving it untouched"
    return
  fi
  ln -s "$main/.env" "$dest/.env" && print "wt: linked .env -> $main/.env"
}

cmd_new() {
  local branch="" base=$default_base want_claude=0 want_code=0 link_env=0 setup_cmd=""
  while (( $# )); do
    case $1 in
      -c|--claude) want_claude=1 ;;
      -e|--code)   want_code=1 ;;
      --base)      shift; [[ -n $1 ]] || die "--base requires a directory"; base=$1 ;;
      --link-env)  link_env=1 ;;
      --setup)     shift; [[ -n $1 ]] || die "--setup requires a command"; setup_cmd=$1 ;;
      -*)          die "unknown flag: $1" ;;
      *)           [[ -z $branch ]] && branch=$1 || die "unexpected argument: $1" ;;
    esac
    shift
  done
  [[ -n $branch ]] || die "usage: wt new <branch> [-c] [-e] [--base DIR] [--link-env] [--setup CMD]"
  git check-ref-format "refs/heads/$branch" || die "invalid branch name: $branch"

  local dest="$base/$branch/$repo_name"
  [[ -e $dest ]] && die "destination already exists: $dest"
  mkdir -p "${dest:h}" || die "could not create ${dest:h}"

  # Existing local branch -> check it out; existing remote -> track it; else create new.
  if git show-ref --verify --quiet "refs/heads/$branch"; then
    git worktree add "$dest" "$branch" || die "git worktree add failed"
  elif git show-ref --verify --quiet "refs/remotes/origin/$branch"; then
    git worktree add "$dest" -b "$branch" --track "origin/$branch" || die "git worktree add failed"
  else
    git worktree add "$dest" -b "$branch" || die "git worktree add failed"
  fi

  print "wt: created $dest  (branch: $branch)"

  # Opt-in post-create setup. Both run only when their flag was passed; plain
  # `wt new <branch>` is unchanged.
  (( link_env )) && link_env "$dest"
  if [[ -n $setup_cmd ]]; then
    print "wt: running setup in $dest"
    ( cd "$dest" && eval "$setup_cmd" ) || die "setup command failed: $setup_cmd"
  fi

  (( want_code ))   && open_code "$dest"
  (( want_claude )) && open_iterm_claude "$dest"
  (( want_claude || want_code )) || print "wt: cd \"$dest\""
}

cmd_open() {
  local branch="" want_claude=0 want_code=0
  while (( $# )); do
    case $1 in
      -c|--claude) want_claude=1 ;;
      -e|--code)   want_code=1 ;;
      -*)          die "unknown flag: $1" ;;
      *)           [[ -z $branch ]] && branch=$1 || die "unexpected argument: $1" ;;
    esac
    shift
  done
  [[ -n $branch ]] || die "usage: wt open <branch> [-c] [-e]"
  local dest
  dest=$(worktree_path_for_branch "$branch") || die "no worktree checked out to branch '$branch' (try: wt ls)"
  (( want_code ))   && open_code "$dest"
  (( want_claude )) && open_iterm_claude "$dest"
  (( want_claude || want_code )) || print -r -- "$dest"
}

cmd_ls() { git worktree list }

cmd_rm() {
  local branch="" force=0 rm_branch=0
  while (( $# )); do
    case $1 in
      -f|--force)   force=1 ;;
      --rm-branch)  rm_branch=1 ;;
      -*)           die "unknown flag: $1" ;;
      *)            [[ -z $branch ]] && branch=$1 || die "unexpected argument: $1" ;;
    esac
    shift
  done
  [[ -n $branch ]] || die "usage: wt rm <branch> [-f] [--rm-branch]"
  local dest
  dest=$(worktree_path_for_branch "$branch") || die "no worktree checked out to branch '$branch' (try: wt ls)"

  if (( force )); then
    git worktree remove --force "$dest" || die "could not remove $dest"
  else
    git worktree remove "$dest" || die "worktree has changes; re-run with -f to force"
  fi
  print "wt: removed $dest"
  # Prune empty parent dirs left by the <branch>/<repo> nesting (incl. slashed
  # branch names). rmdir only deletes empty dirs and stops at the first non-empty
  # one, so this can't climb into a populated directory; capped for safety.
  local pdir=${dest:h}
  local -i depth=0
  while (( depth < 4 )) && rmdir "$pdir" 2>/dev/null; do
    print "wt: removed empty $pdir"
    pdir=${pdir:h}
    (( depth++ ))
  done

  if (( rm_branch )); then
    git branch -D "$branch" && print "wt: deleted branch $branch"
  fi
}

cmd_help() {
  cat <<'HELP'
wt — lightweight git worktree helper

Worktrees are placed at  <base>/<branch>/<repo-name>  (base defaults to the
grouping dir two levels above the repo; override with --base or $WT_BASE).

  wt new <branch> [-c] [-e] [--base DIR] [--link-env] [--setup CMD]
                                           create or check out a worktree
  wt open <branch> [-c] [-e]               open an existing worktree
  wt ls                                    list worktrees
  wt rm <branch> [-f] [--rm-branch]        remove a worktree
  wt help

  -c / --claude   open a new iTerm2 tab in the worktree and run `claude`
  -e / --code     open the worktree in a VSCode window
  -f / --force    force-remove a worktree with uncommitted changes
  --rm-branch     also delete the local branch after removal
  --base DIR      base directory to place the worktree under
  --link-env      symlink the main worktree's root .env into the new worktree
  --setup CMD     run CMD inside the new worktree after creating it
HELP
}

case ${1:-help} in
  new)            shift; cmd_new "$@" ;;
  open)           shift; cmd_open "$@" ;;
  ls|list)        shift; cmd_ls "$@" ;;
  rm|remove)      shift; cmd_rm "$@" ;;
  help|-h|--help) cmd_help ;;
  *)              die "unknown command: $1 (try: wt help)" ;;
esac