git-worktree.zsh

# Git worktree shell functions
# These run in the current shell so `cd` works without spawning subshells.
#
# Assumes each project is cloned in the workspace directory, e.g., `~/workspace/<project>`.
# Worktrees are stored outside the project so editors don't pick them up, e.g., 
# `~/workspace/.worktrees/<project>/<branch>`.
#
# The main workflow is using `gct` to create and switch worktrees. The command for creating and
# switching is intentionally the same for convenience when using the shell history.
#
# `gct` supports multiple formats, e.g., GitHub URLs, pull request numbers, and branch names.
#
#   gct <branch>               Switch to a worktree from the current project
#   gct <project>/<branch>     Switch to a worktree from a specific project
#   gct <pr-number>            Switch to a worktee for a pull request from the current project
#   gct <project>/<pr-number>  Switch to a worktree for a pull request from a specific project
#   gct <github-pr-url>        Switch to a worktree for a pull request by GitHub URL
#   gct <github-issue-url>     Switch to a worktree for an issue by GitHub URL
# 
# `gcm` can be used to return to the main worktree and default branch.
#
# `gco` will switch branches in the current directory. If the target branch has a worktree
# elsewhere, it will remove that worktree first (unless it is dirty, in which case it will error).

WORKSPACE_ROOT="$HOME/workspace"
WORKTREE_ROOT="$WORKSPACE_ROOT/.worktrees"


