worktree.sh

#!/usr/bin/env bash

worktree() {
    local subcommand=$1
    [[ $# -gt 0 ]] && shift

    case $subcommand in
        a|add)
            # Add new worktree (git worktree add)
            __worktree_add "$@"
            ;;
        r|remove)
            # Remove worktree (git worktree remove)
            __worktree_remove "$@"
            ;;
        l|list)
            # List and switch worktrees (interactive)
            __worktree_list "$@"
            ;;
        *)
            echo "Usage: worktree <command>"
            echo ""
            echo "Commands:"
            echo "  a, add       Add a new worktree from branch"
            echo "  r, remove    Remove existing worktree(s)"
            echo "  l, list      List and switch between worktrees"
            return 1
            ;;
    esac
}

__worktree_validate_environment() {
    # Check if fzf is available
    if ! command -v fzf >/dev/null 2>&1; then
        echo "Error: fzf is not installed"
        return 1
    fi
    
    # Check if we're in a git repository
    if ! git rev-parse --git-dir >/dev/null 2>&1; then
        echo "Error: Not in a git repository"
        return 1
    fi
    
    return 0
}

__worktree_add() {
    # Validate environment
    __worktree_validate_environment || return $?

    # Fetch the latest branches from all remotes
    echo "Fetching latest branches from all remotes..."
    git fetch --all

    # Get both local and remote branches, removing duplicates
    local branches
    branches=$(git branch -a --format='%(refname:short)' | grep -v 'HEAD' | sed 's|^origin/||' | sort -u)
    if [[ -z "$branches" ]]; then
        echo "Error: Failed to retrieve branches"
        return 1
    fi

    # Let the user select a branch
    local selected_branch
    selected_branch=$(echo "$branches" | fzf --height=70% --border --layout=reverse --preview 'git log --oneline --graph --color=always {} -20 2>/dev/null || git log --oneline --graph --color=always origin/{} -20' --preview-window=right:60%:wrap)

    if [[ -z "$selected_branch" ]]; then
        echo "No branch selected"
        return 1
    fi

    # Extract the branch name (already cleaned by the format command)
    local branch_name=$selected_branch

    # Get the worktree path
    local current_dir=$(pwd)
    local parent_dir=$(dirname "$current_dir")
    local worktree_path="$parent_dir/$branch_name"
    
    # Validate parent directory is writable
    if [[ ! -w "$parent_dir" ]]; then
        echo "Error: Cannot create worktree in $parent_dir (permission denied)"
        return 1
    fi

    # Check if the worktree already exists
    if [[ -d "$worktree_path" ]]; then
        echo "Worktree already exists at $worktree_path"
        echo -n "Do you want to switch to it? (y/n): "
        read response
        if [[ "$response" = "y" ]]; then
            cd "$worktree_path"
            echo "Switched to worktree: $worktree_path"
        fi
        return 0
    fi

    # Check if the local branch exists
    if git show-ref --verify --quiet "refs/heads/$branch_name"; then
        # Local branch exists, create worktree from it
        git worktree add "$worktree_path" "$branch_name"
    elif git show-ref --verify --quiet "refs/remotes/origin/$branch_name"; then
        # Remote branch exists, create local branch tracking the remote
        git worktree add -b "$branch_name" "$worktree_path" "origin/$branch_name"
    else
        # New branch, create it in the worktree
        git worktree add -b "$branch_name" "$worktree_path"
    fi

    if [[ $? -eq 0 ]]; then
        echo "Created worktree: $worktree_path"
        cd "$worktree_path"
        echo "Switched to branch: $branch_name"
    else
        echo "Failed to create worktree"
        return 1
    fi
}

__worktree_remove() {
    # Validate environment
    __worktree_validate_environment || return $?

    # Get the current worktree path
    local current_worktree=$(git rev-parse --show-toplevel)

    # Get the main worktree path
    local main_worktree=$(git worktree list | head -n1 | awk '{print $1}')

    # Get all worktrees except the current one and the main one
    local worktrees
    worktrees=$(git worktree list | awk '{print $1}' | grep -v "^$current_worktree$" | grep -v "^$main_worktree$")

    if [[ -z "$worktrees" ]]; then
        echo "No worktrees available to delete"
        echo "Note: Current worktree ($(basename "$current_worktree")) and main worktree ($(basename "$main_worktree")) cannot be deleted"
        return 0
    fi

    # Let the user select worktrees to delete (multi-select with TAB)
    local selected_worktrees
    selected_worktrees=$(echo "$worktrees" | fzf --multi --height=70% --border --layout=reverse --preview 'cd {} && git log --oneline --graph --color=always -20' --preview-window=right:60%:wrap --header="Select worktrees to delete (TAB to multi-select)")

    if [[ -z "$selected_worktrees" ]]; then
        echo "No worktrees selected"
        return 0
    fi

    # Confirm deletion
    echo "Selected worktrees to delete:"
    echo "$selected_worktrees"
    echo -n "Are you sure you want to delete these worktrees? (y/n): "
    read response

    if [[ "$response" = "y" ]]; then
        while IFS= read -r worktree; do
            echo "Deleting worktree: $worktree"
            git worktree remove --force "$worktree"
            if [[ $? -eq 0 ]]; then
                echo "Successfully deleted: $worktree"
            else
                echo "Failed to delete: $worktree"
            fi
        done <<< "$selected_worktrees"
    else
        echo "Deletion cancelled"
    fi
}

__worktree_list() {
    # Validate environment
    __worktree_validate_environment || return $?

    # Get the current worktree
    local current_worktree=$(git rev-parse --show-toplevel)

    # Get all worktrees and format them nicely
    local worktrees
    worktrees=$(git worktree list | awk -v current="$current_worktree" '{
        path = $1
        branch = substr($0, index($0, "[") + 1)
        sub(/\].*/, "", branch)
        if (path == current) {
            printf "* %s [%s]\n", path, branch
        } else {
            printf "  %s [%s]\n", path, branch
        }
    }')

    # Select a worktree using fzf
    local selected
    selected=$(echo "$worktrees" | fzf --height=70% --border --layout=reverse \
        --preview 'bash -c "path=\$(echo {} | sed \"s/^[* ] //; s/ \[.*\]//\"); cd \$path && git log --oneline --graph --color=always -20"' \
        --preview-window=right:60%:wrap \
        --header="Current: $current_worktree")

    if [[ -z "$selected" ]]; then
        echo "No worktree selected"
        return 1
    fi

    # Extract the path from the selection
    local selected_path=$(echo "$selected" | sed 's/^[* ] //; s/ \[.*\]//')

    # Change to the selected worktree
    cd "$selected_path"
    echo "Switched to worktree: $selected_path"

    # Show the current branch
    local branch=$(git branch --show-current)
    echo "Current branch: $branch"
}