wti

#!/bin/bash

# A helper to manage worktrees.
#
# Usage:
#   wti add [<branch>] - Add a new worktree.
#   wti rm             - Remove a worktree.
#   wti list           - List worktrees.

# look for a global variable WORKTREE_PATH which will default to ~/.worktrees
WORKTREE_PATH="${WORKTREE_PATH:-$HOME/.worktrees}"

# Ensure the worktree path exists
mkdir -p "$WORKTREE_PATH"

# Add a worktree.
add_worktree() {
    local branch="$1"
    
    if [ -z "$branch" ]; then
        # fzf all branches, and include a "Create new" option
        # The user can select an existing branch or type a new name and press Enter.
        branch=$( (git branch -a | sed 's/remotes\/origin\///' | sed 's/^\* //' | sed 's/  //' | sort -u) | fzf --print-query --prompt="Select a branch or enter a new branch name: " --header="Select an existing branch or type a new name and press Enter." --height 40% --reverse | tail -n1)
    fi

    if [ -z "$branch" ]; then
        echo "No branch selected. Exiting."
        return
    fi

    # a worktree is created in a directory with the same name as the branch
    local worktree_dir="$WORKTREE_PATH/$branch"
    if [ -d "$worktree_dir" ]; then
        echo "Worktree directory '$worktree_dir' already exists. Exiting."
        return
    fi

    echo "Creating worktree for '$branch' at '$worktree_dir'..."

    # if provided a branch which only exists on remote, have this behave similarly to checking out a remote branch
    if git show-ref --verify --quiet "refs/remotes/origin/$branch" && ! git show-ref --verify --quiet "refs/heads/$branch"; then
        echo "Branch '$branch' found on remote. Creating a local tracking branch."
        git worktree add "$worktree_dir" "origin/$branch"
    # if an argument is provided, check whether it's an existing branch
    elif git show-ref --verify --quiet "refs/heads/$branch"; then
        echo "Branch '$branch' exists locally."
        git worktree add "$worktree_dir" "$branch"
    # otherwise assume it's a new branch
    else
        echo "Branch '$branch' does not exist. Creating it."
        git worktree add -b "$branch" "$worktree_dir"
    fi

    # add some tactful user logging here to ensure that the user's aware of what's going on
    if [ $? -eq 0 ]; then
        echo "Worktree for branch '$branch' added at '$worktree_dir'."
    else
        echo "Failed to create worktree."
    fi
}

# Remove a worktree.
remove_worktree() {
    local worktree_to_remove
    # fzf all worktrees and allow removing the selected worktree
    worktree_to_remove=$(git worktree list | awk '{print $1 " " $3}' | fzf --prompt="Select a worktree to remove: " --height 40% --reverse | awk '{print $1}')
    
    if [ -z "$worktree_to_remove" ]; then
        echo "No worktree selected. Exiting."
        return
    fi

    echo "Removing worktree '$worktree_to_remove'..."
    git worktree remove "$worktree_to_remove"
    if [ $? -eq 0 ]; then
        echo "Worktree '$worktree_to_remove' removed."
    else
        echo "Failed to remove worktree."
    fi
}

# List and print on a worktree.
list_worktrees() {
    local worktree_to_switch
    # fzf all worktrees and allow printing info on the selected one
    worktree_to_switch=$(git worktree list | awk '{print $1 " " $3}' | fzf --prompt="Select a worktree: " --height 40% --reverse | awk '{print $1}')
    
    if [ -n "$worktree_to_switch" ]; then
        echo "$worktree_to_switch"
    fi
}

# Main script logic
case "$1" in
    add)
        add_worktree "$2"
        ;;
    rm)
        remove_worktree
        ;;
    list)
        list_worktrees
        ;;
    *)
        echo "Usage: $0 {add|rm|list|cd|code}"
        exit 1
        ;;
esac

exit 0