# _gwt_repo_root [dir] - Resolve the root of the main repo (works from worktrees)
_gwt_repo_root() {
    local dir="${1:-.}"
    local git_common_dir
    git_common_dir=$(git -C "$dir" rev-parse --git-common-dir 2>/dev/null) || return 1
    if [[ "$git_common_dir" != /* ]]; then
        git_common_dir="$(cd -- "$dir" && cd -- "$git_common_dir" && pwd -P)"
    else
        git_common_dir="$(cd -- "$git_common_dir" && pwd -P)"
    fi
    dirname "$git_common_dir"
}

# _gwt_default_branch [dir] - Detect the default branch (main, master, etc.)
_gwt_default_branch() {
    local dir="${1:-.}"
    local branch
    branch=$(git -C "$dir" symbolic-ref refs/remotes/origin/HEAD 2>/dev/null | sed 's|refs/remotes/origin/||')
    echo "${branch:-main}"
}

# _gwt_find_worktree repo_root branch - Find existing (non-prunable) worktree path
# Uses --porcelain output to safely handle paths with spaces and special chars.
_gwt_find_worktree() {
    local repo_root="$1" branch="$2"
    { git -C "$repo_root" worktree list --porcelain; echo; } | awk -v target="refs/heads/$branch" '
        /^worktree / { path = substr($0, 10); prunable = 0; matched = 0 }
        /^branch /   { if (substr($0, 8) == target) matched = 1 }
        /^prunable/  { prunable = 1 }
        /^$/         { if (matched && !prunable) { print path; exit } }
    '
}

# _gwt_switch_branch dir branch - Switch to branch in an existing worktree dir
_gwt_switch_branch() {
    local dir="$1" branch="$2"
    if git -C "$dir" show-ref --verify --quiet "refs/heads/$branch"; then
        git -C "$dir" switch "$branch"
    elif git -C "$dir" show-ref --verify --quiet "refs/remotes/origin/$branch"; then
        git -C "$dir" switch --track "origin/$branch"
    else
        git -C "$dir" switch -c "$branch"
    fi
}

# _gwt_create_worktree repo_root wt_dir branch - Add a new worktree
_gwt_create_worktree() {
    local repo_root="$1" wt_dir="$2" branch="$3"
    if git -C "$repo_root" show-ref --verify --quiet "refs/heads/$branch"; then
        git -C "$repo_root" worktree add "$wt_dir" "$branch"
    elif git -C "$repo_root" show-ref --verify --quiet "refs/remotes/origin/$branch"; then
        git -C "$repo_root" worktree add --track -b "$branch" "$wt_dir" "origin/$branch"
    else
        local remote_ref
        remote_ref=$(git -C "$repo_root" for-each-ref --format='%(refname:short)' "refs/remotes/*/$branch" | head -1)
        if [[ -n "$remote_ref" ]]; then
            git -C "$repo_root" worktree add --track -b "$branch" "$wt_dir" "$remote_ref"
        else
            git -C "$repo_root" worktree add "$wt_dir" -b "$branch"
        fi
    fi
}

# _gwt_switch branch wt_dir repo_root - Find or create worktree, then cd into it
_gwt_switch() {
    local branch="$1" wt_dir="$2" repo_root="$3"

    # Check for existing worktree registered with git
    local existing
    existing=$(_gwt_find_worktree "$repo_root" "$branch")
    if [[ -n "$existing" ]]; then
        if [[ -d "$existing" ]]; then
            builtin cd "$existing"
            return 0
        else
            git -C "$repo_root" worktree remove --force "$existing"
        fi
    fi

    # Check if a PR-prefixed worktree already exists for this branch (e.g., 17969-my-branch)
    local wt_parent
    wt_parent="$(dirname "$wt_dir")"
    local normalized="${branch//\//-}"
    if [[ -d "$wt_parent" ]]; then
        local match_dir
        for match_dir in "$wt_parent"/*-"$normalized"(N); do
            if [[ -d "$match_dir" ]]; then
                builtin cd "$match_dir"
                return 0
            fi
        done
    fi

    # Check if worktree directory already exists (orphaned)
    if [[ -d "$wt_dir" ]]; then
        echo "Switching to existing worktree at: $wt_dir"
        _gwt_switch_branch "$wt_dir" "$branch"
        builtin cd "$wt_dir"
        return 0
    fi

    # Fetch latest refs before creating a new worktree
    git -C "$repo_root" fetch --quiet 2>/dev/null

    mkdir -p "$(dirname "$wt_dir")"
    _gwt_create_worktree "$repo_root" "$wt_dir" "$branch"
    builtin cd "$wt_dir"
}

# _gwt_setup_pr_tracking worktree_path local_branch remote_branch is_fork fork_owner
# Top-level helper to configure remote tracking for PR branches.
_gwt_setup_pr_tracking() {
    local worktree_path="$1" local_branch="$2" remote_branch="$3" is_fork="$4" fork_owner="$5"
    if [[ "$is_fork" = "true" ]]; then
        if ! git -C "$worktree_path" remote get-url "$fork_owner" &>/dev/null; then
            local origin_url fork_url
            origin_url=$(git -C "$worktree_path" remote get-url origin)
            fork_url=$(echo "$origin_url" | sed -E "s|([:/])[^/]+/([^/]+)$|\1$fork_owner/\2|")
            git -C "$worktree_path" remote add "$fork_owner" "$fork_url"
        fi
        git -C "$worktree_path" fetch "$fork_owner" "$remote_branch" 2>/dev/null
        git -C "$worktree_path" branch --set-upstream-to="$fork_owner/$remote_branch" "$local_branch" 2>/dev/null
    else
        git -C "$worktree_path" branch --set-upstream-to="origin/$remote_branch" "$local_branch" 2>/dev/null
    fi
}

# _gh_url_parse: detect GitHub PR/issue URLs and resolve to project + branch
# Sets (typeset -g): _GH_URL_TYPE (pr|issue|""), _GH_URL_PROJECT,
#                     _GH_URL_BRANCH, _GH_URL_PR_NUM
_gh_url_parse() {
    local url="$1"
    typeset -g _GH_URL_TYPE="" _GH_URL_PROJECT="" _GH_URL_BRANCH="" _GH_URL_PR_NUM=""

    # Match https://github.com/owner/repo/pull/123 or /issues/456
    if [[ "$url" =~ ^https?://github\.com/([^/]+)/([^/]+)/(pull|issues)/([0-9]+) ]]; then
        local owner="${match[1]}"
        local repo="${match[2]}"
        local kind="${match[3]}"
        local num="${match[4]}"
        # Strip .git suffix if present
        repo="${repo%.git}"
        _GH_URL_PROJECT="$repo"

        if [[ "$kind" == "pull" ]]; then
            _GH_URL_TYPE="pr"
            _GH_URL_PR_NUM="$num"
            # Get branch name from PR
            _GH_URL_BRANCH=$(gh pr view "$num" --repo "$owner/$repo" --json headRefName -q .headRefName 2>/dev/null)
            if [[ -z "$_GH_URL_BRANCH" ]]; then
                echo "error: could not get branch for pull request #$num in $owner/$repo" >&2
                return 1
            fi
        elif [[ "$kind" == "issues" ]]; then
            _GH_URL_TYPE="issue"
            # Get issue title and slugify it
            local title
            title=$(gh issue view "$num" --repo "$owner/$repo" --json title -q .title 2>/dev/null)
            if [[ -z "$title" ]]; then
                echo "error: could not get title for issue #$num in $owner/$repo" >&2
                return 1
            fi
            # Slugify: lowercase, replace non-alnum with hyphens, collapse, trim
            local slug
            slug=$(echo "$title" | tr '[:upper:]' '[:lower:]' | sed -E 's/[^a-z0-9]+/-/g; s/^-+//; s/-+$//' | cut -c1-60)
            _GH_URL_BRANCH="${num}-${slug}"
        fi
    fi
}

gct() {
    local name="${1:-}"

    if [[ -z "$name" ]]; then
        echo "Usage: gct <branch>" >&2
        echo "       gct <project>/<branch>" >&2
        echo "       gct <github-pr-url>" >&2
        echo "       gct <github-issue-url>" >&2
        return 1
    fi

    # Handle GitHub URLs
    if [[ "$name" =~ ^https?://github\.com/ ]]; then
        _gh_url_parse "$name" || return 1

        if [[ "$_GH_URL_TYPE" == "pr" ]]; then
            if [[ -d "$WORKSPACE_ROOT/$_GH_URL_PROJECT" ]]; then
                _gct_checkout_pr "$_GH_URL_PROJECT/$_GH_URL_PR_NUM"
            else
                _gct_checkout_pr "$_GH_URL_PR_NUM"
            fi
            return $?
        elif [[ "$_GH_URL_TYPE" == "issue" ]]; then
            # Rewrite name to project/branch and fall through
            if [[ -d "$WORKSPACE_ROOT/$_GH_URL_PROJECT" ]]; then
                name="$_GH_URL_PROJECT/$_GH_URL_BRANCH"
            else
                name="$_GH_URL_BRANCH"
            fi
        fi
    fi

    # Strip leading # (e.g., "#17969" -> "17969")
    name="${name###}"

    # Handle bare PR numbers (digits only)
    if [[ "$name" =~ ^[0-9]+$ ]]; then
        _gct_checkout_pr "$name"
        return $?
    fi

    # Handle project/number
    if [[ "$name" == *"/"* ]] && [[ "${name#*/}" =~ ^[0-9]+$ ]]; then
        _gct_checkout_pr "$name"
        return $?
    fi

    # Check if name contains a slash and we're not in a git repo (or it looks like project/branch)
    if [[ "$name" == *"/"* ]] && { ! git rev-parse --git-dir &>/dev/null || [[ -d "$WORKSPACE_ROOT/${name%%/*}" ]]; }; then
        local project="${name%%/*}"
        local branch="${name#*/}"
        local repo_root="$WORKSPACE_ROOT/$project"

        if [[ ! -d "$repo_root/.git" ]] && [[ ! -d "$repo_root" ]]; then
            echo "error: project '$project' not found at $repo_root" >&2
            return 1
        fi

        local normalized_branch="${branch//\//-}"
        local wt_dir="$WORKTREE_ROOT/$project/$normalized_branch"

        _gwt_switch "$branch" "$wt_dir" "$repo_root"
        return $?
    fi

    # Use current git repo
    local repo_root
    repo_root=$(_gwt_repo_root) || { echo "error: not in a git repository" >&2; return 1; }
    local repo_name
    repo_name=$(basename "$repo_root")
    local normalized_name="${name//\//-}"
    local wt_dir="$(dirname "$repo_root")/.worktrees/$repo_name/$normalized_name"

    _gwt_switch "$name" "$wt_dir" "$repo_root"
}

gcm() {
    local git_dir git_common_dir
    git_dir=$(git rev-parse --git-dir 2>/dev/null)
    git_common_dir=$(git rev-parse --git-common-dir 2>/dev/null)
    local default_branch
    default_branch=$(_gwt_default_branch)

    if [[ "$git_dir" != "$git_common_dir" ]]; then
        # We're in a worktree — go back to the main workspace
        local main_worktree
        main_worktree=$(git worktree list --porcelain | awk '/^worktree / { print substr($0, 10); exit }')
        builtin cd "$main_worktree"
        git checkout "$default_branch"
    else
        git checkout "$default_branch"
    fi
}

# gco <branch> - Switch to branch in current directory
# If branch has a worktree elsewhere, remove it first (requires clean state)
gco() {
    local name="${1:-}"

    if [[ -z "$name" ]]; then
        echo "Usage: gco <branch> [git-checkout args]" >&2
        return 1
    fi

    shift
    local -a rest=("$@")

    # Pass-through mode for file checkout
    if (( ${#rest} > 0 )); then
        git checkout "$name" "${rest[@]}"
        return
    fi

    local repo_root
    repo_root=$(_gwt_repo_root) || { echo "error: not in a git repository" >&2; return 1; }

    # Check for existing worktree
    local worktree_path
    worktree_path=$(_gwt_find_worktree "$repo_root" "$name")

    if [[ -n "$worktree_path" ]] && [[ -d "$worktree_path" ]]; then
        if [[ "$PWD" = "$worktree_path" ]]; then
            echo "Already in worktree for $name"
            return 0
        fi
        # If we're in a worktree, switch to the target worktree
        if [[ "$(git rev-parse --git-dir 2>/dev/null)" != "$(git rev-parse --git-common-dir 2>/dev/null)" ]]; then
            echo "Switching to worktree at $worktree_path"
            builtin cd "$worktree_path"
            return 0
        fi
        # In the main worktree: remove the other worktree and check out here
        if [[ -n "$(git -C "$worktree_path" status --porcelain | grep -v '^??')" ]]; then
            echo "error: worktree at '$worktree_path' has uncommitted changes" >&2
            return 1
        fi
        if [[ -n "$(git status --porcelain | grep -v '^??')" ]]; then
            echo "error: current directory has uncommitted changes" >&2
            return 1
        fi
        echo "Removing worktree at $worktree_path"
        git worktree remove "$worktree_path"
    elif [[ -n "$worktree_path" ]]; then
        git worktree remove --force "$worktree_path"
    fi

    # Switch or create branch
    if git show-ref --verify --quiet "refs/heads/$name"; then
        git switch "$name"
    elif git show-ref --verify --quiet "refs/remotes/origin/$name"; then
        git switch --track "origin/$name"
    else
        local remote_ref
        remote_ref=$(git for-each-ref --format='%(refname:short)' "refs/remotes/*/$name" | head -1)
        if [[ -n "$remote_ref" ]]; then
            git switch --track "$remote_ref"
        else
            git switch -c "$name"
        fi
    fi
}

# _gct_checkout_pr: internal helper for PR worktree checkout
_gct_checkout_pr() {
    local arg="${1:-}"

    if [[ -z "$arg" ]]; then
        echo "Usage: gcpr <pr-number>" >&2
        echo "       gcpr <project>/<pr-number>" >&2
        return 1
    fi

    local project pr_num repo_root

    if [[ "$arg" == *"/"* ]]; then
        project="${arg%%/*}"
        pr_num="${arg#*/}"
        repo_root="$WORKSPACE_ROOT/$project"

        if [[ ! -d "$repo_root" ]]; then
            echo "error: project '$project' not found at $repo_root" >&2
            return 1
        fi
    else
        pr_num="$arg"
        repo_root=$(_gwt_repo_root) || { echo "error: not in a git repository" >&2; return 1; }
        project=$(basename "$repo_root")
    fi

    # Get PR info
    local pr_info branch is_fork fork_owner
    pr_info=$(gh pr view "$pr_num" --json headRefName,headRepositoryOwner,isCrossRepository --repo "$(git -C "$repo_root" remote get-url origin)")
    branch=$(echo "$pr_info" | jq -r .headRefName)
    is_fork=$(echo "$pr_info" | jq -r .isCrossRepository)
    fork_owner=$(echo "$pr_info" | jq -r .headRepositoryOwner.login)

    if [[ -z "$branch" ]]; then
        echo "error: could not get branch for pull request #$pr_num" >&2
        return 1
    fi

    # Use pr_num-branch as the local branch name to avoid ref conflicts
    # (e.g., a branch named "fix" would conflict with "fix/something")
    local local_branch="$pr_num-${branch//\//-}"
    local wt_dir="$WORKTREE_ROOT/$project/$local_branch"

    # Check for existing worktree — match either the local (pr_num-prefixed) or remote branch name
    local existing
    existing=$(_gwt_find_worktree "$repo_root" "$local_branch")
    if [[ -z "$existing" ]]; then
        existing=$(_gwt_find_worktree "$repo_root" "$branch")
    fi

    if [[ -n "$existing" ]]; then
        if [[ -d "$existing" ]]; then
            builtin cd "$existing"
            return 0
        else
            git -C "$repo_root" worktree remove --force "$existing"
        fi
    fi

    if [[ -d "$wt_dir" ]]; then
        builtin cd "$wt_dir"
        return 0
    fi

    mkdir -p "$(dirname "$wt_dir")"

    # Force-update local branch from pull request head; fall back to fetching the remote branch
    local fetch_err
    fetch_err=$(git -C "$repo_root" fetch origin "+pull/$pr_num/head:refs/heads/$local_branch" 2>&1) \
        || fetch_err=$(git -C "$repo_root" fetch origin "$branch" 2>&1)

    if ! git -C "$repo_root" show-ref --verify --quiet "refs/heads/$local_branch"; then
        echo "error: could not fetch branch '$branch' for pull request #$pr_num" >&2
        if [[ -n "$fetch_err" ]]; then
            echo "$fetch_err" >&2
        fi
        return 1
    fi

    git -C "$repo_root" worktree add "$wt_dir" "$local_branch"

    _gwt_setup_pr_tracking "$wt_dir" "$local_branch" "$branch" "$is_fork" "$fork_owner"
    builtin cd "$wt_dir"
}