Robust zsh git worktree management function

wt

# Git worktree management function
# Creates worktrees in ../<repo-name>-wip/<branch> structure
# Example: In ./ayamara-backend, 'wt feature/auth' creates ../ayamara-backend-wip/feature/auth
#
# Usage:
#   wt <branch>           - Create new worktree with new branch and cd into it
#   wt -e <branch>        - Create worktree from existing branch
#   wt -b <base> <branch> - Create new branch from <base>
#   wt ls                 - List all worktrees
#   wt rm <branch>        - Remove worktree
#   wt sw <branch>        - Switch to existing worktree

emulate -L zsh
setopt local_options extended_glob

local -a opts_e opts_b
local base_branch existing_branch subcommand branch_name worktree_path
local git_root parent_dir repo_name wip_dir

# Parse options
zparseopts -D -E - e=opts_e b:=opts_b

# Check if we're in a git repo
if ! git rev-parse --git-dir &>/dev/null; then
  print -u2 "Error: not in a git repository"
  return 1
fi

# Get git root for path calculations
git_root=$(git rev-parse --show-toplevel)
parent_dir="${git_root:h}"
repo_name="${git_root:t}"
wip_dir="${parent_dir}/${repo_name}-wip"

# Handle subcommands
if [[ $# -eq 0 ]]; then
  print -u2 "Usage: wt <branch> | wt [-e] [-b base] <branch> | wt {ls|rm|sw} [args]"
  return 1
fi

subcommand="$1"

case "$subcommand" in
  ls)
    # List all worktrees
    git worktree list
    ;;

  rm)
    # Remove worktree
    if [[ $# -lt 2 ]]; then
      print -u2 "Error: wt rm requires a branch name"
      return 1
    fi
    branch_name="$2"
    worktree_path="${wip_dir}/${branch_name}"

    # Check if worktree exists
    if ! git worktree list | grep -q "${worktree_path}"; then
      print -u2 "Error: worktree '${branch_name}' not found"
      return 1
    fi

    git worktree remove "$worktree_path" && \
      print "Removed worktree: ${branch_name}"
    ;;

  sw)
    # Switch to existing worktree
    if [[ $# -lt 2 ]]; then
      print -u2 "Error: wt sw requires a branch name"
      return 1
    fi
    branch_name="$2"
    worktree_path="${wip_dir}/${branch_name}"

    # Check if worktree exists
    if [[ ! -d "$worktree_path" ]]; then
      print -u2 "Error: worktree directory '${worktree_path}' does not exist"
      return 1
    fi

    builtin cd "$worktree_path" && \
      print "Switched to worktree: ${branch_name}"
    ;;

  *)
    # Default: create new worktree
    branch_name="$1"

    # Validate branch name (no spaces)
    if [[ "$branch_name" == *" "* ]]; then
      print -u2 "Error: branch name cannot contain spaces"
      return 1
    fi

    worktree_path="${wip_dir}/${branch_name}"

    # Check if worktree path already exists
    if [[ -e "$worktree_path" ]]; then
      print -u2 "Error: path '${worktree_path}' already exists"
      return 1
    fi

    # Check if worktree already exists for this branch
    if git worktree list | grep -q "\\[${branch_name}\\]"; then
      print -u2 "Error: worktree for branch '${branch_name}' already exists"
      return 1
    fi

    # Create wip directory and any parent directories needed for nested branches
    mkdir -p "${worktree_path:h}"

    # Build git worktree add command
    local -a git_cmd
    git_cmd=(git worktree add)

    if (( ${#opts_e} )); then
      # Use existing branch
      if ! git show-ref --verify --quiet "refs/heads/${branch_name}"; then
        print -u2 "Error: branch '${branch_name}' does not exist"
        return 1
      fi
      git_cmd+=("$worktree_path" "$branch_name")
    elif (( ${#opts_b} )); then
      # Create from specific base branch
      base_branch="${opts_b[2]}"
      if ! git show-ref --verify --quiet "refs/heads/${base_branch}"; then
        print -u2 "Error: base branch '${base_branch}' does not exist"
        return 1
      fi
      git_cmd+=("$worktree_path" -b "$branch_name" "$base_branch")
    else
      # Create new branch from current HEAD
      git_cmd+=("$worktree_path" -b "$branch_name")
    fi

    # Create worktree
    if "${git_cmd[@]}"; then
      print "Created worktree: ${branch_name} at ${worktree_path}"

      if builtin cd "$worktree_path"; then
        print "Switched to: ${worktree_path}"

        # Link gitignored .env* files from main worktree
        local envfile filename target
        find "$git_root" -maxdepth 1 \( -type f -o -type l \) -name ".env*" 2>/dev/null | while IFS= read -r envfile; do
          filename="${envfile:t}"

          # Check if this file pattern is gitignored in main
          if (builtin cd "$git_root" && git check-ignore "$filename" &>/dev/null); then
            if [[ -L "$envfile" ]]; then
              # It's a symlink - recreate the same symlink
              target=$(readlink "$envfile")
              ln -sf "$target" "$filename"
              print "✓ Linked ${filename} -> ${target} (symlink, gitignored)"
            else
              # It's a regular file - create symlink to it
              ln -sf "$envfile" "$filename"
              print "✓ Linked ${filename} -> ${envfile} (file, gitignored)"
            fi
          fi
        done
      else
        return 1
      fi
    else
      return 1
    fi
    ;;
esac