tree-me

#!/usr/bin/env bash
# Minimal git worktree helper - leverages git's native capabilities
# Usage: tree-me <command> [options]
#
# For auto-cd and tab completion, add this to your ~/.bashrc or ~/.zshrc:
#   source <(tree-me shellenv)
#
# This enables:
# - Automatic cd to worktree after switch/create/pr commands
# - Tab completion for commands and branch names

set -e

# Worktree organization: ~/dev/worktrees/<repo>/<branch>
WORKTREE_ROOT="${WORKTREE_ROOT:-$HOME/dev/worktrees}"

# Check if Graphite is initialized in this repo
is_graphite_repo() {
    command -v gt >/dev/null 2>&1 && gt repo name >/dev/null 2>&1
}

# Track branch with Graphite, specifying parent
graphite_track() {
    local branch="$1"
    local parent="$2"
    if is_graphite_repo; then
        echo "📊 Tracking branch with Graphite…"
        if [ -n "$parent" ]; then
            gt track "$branch" --parent "$parent" --force 2>/dev/null || gt track "$branch" --force 2>/dev/null || true
        else
            gt track "$branch" --force 2>/dev/null || true
        fi
    fi
}

# Show help if no arguments
show_help() {
    cat << 'EOF'
Usage: tree-me <command> [options]

Git-like worktree management with organized directory structure.
Automatically integrates with Graphite (gt) when available.

Commands:
  switch, sw <branch|pr>        Switch to a branch, or check out a PR (number/url), in new worktree
  create <branch> [base]        Create new branch in worktree (default: main/master)
  pr <number|url>               Checkout GitHub PR in worktree (uses gh)
  list, ls                      List all worktrees
  remove, rm [-f] <branch|pat>  Remove a worktree (quote patterns: "review-*")
  prune                         Remove worktree administrative files
  shellenv                      Output shell function for auto-cd (source this)

Options:
  -f, --force                   For remove: drop a worktree with uncommitted
                                changes, and remove every match of a pattern
                                without confirming. Repeat it (-f -f) to remove
                                a locked worktree, as git itself requires.

Examples:
  tree-me switch feature-branch
  tree-me switch 123
  tree-me create my-feature
  tree-me create my-feature develop
  tree-me pr 123
  tree-me pr https://github.com/org/repo/pull/123
  tree-me list
  tree-me remove old-branch
  tree-me remove -f old-branch
  tree-me remove -f -f locked-branch
  tree-me remove "review-*"
  tree-me prune

Setup auto-cd:
  Add to ~/.bashrc or ~/.zshrc:
    source <(tree-me shellenv)

Worktrees are organized at: $WORKTREE_ROOT/<repo>/<branch>
Set WORKTREE_ROOT to customize the location (default: ~/dev/worktrees)

Graphite Integration:
  When Graphite (gt) is installed and initialized in a repo, tree-me
  automatically tracks new branches with Graphite after creation.
  This enables seamless use of gt commands (modify, submit, sync, etc.)
  in worktrees without manual 'gt track' calls.
EOF
}

