worktree.sh

#!/usr/bin/env bash
# worktree.sh — create/list/remove per-branch git worktrees with per-worktree envs
#
# Subcommands:
#   worktree.sh new    <branch> [--start <ref>] [--into <dir>] [--no-poetry] [--no-node] [--fast]
#                                [--poetry-args "<args>"] [--python "python3.11"]
#   worktree.sh list
#   worktree.sh remove [<branch>] [--into <dir>] [--path <worktree-path>] [--force]
#
# Behavior (new):
#   - Creates git worktree, then:
#       * If pyproject.toml -> POETRY_VIRTUALENVS_IN_PROJECT=1 poetry install (venv at ./.venv)
#       * Else if requirements*.txt -> python -m venv .venv && pip install -r ...
#       * If package.json -> install via pnpm/yarn/npm (auto-detected)
#   - Per-worktree .venv and node_modules live inside the worktree and are removed with it.

set -euo pipefail

die()  { echo "❌ $*" >&2; exit 1; }
info() { echo "ℹ️  $*"; }
ok()   { echo "✅ $*"; }

print_help() {
  cat <<'EOF'
worktree.sh

Subcommands:
  new     Create a new git worktree + per-worktree environments
  list    List all active git worktrees
  remove  Remove a git worktree and its directory

Usage:
  worktree.sh new <branch> [--start <ref>] [--into <dir>]
                             [--no-poetry] [--no-node] [--fast]
                             [--poetry-args "<args>"] [--python "python3.11"]

  worktree.sh list

  worktree.sh remove [<branch>] [--into <dir>] [--path <worktree-path>] [--force]

Notes:
  • If you run 'remove' with no args, it prints ready-to-copy commands for each active worktree and exits.
  • Per-worktree Python venv is created at ./.venv (Poetry or venv fallback).
  • Node modules are installed in the worktree folder.
  • Removing a worktree removes its .venv and node_modules automatically.
EOF
}

sanitize_dirname() {
  local result
  result="$(printf "%s" "$1" \
    | sed -E 's/[[:space:]]+/_/g; s|/|-|g; s/[^A-Za-z0-9._-]/-/g; s/-{2,}/-/g; s/^[-._]+//; s/[-._]+$//')"
  [[ "$result" =~ \.\. ]] && die "Invalid branch name (contains ..)"
  [[ -z "$result" ]] && die "Invalid branch name (sanitizes to empty)"
  printf "%s" "$result"
}

branch_exists()       { git show-ref --verify --quiet "refs/heads/$1"; }
branch_has_worktree() { git worktree list --porcelain | awk '/^branch /{print $2}' | grep -qx "refs/heads/$1"; }
ref_exists()          { git rev-parse --verify --quiet "$1^{commit}" >/dev/null 2>&1; }

pick_start_ref() {
  if git show-ref --verify --quiet refs/remotes/origin/main; then
    echo "origin/main"
  elif git show-ref --verify --quiet refs/remotes/origin/master; then
    echo "origin/master"
  else
    echo "HEAD"
  fi
}

relpath_or_fallback() {
  local from="$1" to="$2" fallback="$3"
  if command -v realpath >/dev/null 2>&1; then
    realpath --relative-to="$from" "$to" 2>/dev/null || echo "$fallback"
  else
    echo "$fallback"
  fi
}

# -------- installers (used by `new`) ---------------------------------------

install_python_env() {
  # In worktree dir; prefer Poetry+pyproject; else requirements fallback
  local poetry_args="${1:-}"
  local pybin="${PYBIN:-python3}"

  if [ -f pyproject.toml ]; then
    if command -v poetry >/dev/null 2>&1; then
      info "Installing Python deps via Poetry (in-project venv)..."
      POETRY_VIRTUALENVS_IN_PROJECT=1 poetry install --no-interaction ${poetry_args}
      chmod 700 .venv 2>/dev/null || true
      [ -f poetry.lock ] && chmod 600 poetry.lock 2>/dev/null || true
    else
      die "Poetry not found but pyproject.toml present. Install Poetry or use --no-poetry."
    fi
  else
    # Fallback: requirements
    local req=""
    for cand in requirements.txt requirements-dev.txt requirements.in; do
      [ -f "$cand" ] && req="$cand" && break
    done
    if [ -n "$req" ]; then
      info "Creating venv with ${pybin} and installing from ${req}..."
      "${pybin}" -m venv .venv
      # shellcheck disable=SC1091
      source .venv/bin/activate
      python -m pip install --upgrade pip
      pip install -r "$req"
      deactivate || true
      chmod 700 .venv 2>/dev/null || true
    fi
  fi
}

