Shell function for managing git worktrees with credential key symlinking

.worktree.sh

# Shell function for creating and managing git worktrees
# Usage: worktree add <branch-name> [--setup]
#        worktree setup
#        worktree start
#        worktree stop
#        worktree list
#        worktree remove <branch-name> [--wip|--force]

WORKTREE_PORT_REGISTRY="$HOME/.worktree-ports"
WORKTREE_BASE_RAILS_PORT=3000
WORKTREE_BASE_VITE_PORT=3036
WORKTREE_MAX_DB_NAME_LENGTH=63

worktree() {
    local ACTION="$1"
    shift

    case "$ACTION" in
        add)
            _worktree_add "$@"
            ;;
        setup)
            _worktree_setup "$@"
            ;;
        start)
            _worktree_start "$@"
            ;;
        stop)
            _worktree_stop "$@"
            ;;
        restart)
            _worktree_restart "$@"
            ;;
        info|status)
            _worktree_info "$@"
            ;;
        run)
            _worktree_run "$@"
            ;;
        console)
            _worktree_console "$@"
            ;;
        open)
            _worktree_open "$@"
            ;;
        logs)
            _worktree_logs "$@"
            ;;
        connect)
            _worktree_connect "$@"
            ;;
        cd)
            _worktree_cd "$@"
            ;;
        list|ls)
            _worktree_list "$@"
            ;;
        prune)
            _worktree_prune "$@"
            ;;
        remove|rm)
            _worktree_remove "$@"
            ;;
        help|--help|-h|"")
            _worktree_usage
            ;;
        *)
            echo "Unknown command: $ACTION"
            echo ""
            _worktree_usage
            return 1
            ;;
    esac
}

_worktree_usage() {
    echo "Usage: worktree <command> [options]"
    echo ""
    echo "Commands:"
    echo "  add <branch> [--setup]  Create a new worktree (optionally run full setup)"
    echo "  setup                   Clone DB and run bin/update (run from within worktree)"
    echo "  start [-D]              Start dev server with unique ports (-D to daemonize)"
    echo "  stop                    Stop dev server (overmind)"
    echo "  restart                 Stop and start dev server"
    echo "  info                    Show current worktree's URL and config"
    echo "  run <command>           Run a command with worktree env vars loaded"
    echo "  console                 Open Rails console with worktree env"
    echo "  open                    Open Rails URL in browser"
    echo "  logs                    View overmind logs"
    echo "  connect <process>       Connect to overmind process (web, vite, worker)"
    echo "  cd <name>               Jump to a worktree by name"
    echo "  list                    List all worktrees with their ports and databases"
    echo "  prune                   Clean up stale entries from port registry"
    echo "  remove <branch>         Remove worktree [--wip|--force]"
    echo "  help                    Show this help message"
    echo ""
    echo "Workflow:"
    echo "  worktree add my-feature      # Create worktree, cd into it"
    echo "  worktree setup               # Clone DB, run bin/update"
    echo "  worktree start               # Start Rails + Vite on unique ports"
    echo "  worktree stop                # Stop servers"
    echo "  worktree info                # Show URL and config"
    echo "  worktree list                # Show all worktrees"
    echo "  worktree remove my-feature   # Remove worktree, optionally drop DB"
    echo ""
    echo "Quick start:"
    echo "  worktree add my-feature --setup   # Create + full setup in one command"
}