if [ $# -eq 0 ]; then
    show_help
    exit 0
fi

# Get repo name from git
get_repo_name() {
    basename "$(git remote get-url origin 2>/dev/null || git rev-parse --show-toplevel)" .git
}

# Source shared helpers
SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
source "$SCRIPT_DIR/lib/git-default-branch.sh"
source "$SCRIPT_DIR/lib/git-worktree.sh"

# Extract PR number from URL or number
get_pr_number() {
    local input="$1"
    if [[ "$input" =~ ^https://github.com/.*/pull/([0-9]+) ]]; then
        echo "${BASH_REMATCH[1]}"
    elif [[ "$input" =~ ^[0-9]+$ ]]; then
        echo "$input"
    else
        echo "Error: Invalid PR number or URL: $input" >&2
        exit 1
    fi
}

# True if input looks like something get_pr_number can parse (a GitHub PR URL
# or a bare number), rather than a branch name.
is_pr_reference() {
    [[ "$1" =~ ^https://github.com/.*/pull/[0-9]+ ]] || [[ "$1" =~ ^[0-9]+$ ]]
}

# Check out a GitHub PR (by number or URL) into a worktree, trying a series of
# fallback names if the preferred one is taken. Shared by the `pr` command and
# by `checkout`/`switch` when given a PR reference instead of a branch name.
checkout_pr() {
    local input="$1"
    local pr_number pr_info head_branch is_fork
    local -a try_names
    local last_candidate max_suffix n candidate existing
    local path branch candidate_path

    pr_number=$(get_pr_number "$input")

    # Check if gh is installed
    if ! command -v gh >/dev/null 2>&1; then
        echo "Error: 'gh' CLI not found. Install it from https://cli.github.com" >&2
        exit 1
    fi

    # Look up the PR's head branch and whether it comes from a fork. Branch
    # names live in a single namespace per repo, so a same-repo PR's head
    # branch is already a unique, stable identity; a fork's branch name has
    # no such guarantee (two forks can both have a branch called "fix"), so
    # it's always prefixed with the PR number. A failed lookup (bad PR
    # number, no auth, no network) is fatal here rather than silently
    # falling through to the retry loop below, which would otherwise burn
    # through every fallback name hitting the same underlying failure.
    pr_info=$(gh pr view "$pr_number" --json headRefName,isCrossRepository \
        -q '[.headRefName, (.isCrossRepository | tostring)] | @tsv' 2>/dev/null || true)
    if [ -z "$pr_info" ]; then
        echo "Error: Failed to look up PR #$pr_number (check the number and your 'gh' auth)" >&2
        exit 1
    fi
    # Split on the first tab with parameter expansion rather than `read`,
    # which strips a leading/trailing tab as IFS whitespace and would
    # misparse an empty headRefName by shifting isCrossRepository into it.
    head_branch="${pr_info%%$'\t'*}"
    is_fork="${pr_info#*$'\t'}"

    if [ -z "$head_branch" ]; then
        try_names=("pr-$pr_number")
    elif [ "$is_fork" = "true" ]; then
        try_names=("pr-$pr_number-$head_branch")
    else
        try_names=("$head_branch" "pr-$pr_number-$head_branch")
    fi

    # Extend with a numbered suffix on the last candidate in case even
    # "pr-<n>-<branch>" is somehow already taken, which happens when a
    # leftover local branch from a prior `tree-me rm` has diverged from
    # the PR: `gh pr checkout` below can't fast-forward it, so the retry
    # loop falls through to a fresh numbered name. Indexed by length
    # rather than "${try_names[-1]}" for portability (negative array
    # subscripts need bash 4.2+).
    last_candidate="${try_names[$((${#try_names[@]} - 1))]}"
    max_suffix=11
    for ((n = 2; n <= max_suffix; n++)); do
        try_names+=("${last_candidate}-${n}")
    done

    # Re-running `tree-me pr` for the same PR should cd back into whichever
    # name it landed on last time, not create a second worktree.
    for candidate in "${try_names[@]}"; do
        existing=$(worktree_path_for "$candidate")
        if [ -n "$existing" ]; then
            echo "✓ Worktree already exists: $existing"
            echo "TREE_ME_CD:$existing"
            exit 0
        fi
    done

    # Create a detached worktree, then let `gh pr checkout` create the
    # branch inside it. Delegating to gh sets up the push/tracking config so
    # that PRs from forks with "Allow edits from maintainers" enabled push
    # back to the contributor's branch. A plain `git fetch pull/N/head`
    # leaves the branch disconnected from the PR.
    #
    # A candidate name can be taken by an unrelated branch or worktree; on
    # failure, fall back to the next name in the list.
    path=""
    branch=""
    for candidate in "${try_names[@]}"; do
        if [ "$candidate" != "${try_names[0]}" ]; then
            echo "ℹ Retrying checkout as '$candidate'…" >&2
        fi
        candidate_path="$WORKTREE_ROOT/$repo/$candidate"
        if ! git worktree add --detach "$candidate_path" >&2; then
            continue
        fi
        if (cd "$candidate_path" && gh pr checkout "$pr_number" --branch "$candidate" >&2); then
            branch="$candidate"
            path="$candidate_path"
            break
        fi
        git worktree remove --force "$candidate_path" 2>/dev/null || true
    done

    if [ -z "$path" ]; then
        echo "Error: Failed to check out PR #$pr_number (tried: ${try_names[*]})" >&2
        exit 1
    fi
    echo "✓ PR #$pr_number checked out at: $path"

    # `gh pr checkout` configures tracking for same-repo PRs and forks with
    # "Allow edits from maintainers" enabled. It leaves the branch with no
    # upstream for forks without that flag, so a bare `git pull` then fails
    # with "no tracking information for the current branch". Backstop it:
    # every PR head is reachable as refs/pull/<N>/head on origin, so point
    # the branch there when nothing else set an upstream. This is a no-op
    # when gh already configured tracking (the common case).
    if [ -z "$(git -C "$path" config --get "branch.$branch.merge")" ]; then
        git -C "$path" config "branch.$branch.remote" origin
        git -C "$path" config "branch.$branch.merge" "refs/pull/$pr_number/head"
        echo "ℹ Set '$branch' to pull from origin refs/pull/$pr_number/head" >&2
    fi

    # Track with Graphite if available (use --force to auto-detect parent)
    (cd "$path" && graphite_track "$branch" "")

    echo "TREE_ME_CD:$path"
}

repo=$(get_repo_name)
command="$1"
shift

case "$command" in
    shellenv)
        cat << 'EOF'
tree-me() {
    local output
    output=$(command tree-me "$@")
    local exit_code=$?
    echo "$output"
    if [ $exit_code -eq 0 ]; then
        local cd_path=$(echo "$output" | grep "^TREE_ME_CD:" | cut -d: -f2-)
        [ -n "$cd_path" ] && cd "$cd_path"
    fi
    return $exit_code
}

# Bash completion
if [ -n "$BASH_VERSION" ]; then
    _tree_me_complete() {
        local cur commands
        COMPREPLY=()
        cur="${COMP_WORDS[COMP_CWORD]}"
        commands="switch sw create pr list ls remove rm prune help shellenv"

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

        # Complete branch names for switch/remove/rm. Keyed off the subcommand
        # rather than the previous word so a flag before the branch (rm -f) does
        # not turn completion off.
        case "${COMP_WORDS[1]}" in
            co|checkout|switch|sw|remove|rm)
                local candidates
                candidates=$(git worktree list 2>/dev/null | sed -n 's/.*\[\([^]]*\)\].*/\1/p' | tail -n +2)
                # Offered only once a dash is typed, so a lone branch still
                # completes on its own rather than tying with the flag.
                case "${COMP_WORDS[1]}" in
                    remove|rm)
                        [[ "$cur" == -* ]] && candidates="$candidates -f"
                        ;;
                esac
                COMPREPLY=( $(compgen -W "$candidates" -- "$cur") )
                return 0
                ;;
        esac
    }
    complete -F _tree_me_complete tree-me
fi

# Zsh completion
if [ -n "$ZSH_VERSION" ]; then
    _tree_me_complete_zsh() {
        local -a commands branches
        commands=(
            'switch:Switch to an existing branch in new worktree'
            'sw:Switch to an existing branch in new worktree'
            'create:Create new branch in worktree'
            'pr:Checkout GitHub PR in worktree'
            'list:List all worktrees'
            'ls:List all worktrees'
            'remove:Remove a worktree'
            'rm:Remove a worktree'
            'prune:Remove worktree administrative files'
            'help:Show help'
            'shellenv:Output shell function for auto-cd'
        )

        if (( CURRENT == 2 )); then
            _describe 'command' commands
        else
            # Not pinned to CURRENT == 3, so a flag before the branch (rm -f)
            # does not turn completion off.
            case "$words[2]" in
                checkout|co|switch|sw|remove|rm)
                    branches=(${(f)"$(git worktree list 2>/dev/null | sed -n 's/.*\[\([^]]*\)\].*/\1/p' | tail -n +2)"})
                    # Offered only once a dash is typed, so a lone branch still
                    # completes on its own rather than tying with the flag.
                    case "$words[2]" in
                        remove|rm)
                            [[ "$words[CURRENT]" == -* ]] && \
                                branches+=('-f:Force removal; repeat for a locked worktree')
                            ;;
                    esac
                    _describe 'branch' branches
                    ;;
            esac
        fi
    }
    compdef _tree_me_complete_zsh tree-me