install_node_packages() {
  # In worktree dir
  if [ -f package.json ]; then
    if [ -f pnpm-lock.yaml ] && command -v pnpm >/dev/null 2>&1; then
      info "Installing Node deps via pnpm..."
      pnpm install
    elif [ -f yarn.lock ] && command -v yarn >/dev/null 2>&1; then
      info "Installing Node deps via yarn..."
      yarn install
    elif [ -f package-lock.json ] && command -v npm >/dev/null 2>&1; then
      info "Installing Node deps via npm ci..."
      npm ci
    elif command -v npm >/dev/null 2>&1; then
      info "Installing Node deps via npm..."
      npm install
    else
      die "No supported Node package manager found (npm/yarn/pnpm). Install one or use --no-node."
    fi
  fi
}

# -------- list --------------------------------------------------------------

list_worktrees() {
  local line path branch has_any=0 i=1
  printf "%-4s %-60s %s\n" "#" "PATH" "BRANCH"
  printf "%-4s %-60s %s\n" "----" "------------------------------------------------------------" "---------------------------"
  while IFS= read -r line; do
    case "$line" in
      worktree\ *) path="${line#worktree }" ;;
      branch\ *)   branch="${line#branch }"; branch="${branch#refs/heads/}"
                   printf "%-4s %-60s %s\n" "$i" "$path" "$branch"
                   i=$((i+1)); has_any=1 ;;
      *) ;;
    esac
  done < <(git worktree list --porcelain)
  [ "$has_any" -eq 1 ] || echo "(no worktrees found)"
}

# also used by `remove` with no args
print_remove_suggestions() {
  local line path branch has_any=0
  echo "Active worktrees — copy one of these commands:"
  while IFS= read -r line; do
    case "$line" in
      worktree\ *) path="${line#worktree }" ;;
      branch\ *)   branch="${line#branch }"; branch="${branch#refs/heads/}"
                   echo "  worktree.sh remove $branch"
                   echo "  worktree.sh remove --path \"$path\""
                   echo "" ; has_any=1 ;;
      *) ;;
    esac
  done < <(git worktree list --porcelain)
  [ "$has_any" -eq 1 ] || echo "(no worktrees found)"
}

# -------- new ---------------------------------------------------------------