# Sanitize branch name for use in database name
_worktree_sanitize_branch() {
    local sanitized=$(echo "$1" | sed 's/[^a-zA-Z0-9]/_/g')

    # Truncate if needed to fit within PostgreSQL's 63-char limit
    # Account for prefix "unclecharlie_development_" (25 chars)
    local max_branch_length=$((WORKTREE_MAX_DB_NAME_LENGTH - 25))

    if [ ${#sanitized} -gt $max_branch_length ]; then
        # Use first part + hash of full name for uniqueness
        local hash=$(echo "$1" | md5sum | cut -c1-8)
        local truncated_length=$((max_branch_length - 9)) # 8 for hash + 1 for underscore
        sanitized="${sanitized:0:$truncated_length}_${hash}"
    fi

    echo "$sanitized"
}

# Get or assign a port slot for a worktree
_worktree_get_slot() {
    local WORKTREE_PATH="$1"
    local SLOT

    # Create registry if it doesn't exist
    touch "$WORKTREE_PORT_REGISTRY"

    # Check if this worktree already has a slot (use exact match with delimiter)
    SLOT=$(awk -F: -v path="$WORKTREE_PATH" '$1 == path {print $2}' "$WORKTREE_PORT_REGISTRY")

    if [ -z "$SLOT" ]; then
        # Find the next available slot (1-99)
        SLOT=1
        while awk -F: '{print $2}' "$WORKTREE_PORT_REGISTRY" | grep -qx "$SLOT"; do
            SLOT=$((SLOT + 1))
        done
        # Register this worktree
        echo "${WORKTREE_PATH}:${SLOT}" >> "$WORKTREE_PORT_REGISTRY"
    fi

    echo "$SLOT"
}

# Remove a worktree from the port registry
_worktree_release_slot() {
    local WORKTREE_PATH="$1"
    if [ -f "$WORKTREE_PORT_REGISTRY" ]; then
        awk -F: -v path="$WORKTREE_PATH" '$1 != path' "$WORKTREE_PORT_REGISTRY" > "${WORKTREE_PORT_REGISTRY}.tmp"
        mv "${WORKTREE_PORT_REGISTRY}.tmp" "$WORKTREE_PORT_REGISTRY"
    fi
}

# Check if we're in the main repo (not a worktree)
_worktree_is_main_repo() {
    local git_dir=$(git rev-parse --git-dir 2>/dev/null)
    # In a worktree, git-dir is .git/worktrees/<name>, in main repo it's .git
    [[ "$git_dir" == ".git" ]]
}

# Get the main repo path from a worktree
_worktree_get_main_repo() {
    git rev-parse --path-format=absolute --git-common-dir 2>/dev/null | sed 's/\/.git$//'
}

_worktree_add() {
    local BRANCH_NAME="$1"
    local FLAG="$2"

    if [ -z "$BRANCH_NAME" ]; then
        echo "Usage: worktree add <branch-name> [--setup]"
        return 1
    fi

    # Check if we're in a git repo
    if ! git rev-parse --is-inside-work-tree &>/dev/null; then
        echo "Error: Not inside a git repository"
        return 1
    fi

    # Check if we're in the main repo, not a worktree
    if ! _worktree_is_main_repo; then
        echo "Error: You're inside a worktree, not the main repository."
        echo "Please run 'worktree add' from the main repository:"
        echo "  cd $(_worktree_get_main_repo)"
        return 1
    fi

    local CURRENT_DIR=$(basename "$(pwd)")
    local WORKTREE_PATH="../${CURRENT_DIR}-${BRANCH_NAME}"
    local MAIN_BRANCH
    local CURRENT_BRANCH

    # Determine the main branch (master or main)
    if git show-ref --verify --quiet refs/heads/main; then
        MAIN_BRANCH="main"
    elif git show-ref --verify --quiet refs/heads/master; then
        MAIN_BRANCH="master"
    else
        echo "Error: Could not find 'main' or 'master' branch"
        return 1
    fi

    # Check if we're on the main branch
    CURRENT_BRANCH=$(git branch --show-current)
    if [ "$CURRENT_BRANCH" != "$MAIN_BRANCH" ]; then
        echo "Switching to $MAIN_BRANCH branch..."
        git checkout "$MAIN_BRANCH" || return 1
    fi

    # Check if worktree path already exists
    if [ -d "$WORKTREE_PATH" ]; then
        echo "Error: Worktree path already exists: $WORKTREE_PATH"
        return 1
    fi

    # Create worktree with new or existing branch
    if git show-ref --verify --quiet "refs/heads/$BRANCH_NAME"; then
        echo "Branch '$BRANCH_NAME' exists, checking it out..."
        git worktree add "$WORKTREE_PATH" "$BRANCH_NAME" || return 1
    else
        echo "Creating new branch '$BRANCH_NAME'..."
        git worktree add -b "$BRANCH_NAME" "$WORKTREE_PATH" || return 1
    fi

    # Symlink .key files from config/credentials (they're gitignored)
    local CREDENTIALS_DIR="config/credentials"
    if [ -d "$CREDENTIALS_DIR" ]; then
        for keyfile in "$CREDENTIALS_DIR"/*.key; do
            if [ -f "$keyfile" ]; then
                local keyfile_name=$(basename "$keyfile")
                local keyfile_abs=$(cd "$CREDENTIALS_DIR" && pwd)/"$keyfile_name"
                ln -s "$keyfile_abs" "$WORKTREE_PATH/$CREDENTIALS_DIR/$keyfile_name"
                echo "Symlinked: $keyfile_name"
            fi
        done
    fi

    # Create .overmind.env with port and DB config
    local WORKTREE_ABS_PATH=$(cd "$WORKTREE_PATH" && pwd)
    local SLOT=$(_worktree_get_slot "$WORKTREE_ABS_PATH")
    local RAILS_PORT=$((WORKTREE_BASE_RAILS_PORT + SLOT))
    local VITE_PORT=$((WORKTREE_BASE_VITE_PORT + SLOT))
    local SANITIZED_BRANCH=$(_worktree_sanitize_branch "$BRANCH_NAME")
    local DEV_DB="unclecharlie_development_${SANITIZED_BRANCH}"
    local TEST_DB="unclecharlie_test_${SANITIZED_BRANCH}"

    cat > "$WORKTREE_PATH/.overmind.env" << EOF
# Worktree-specific configuration (auto-generated by worktree add)
DB_NAME=$DEV_DB
TEST_DB_NAME=$TEST_DB
PORT=$RAILS_PORT
VITE_RUBY_PORT=$VITE_PORT
EOF

    # Create Procfile.local that uses env vars instead of hardcoded ports
    # Note: Redis is excluded - it should run separately as a shared service
    cat > "$WORKTREE_PATH/Procfile.local" << 'EOF'
web: WEB_CONCURRENCY=1 RUBY_DEBUG_OPEN=true bin/rails s -p ${PORT:-3000}
vite: bin/vite dev --clobber
worker: bin/sidekiq
EOF

    echo ""
    echo "Worktree created at: $WORKTREE_PATH"
    echo "Branch: $BRANCH_NAME"
    echo "  Database: $DEV_DB"
    echo "  Test DB:  $TEST_DB"
    echo "  Rails port: $RAILS_PORT"
    echo "  Vite port: $VITE_PORT"

    # Change to the new worktree directory
    cd "$WORKTREE_PATH"

    # Run setup if --setup flag provided
    if [ "$FLAG" = "--setup" ]; then
        echo ""
        _worktree_setup
    else
        echo ""
        echo "Tests can be run now. Run 'worktree setup' to clone the dev database and install dependencies."
    fi
}

_worktree_setup() {
    # Check for .overmind.env (created by worktree add)
    if [ ! -f ".overmind.env" ]; then
        echo "Error: .overmind.env not found. Are you in a worktree created with 'worktree add'?"
        return 1
    fi

    # Source .overmind.env to get DB name
    source .overmind.env

    local SOURCE_DB="unclecharlie_development"
    local TARGET_DB="$DB_NAME"

    echo "Setting up worktree..."
    echo "  Cloning database: $TARGET_DB"
    echo ""

    # Check if target database already exists
    if psql -lqt | cut -d \| -f 1 | grep -qw "$TARGET_DB"; then
        echo "Database '$TARGET_DB' already exists, skipping clone."
    else
        echo "Cloning database '$SOURCE_DB' to '$TARGET_DB'..."

        # Create the database
        createdb "$TARGET_DB" || return 1

        # Dump and restore
        pg_dump "$SOURCE_DB" | psql -q "$TARGET_DB" || return 1

        echo "Database cloned successfully."
    fi

    # Run bin/update
    echo ""
    echo "Running bin/update..."
    bin/update || return 1

    echo ""
    echo "Setup complete!"
    echo "Run 'worktree start' to start the dev server."
}

_worktree_start() {
    local DAEMONIZE=""
    if [ "$1" = "-D" ] || [ "$1" = "-d" ]; then
        DAEMONIZE="-D"
    fi

    # Check for .overmind.env
    if [ ! -f ".overmind.env" ]; then
        echo "Error: .overmind.env not found. Are you in a worktree created with 'worktree add'?"
        return 1
    fi

    # Check for Procfile.local
    if [ ! -f "Procfile.local" ]; then
        echo "Error: Procfile.local not found. Are you in a worktree created with 'worktree add'?"
        return 1
    fi

    # Ensure Redis is running (shared service)
    if ! redis-cli ping &>/dev/null; then
        echo "Starting Redis..."
        redis-server /usr/local/etc/redis.conf --daemonize yes
        sleep 1
        if redis-cli ping &>/dev/null; then
            echo "Redis started."
        else
            echo "Warning: Failed to start Redis. Some features may not work."
        fi
    fi

    # Source .overmind.env to get ports for display
    source .overmind.env

    echo "Starting dev server..."
    echo "  Rails: http://localhost:${PORT}"
    echo "  Vite:  http://localhost:${VITE_RUBY_PORT}"
    if [ -n "$DAEMONIZE" ]; then
        echo "  Mode:  daemonized (use 'worktree stop' to stop)"
    fi
    echo ""

    # Check for overmind or foreman
    if command -v overmind &> /dev/null; then
        # Use short names to avoid tmux path length limits
        local SHORT_ID=$(echo "$(pwd)" | md5sum | cut -c1-8)
        local OVERMIND_SOCK="./.overmind-${SHORT_ID}.sock"
        local TMUX_SOCK="/tmp/overmind-${SHORT_ID}"

        # Clean up stale socket if overmind isn't actually running
        if [ -e "$OVERMIND_SOCK" ]; then
            if ! OVERMIND_SOCKET="$OVERMIND_SOCK" overmind echo &>/dev/null; then
                echo "Cleaning up stale socket..."
                rm -f "$OVERMIND_SOCK"
            else
                echo "Error: Overmind is already running. Use 'worktree stop' first."
                return 1
            fi
        fi

        # Use Procfile.local which has env var substitution for ports
        # Use short title to avoid tmux socket path length limit
        OVERMIND_SOCKET="$OVERMIND_SOCK" overmind start -f Procfile.local -w "wt-${SHORT_ID}" $DAEMONIZE
    elif command -v foreman &> /dev/null; then
        if [ -n "$DAEMONIZE" ]; then
            echo "Warning: Foreman does not support daemonization. Running in foreground."
        fi
        foreman start -f Procfile.local --env .overmind.env
    else
        echo "Error: Neither overmind nor foreman found. Please install one."
        return 1
    fi
}

_worktree_stop() {
    # Use the same short socket names as start
    local SHORT_ID=$(echo "$(pwd)" | md5sum | cut -c1-8)
    local OVERMIND_SOCK="./.overmind-${SHORT_ID}.sock"

    if command -v overmind &> /dev/null && [ -e "$OVERMIND_SOCK" ]; then
        OVERMIND_SOCKET="$OVERMIND_SOCK" overmind quit
        echo "Stopped overmind."
    elif command -v overmind &> /dev/null && [ -e ".overmind.sock" ]; then
        # Fallback for old socket naming
        overmind quit
        echo "Stopped overmind."
    else
        echo "No running overmind process found."
        echo "If using foreman, press Ctrl+C in the terminal running it."
    fi
}

_worktree_restart() {
    _worktree_stop
    sleep 1
    _worktree_start "$@"
}

_worktree_console() {
    _worktree_run bin/rails console "$@"
}

_worktree_open() {
    # Check for .overmind.env
    if [ ! -f ".overmind.env" ]; then
        echo "Error: .overmind.env not found. Are you in a worktree created with 'worktree add'?"
        return 1
    fi

    source .overmind.env
    local url="http://localhost:${PORT}"
    echo "Opening $url"
    open "$url"
}

_worktree_logs() {
    local SHORT_ID=$(echo "$(pwd)" | md5sum | cut -c1-8)
    local OVERMIND_SOCK="./.overmind-${SHORT_ID}.sock"

    if [ ! -e "$OVERMIND_SOCK" ]; then
        echo "Error: Overmind is not running. Use 'worktree start' first."
        return 1
    fi

    OVERMIND_SOCKET="$OVERMIND_SOCK" overmind echo
}

_worktree_connect() {
    local PROCESS="$1"

    if [ -z "$PROCESS" ]; then
        echo "Usage: worktree connect <process>"
        echo "Processes: web, vite, worker"
        return 1
    fi

    local SHORT_ID=$(echo "$(pwd)" | md5sum | cut -c1-8)
    local OVERMIND_SOCK="./.overmind-${SHORT_ID}.sock"

    if [ ! -e "$OVERMIND_SOCK" ]; then
        echo "Error: Overmind is not running. Use 'worktree start' first."
        return 1
    fi

    OVERMIND_SOCKET="$OVERMIND_SOCK" overmind connect "$PROCESS"
}

_worktree_cd() {
    local INPUT="$1"

    if [ -z "$INPUT" ]; then
        echo "Usage: worktree cd <worktree-name>"
        return 1
    fi

    # Find the worktree path from registry
    if [ ! -f "$WORKTREE_PORT_REGISTRY" ]; then
        echo "Error: No worktrees registered."
        return 1
    fi

    local path slot found=""
    while IFS=: read -r path slot; do
        local short_path="${path##*/}"
        if [ "$short_path" = "$INPUT" ] || [ "${path##*/}" = "$INPUT" ]; then
            if [ -d "$path" ]; then
                found="$path"
                break
            fi
        fi
    done < "$WORKTREE_PORT_REGISTRY"

    if [ -z "$found" ]; then
        echo "Error: Worktree '$INPUT' not found."
        echo "Use 'worktree list' to see available worktrees."
        return 1
    fi

    cd "$found"
}

_worktree_prune() {
    if [ ! -f "$WORKTREE_PORT_REGISTRY" ]; then
        echo "No worktrees registered."
        return 0
    fi

    local path slot removed=0
    local temp_file="${WORKTREE_PORT_REGISTRY}.tmp"
    > "$temp_file"

    while IFS=: read -r path slot; do
        if [ -d "$path" ]; then
            echo "${path}:${slot}" >> "$temp_file"
        else
            echo "Removing stale entry: ${path##*/}"
            removed=$((removed + 1))
        fi
    done < "$WORKTREE_PORT_REGISTRY"

    mv "$temp_file" "$WORKTREE_PORT_REGISTRY"
    echo "Pruned $removed stale entries."
}

_worktree_info() {
    # Check for .overmind.env
    if [ ! -f ".overmind.env" ]; then
        if _worktree_is_main_repo; then
            echo "You're in the main repo. Use 'worktree list' to see all worktrees."
        else
            echo "Error: .overmind.env not found. Are you in a worktree created with 'worktree add'?"
        fi
        return 1
    fi

    # Source .overmind.env to get config
    source .overmind.env

    local BRANCH=$(git branch --show-current)

    echo ""
    echo "Worktree: $(basename "$(pwd)")"
    echo "Branch:   $BRANCH"
    echo ""
    echo "URL:      http://localhost:${PORT}"
    echo ""
    echo "Database: $DB_NAME"
    echo "Test DB:  $TEST_DB_NAME"
    echo "Vite:     http://localhost:${VITE_RUBY_PORT}"
    echo ""
}

_worktree_run() {
    if [ $# -eq 0 ]; then
        echo "Usage: worktree run <command>"
        echo "Example: worktree run bin/rails db:migrate"
        return 1
    fi

    # Check for .overmind.env
    if [ ! -f ".overmind.env" ]; then
        echo "Error: .overmind.env not found. Are you in a worktree created with 'worktree add'?"
        return 1
    fi

    # Source .overmind.env and export all variables, then run the command
    set -a
    source .overmind.env
    set +a
    "$@"
}

_worktree_list() {
    echo "Worktrees:"
    echo ""

    if [ ! -f "$WORKTREE_PORT_REGISTRY" ] || [ ! -s "$WORKTREE_PORT_REGISTRY" ]; then
        echo "  No worktrees registered."
        return 0
    fi

    local path slot rails_port vite_port db_name short_path max_path_len=4

    # First pass: find the longest path name
    while IFS=: read -r path slot; do
        if [ -d "$path" ]; then
            short_path="${path##*/}"
            if [ ${#short_path} -gt $max_path_len ]; then
                max_path_len=${#short_path}
            fi
        fi
    done < "$WORKTREE_PORT_REGISTRY"

    # Print header with dynamic width
    printf "  %-${max_path_len}s  %-12s %-12s %s\n" "PATH" "RAILS PORT" "VITE PORT" "DATABASE"
    printf "  %-${max_path_len}s  %-12s %-12s %s\n" "----" "----------" "---------" "--------"

    # Second pass: print the data
    while IFS=: read -r path slot; do
        if [ -d "$path" ]; then
            rails_port=$((WORKTREE_BASE_RAILS_PORT + slot))
            vite_port=$((WORKTREE_BASE_VITE_PORT + slot))
            db_name=""

            # Try to read DB name from .overmind.env
            if [ -f "$path/.overmind.env" ]; then
                db_name=$(/usr/bin/grep "^DB_NAME=" "$path/.overmind.env" | /usr/bin/cut -d= -f2)
            fi

            # Shorten path for display
            short_path="${path##*/}"

            printf "  %-${max_path_len}s  %-12s %-12s %s\n" "$short_path" "$rails_port" "$vite_port" "$db_name"
        fi
    done < "$WORKTREE_PORT_REGISTRY"

    echo ""
}

_worktree_remove() {
    local INPUT="$1"
    local FLAG="$2"

    if [ -z "$INPUT" ]; then
        echo "Usage: worktree remove <branch-name|worktree-dir> [--wip|--force]"
        return 1
    fi

    local CURRENT_DIR=$(basename "$(pwd)")
    local WORKTREE_PATH
    local MAIN_REPO_PATH
    local BRANCH_NAME
    local HAD_CHANGES=false

    # Determine the main repo name (either current dir or derived from worktree name)
    local MAIN_REPO_NAME
    if _worktree_is_main_repo; then
        MAIN_REPO_NAME="$CURRENT_DIR"
    else
        # We're in a worktree, get main repo name
        MAIN_REPO_NAME=$(basename "$(_worktree_get_main_repo)")
    fi

    # Check if input is a full directory name (starts with main repo prefix)
    if [[ "$INPUT" == "${MAIN_REPO_NAME}-"* ]]; then
        # Extract branch name by removing the prefix
        BRANCH_NAME="${INPUT#${MAIN_REPO_NAME}-}"
    else
        BRANCH_NAME="$INPUT"
    fi

    # Check if we're in the worktree we want to remove
    if [[ "$CURRENT_DIR" == "${MAIN_REPO_NAME}-${BRANCH_NAME}" ]]; then
        # We're in the worktree, derive main repo path
        MAIN_REPO_PATH="../${MAIN_REPO_NAME}"
        WORKTREE_PATH="$(pwd)"

        echo "Currently in worktree, switching to main repo..."
        cd "$MAIN_REPO_PATH" || return 1
    else
        # We're in the main repo or a different worktree
        if _worktree_is_main_repo; then
            WORKTREE_PATH="../${CURRENT_DIR}-${BRANCH_NAME}"
            MAIN_REPO_PATH="$(pwd)"
        else
            MAIN_REPO_PATH="$(_worktree_get_main_repo)"
            WORKTREE_PATH="${MAIN_REPO_PATH}/../${MAIN_REPO_NAME}-${BRANCH_NAME}"
            cd "$MAIN_REPO_PATH" || return 1
        fi
    fi

    # Convert to absolute path for registry
    WORKTREE_PATH=$(cd "$WORKTREE_PATH" 2>/dev/null && pwd) || {
        echo "Error: Worktree does not exist"
        return 1
    }

    # Check for uncommitted changes
    if git -C "$WORKTREE_PATH" status --porcelain | grep -q .; then
        HAD_CHANGES=true

        # If no flag provided, show changes and prompt
        if [ -z "$FLAG" ]; then
            echo "Worktree has uncommitted changes:"
            git -C "$WORKTREE_PATH" status --short
            echo ""
            echo "WARNING: Any stashes in this worktree will be lost!"
            echo ""
            echo "How would you like to proceed?"
            echo "  1) wip    - Create a WIP commit before removing (recommended)"
            echo "  2) force  - Discard all changes and remove"
            echo "  3) cancel - Abort removal"
            echo ""
            read "?Choose [1-3]: " choice

            case "$choice" in
                1|wip)    FLAG="--wip" ;;
                2|force)  FLAG="--force" ;;
                3|cancel) echo "Aborted."; return 0 ;;
                *)        echo "Invalid choice. Aborted."; return 1 ;;
            esac
        fi

        # Handle the flag
        case "$FLAG" in
            --wip)
                echo "Creating WIP commit..."
                git -C "$WORKTREE_PATH" add -A || return 1
                git -C "$WORKTREE_PATH" commit -m "WIP: Work in progress on $BRANCH_NAME" || return 1
                ;;
            --force)
                echo "Warning: Discarding uncommitted changes and any stashes"
                ;;
            *)
                echo "Unknown flag: $FLAG"
                echo "Use --wip or --force"
                return 1
                ;;
        esac
    fi

    # Check for worktree databases
    local SANITIZED_BRANCH=$(_worktree_sanitize_branch "$BRANCH_NAME")
    local WORKTREE_DB="unclecharlie_development_${SANITIZED_BRANCH}"
    local WORKTREE_TEST_DB="unclecharlie_test_${SANITIZED_BRANCH}"
    local DBS_TO_DROP=()

    if psql -lqt | cut -d \| -f 1 | grep -qw "$WORKTREE_DB"; then
        DBS_TO_DROP+=("$WORKTREE_DB")
    fi
    if psql -lqt | cut -d \| -f 1 | grep -qw "$WORKTREE_TEST_DB"; then
        DBS_TO_DROP+=("$WORKTREE_TEST_DB")
    fi

    if [ ${#DBS_TO_DROP[@]} -gt 0 ]; then
        echo ""
        echo "Found worktree databases:"
        for db in "${DBS_TO_DROP[@]}"; do
            echo "  - $db"
        done
        read "?Drop these databases? [y/N]: " drop_db
        if [[ "$drop_db" =~ ^[Yy]$ ]]; then
            for db in "${DBS_TO_DROP[@]}"; do
                dropdb "$db"
                echo "Dropped: $db"
            done
        else
            echo "Databases kept."
        fi
    fi

    # Release the port slot
    _worktree_release_slot "$WORKTREE_PATH"

    # Remove the worktree
    echo "Removing worktree at: $WORKTREE_PATH"
    if [ "$FLAG" = "--force" ]; then
        git worktree remove --force "$WORKTREE_PATH" || return 1
    else
        git worktree remove "$WORKTREE_PATH" || return 1
    fi

    echo ""
    echo "Worktree removed: $WORKTREE_PATH"
    echo "Branch '$BRANCH_NAME' still exists. To delete it: git branch -d $BRANCH_NAME"
}

# Tab completion for zsh
_worktree_completions() {
    local cmd="${words[2]}"

    case "$CURRENT" in
        2)
            # Complete commands
            compadd add setup start stop restart info run console open logs connect cd list ls prune remove rm help
            ;;
        3)
            case "$cmd" in
                remove|rm|cd)
                    # Complete with worktree directory names from registry
                    if [ -f "$WORKTREE_PORT_REGISTRY" ]; then
                        local worktrees=()
                        while IFS=: read -r path slot; do
                            if [ -d "$path" ]; then
                                worktrees+=("${path##*/}")
                            fi
                        done < "$WORKTREE_PORT_REGISTRY"
                        [ ${#worktrees[@]} -gt 0 ] && compadd -a worktrees
                    fi
                    ;;
                add)
                    # Complete with git branch names
                    local branches=($(git branch --format='%(refname:short)' 2>/dev/null))
                    [ ${#branches[@]} -gt 0 ] && compadd -a branches
                    ;;
                connect)
                    # Complete with process names
                    compadd web vite worker
                    ;;
            esac
            ;;
    esac
}

compdef _worktree_completions worktree

# Aliases
alias wt='worktree'
alias wta='worktree add'
alias wts='worktree start'
alias wtp='worktree stop'
alias wtrs='worktree restart'
alias wtr='worktree run'
alias wtc='worktree console'
alias wto='worktree open'
alias wtlg='worktree logs'
alias wtcn='worktree connect'
alias wtcd='worktree cd'
alias wti='worktree info'
alias wtl='worktree list'
alias wtpr='worktree prune'
alias wtrm='worktree remove'