fi
EOF
        ;;

    checkout|co|switch|sw)
        branch="${1:?Branch name required. Usage: tree-me switch <branch>}"
        path="$WORKTREE_ROOT/$repo/$branch"

        existing=$(worktree_path_for "$branch")
        if [ -n "$existing" ]; then
            echo "✓ Worktree already exists: $existing"
            echo "TREE_ME_CD:$existing"
            exit 0
        fi

        # Check if branch exists
        if git show-ref --verify --quiet "refs/heads/$branch" || \
           git show-ref --verify --quiet "refs/remotes/origin/$branch"; then
            git worktree add "$path" "$branch"
            echo "✓ Worktree created at: $path"

            # Track with Graphite if available (use --force to auto-detect parent)
            (cd "$path" && graphite_track "$branch" "")

            echo "TREE_ME_CD:$path"
        elif is_pr_reference "$branch"; then
            # No matching branch, but this looks like a PR reference, so
            # check it out the same way `tree-me pr` does.
            checkout_pr "$branch"
        else
            echo "Error: Branch '$branch' does not exist" >&2
            echo "Use 'tree-me create $branch' to create a new branch" >&2
            exit 1
        fi
        ;;

    create)
        branch="${1:?Branch name required. Usage: tree-me create <branch> [base-branch]}"
        base="${2:-$(get_default_branch)}"
        path="$WORKTREE_ROOT/$repo/$branch"

        existing=$(worktree_path_for "$branch")
        if [ -n "$existing" ]; then
            echo "✓ Worktree already exists: $existing"
            echo "TREE_ME_CD:$existing"
            exit 0
        fi

        # Create new branch
        git worktree add "$path" -b "$branch" "$base"
        echo "✓ Worktree created at: $path"

        # Track with Graphite if available
        (cd "$path" && graphite_track "$branch" "$base")

        echo "TREE_ME_CD:$path"
        ;;

    pr)
        input="${1:?PR number or URL required. Usage: tree-me pr <number|url>}"
        checkout_pr "$input"
        ;;

    list|ls)
        git worktree list
        ;;

    remove|rm)
        # -f is repeatable and forwarded to git as-is, so git's own advice
        # ("use 'remove -f -f' to override") is directly runnable here: one -f
        # removes a worktree with uncommitted changes, two also removes a
        # locked one. It doubles as the "don't ask" flag for patterns.
        force=0
        args=()
        while [ $# -gt 0 ]; do
            case "$1" in
                -f|--force) force=$((force + 1)) ;;
                -ff) force=$((force + 2)) ;;
                --) shift; args+=("$@"); break ;;
                -*)
                    echo "Error: Unknown option '$1'" >&2
                    echo "Usage: tree-me remove [-f] <branch|pattern>" >&2
                    exit 1
                    ;;
                *) args+=("$1") ;;
            esac
            shift
        done

        if [ ${#args[@]} -eq 0 ]; then
            echo "Error: Branch name or pattern required. Usage: tree-me remove [-f] <branch|pattern>" >&2
            exit 1
        fi
        if [ ${#args[@]} -gt 1 ]; then
            echo "Error: Unexpected argument '${args[1]}' (quote patterns: \"review-*\")" >&2
            echo "Usage: tree-me remove [-f] <branch|pattern>" >&2
            exit 1
        fi
        pattern="${args[0]}"

        force_args=()
        [ "$force" -ge 1 ] && force_args+=(--force)
        [ "$force" -ge 2 ] && force_args+=(--force)

        declare -A branch_to_path=()
        branches=()
        while IFS=$'\t' read -r b p; do
            branch_to_path[$b]=$p
            branches+=("$b")
        done < <(list_worktrees_excluding_main)

        is_glob=0
        [[ "$pattern" == *[*?\[]* ]] && is_glob=1

        matches=()
        for b in "${branches[@]}"; do
            # shellcheck disable=SC2053
            [[ "$b" == $pattern ]] && matches+=("$b")
        done

        if [ ${#matches[@]} -eq 0 ]; then
            if [ "$is_glob" = 1 ]; then
                echo "Error: No worktrees match pattern: $pattern" >&2
            else
                echo "Error: No worktree found for branch: $pattern" >&2
            fi
            exit 1
        fi

        if [ "$is_glob" = 1 ] && [ "$force" -eq 0 ]; then
            echo "Matching worktrees:"
            for b in "${matches[@]}"; do
                echo "  $b → ${branch_to_path[$b]}"
            done
            # `read` reports EOF as a failure even when it read an answer that
            # simply lacked a trailing newline, so keep whatever it did read
            # and let the default-to-no check below decide. Under `set -e` a
            # bare EOF would otherwise abort the script instead of declining.
            read -r -p "Remove ${#matches[@]} worktree(s)? [y/N] " response || true
            if [[ ! "$response" =~ ^[Yy]$ ]]; then
                echo "Aborted"
                exit 0
            fi
        fi

        # One match git refuses (locked, or dirty without -f) shouldn't strand
        # the rest of the batch unattempted, so report it and carry on, with
        # the final status reflecting that something was left behind.
        failed=0
        for b in "${matches[@]}"; do
            path="${branch_to_path[$b]}"
            if git worktree remove "${force_args[@]}" "$path"; then
                echo "✓ Removed worktree: $path"
            else
                echo "✗ Skipped worktree: $path" >&2
                failed=1
            fi
        done
        [ "$failed" -eq 0 ] || exit 1
        ;;

    prune)
        git worktree prune
        echo "✓ Pruned stale worktree administrative files"
        ;;

    help|--help|-h)
        show_help
        ;;

    *)
        echo "Error: Unknown command '$command'" >&2
        echo "Run 'tree-me help' for usage information" >&2
        exit 1
        ;;
esac