cmd_new() {
  local BRANCH="$1"; shift
  local START_OVERRIDE="" TARGET_DIR=""
  local DO_POETRY=1 DO_NODE=1 FAST=0 POETRY_ARGS="" PYBIN="python3"

  while [ $# -gt 0 ]; do
    case "$1" in
      --start)        [ $# -lt 2 ] && die "--start requires <ref>"; shift; START_OVERRIDE="$1" ;;
      --into)         [ $# -lt 2 ] && die "--into requires <dir>"; shift; TARGET_DIR="$1" ;;
      --no-poetry)    DO_POETRY=0 ;;
      --no-node)      DO_NODE=0 ;;
      --fast)         FAST=1; DO_POETRY=0; DO_NODE=0 ;;
      --poetry-args)  [ $# -lt 2 ] && die "--poetry-args requires a quoted string"; shift; POETRY_ARGS="$1" ;;
      --python)       [ $# -lt 2 ] && die "--python requires interpreter name/path"; shift; PYBIN="$1" ;;
      -h|--help)      print_help; exit 0 ;;
      *) die "Unknown option to 'new': $1" ;;
    esac
    shift
  done

  local REPO_ROOT="$(git rev-parse --show-toplevel)"
  cd "$REPO_ROOT"

  local DIR_NAME="$(sanitize_dirname "$BRANCH")"
  local WORKTREE_PATH="$REPO_ROOT/${TARGET_DIR:+$TARGET_DIR/}$DIR_NAME"

  local START_REF="${START_OVERRIDE:-$(pick_start_ref)}"
  [ "$START_REF" = "HEAD" ] || ref_exists "$START_REF" || die "Invalid start ref $START_REF"

  [ -e "$WORKTREE_PATH" ] && die "Worktree path already exists: $WORKTREE_PATH"

  cleanup_new() { git worktree remove "$WORKTREE_PATH" 2>/dev/null || true; }
  trap cleanup_new ERR

  if branch_exists "$BRANCH"; then
    branch_has_worktree "$BRANCH" && die "Branch '$BRANCH' already attached to a worktree"
    git worktree add "$WORKTREE_PATH" "$BRANCH"
  else
    git worktree add -b "$BRANCH" "$WORKTREE_PATH" "$START_REF"
  fi

  cd "$WORKTREE_PATH"

  if [ "$FAST" -eq 0 ]; then
    if [ "$DO_POETRY" -eq 1 ] && { [ -f pyproject.toml ] || [ -f requirements.txt ] || [ -f requirements-dev.txt ] || [ -f requirements.in ]; }; then
      PYBIN="$PYBIN" install_python_env "$POETRY_ARGS"
    fi
    if [ "$DO_NODE" -eq 1 ] && [ -f package.json ]; then
      install_node_packages
    fi
  fi

  # Best-effort perms for autoswitchers
  [ -d .venv ] && chmod 700 .venv 2>/dev/null || true
  [ -f poetry.lock ] && chmod 600 poetry.lock 2>/dev/null || true

  trap - ERR
  ok "Worktree ready at $WORKTREE_PATH"
  echo "   • Base: $START_REF"
  [ "$FAST" -eq 1 ] && echo "   • Skipped installs (--fast)"
}

# -------- remove ------------------------------------------------------------

cmd_remove() {
  local BRANCH="" TARGET_DIR="" FORCE=0 PATH_ARG=""

  # Parse args (branch/path optional)
  while [ $# -gt 0 ]; do
    case "$1" in
      --into)  [ $# -lt 2 ] && die "--into requires <dir>"; shift; TARGET_DIR="$1" ;;
      --path)  [ $# -lt 2 ] && die "--path requires <path>"; shift; PATH_ARG="$1" ;;
      --force) FORCE=1 ;;
      -h|--help) print_help; exit 0 ;;
      *) [ -z "$BRANCH" ] && BRANCH="$1" || die "Unexpected arg to 'remove': $1" ;;
    esac
    shift
  done

  local REPO_ROOT="$(git rev-parse --show-toplevel)"
  cd "$REPO_ROOT"

  if [ -z "$BRANCH" ] && [ -z "$PATH_ARG" ]; then
    print_remove_suggestions
    exit 1
  fi

  local WORKTREE_PATH=""
  if [ -n "$PATH_ARG" ]; then
    WORKTREE_PATH="$PATH_ARG"
  else
    local DIR_NAME="$(sanitize_dirname "$BRANCH")"
    WORKTREE_PATH="$REPO_ROOT/${TARGET_DIR:+$TARGET_DIR/}$DIR_NAME"
  fi

  [ -e "$WORKTREE_PATH" ] || die "Worktree path not found: $WORKTREE_PATH"

  if [ "$FORCE" -eq 1 ]; then
    info "git worktree remove --force \"$WORKTREE_PATH\""
    git worktree remove --force "$WORKTREE_PATH"
  else
    info "git worktree remove \"$WORKTREE_PATH\""
    git worktree remove "$WORKTREE_PATH"
  fi

  # Remove directory if it still exists (git usually removes it)
  [ -e "$WORKTREE_PATH" ] && rm -rf "$WORKTREE_PATH"

  git worktree prune >/dev/null 2>&1 || true
  ok "Removed worktree at $WORKTREE_PATH (branch not deleted)"
}

# -------- entry -------------------------------------------------------------

[ $# -ge 1 ] || { print_help; exit 1; }
case "$1" in
  new)     shift; [ $# -ge 1 ] || die "Usage: worktree.sh new <branch> [options]"; cmd_new "$@";;
  list)    shift; list_worktrees ;;
  remove)  shift; cmd_remove "$@" ;;
  -h|--help) print_help ;;
  *) die "Unknown subcommand: $1 (use 'new', 'list', or 'remove')" ;;
esac