github__gh-stack/cmd/sync.go

package cmd

import (
	"errors"
	"fmt"
	"strings"

	"github.com/cli/go-gh/v2/pkg/prompter"
	"github.com/github/gh-stack/internal/config"
	"github.com/github/gh-stack/internal/git"
	"github.com/github/gh-stack/internal/modify"
	"github.com/github/gh-stack/internal/stack"
	"github.com/spf13/cobra"
)

type syncOptions struct {
	remote string
	prune  bool
}

func SyncCmd(cfg *config.Config) *cobra.Command {
	opts := &syncOptions{}

	cmd := &cobra.Command{
		Use:   "sync",
		Short: "Sync the current stack with the remote",
		Long: `Fetch, rebase, push, and sync PR state for the current stack.

This command performs a safe, non-interactive synchronization:

  1. Fetches the latest changes from the remote
  2. Fast-forwards the trunk branch to match the remote
  3. Cascade-rebases stack branches onto their updated parents
  4. Pushes all branches atomically (using --force-with-lease --atomic)
  5. Syncs PR state from GitHub

If a rebase conflict is detected, all branches are restored to their
original state and you are advised to run "gh stack rebase" to resolve
conflicts interactively.

Use --prune to delete local branches for merged PRs. Stack metadata is
preserved so that rebase and display logic continue to work correctly.
If you are on a branch that would be pruned, your checkout is moved to
the first active branch in the stack, or the trunk if all are merged.`,
		RunE: func(cmd *cobra.Command, args []string) error {
			return runSync(cfg, opts)
		},
	}

	cmd.Flags().StringVar(&opts.remote, "remote", "", "Remote to fetch from and push to (defaults to auto-detected remote)")
	cmd.Flags().BoolVar(&opts.prune, "prune", false, "Delete local branches for merged PRs")

	return cmd
}

func runSync(cfg *config.Config, opts *syncOptions) error {
	result, err := loadStack(cfg, "")
	if err != nil {
		return ErrNotInStack
	}
	gitDir := result.GitDir

	if err := modify.CheckStateGuard(gitDir); err != nil {
		cfg.Errorf("%s", err)
		return ErrModifyRecovery
	}

	sf := result.StackFile
	s := result.Stack
	currentBranch := result.CurrentBranch

	// Resolve remote once for fetch and push
	remote, err := pickRemote(cfg, currentBranch, opts.remote)
	if err != nil {
		if !errors.Is(err, errInterrupt) {
			cfg.Errorf("%s", err)
		}
		return ErrSilent
	}

	// --- Step 1: Fetch ---
	// Enable git rerere so conflict resolutions are remembered.
	if err := ensureRerere(cfg); errors.Is(err, errInterrupt) {
		return ErrSilent
	}

	// Fetch trunk + active branches so tracking refs are current for
	// fast-forward detection (Step 2) and --force-with-lease (Step 4).
	fetchTargets := append([]string{s.Trunk.Branch}, activeBranchNames(s)...)
	_ = git.FetchBranches(remote, fetchTargets)
	cfg.Successf("Fetched latest changes from %s", remote)

	// --- Step 2: Fast-forward trunk ---
	trunk := s.Trunk.Branch
	trunkUpdated := false

	localSHA, remoteSHA := "", ""
	trunkRefs, trunkErr := git.RevParseMulti([]string{trunk, remote + "/" + trunk})
	if trunkErr == nil {
		localSHA, remoteSHA = trunkRefs[0], trunkRefs[1]
	}

	if trunkErr != nil {
		cfg.Warningf("Could not compare trunk %s with remote — skipping trunk update", trunk)
	} else if localSHA == remoteSHA {
		cfg.Successf("Trunk %s is already up to date", trunk)
	} else {
		isAncestor, err := git.IsAncestor(localSHA, remoteSHA)
		if err != nil {
			cfg.Warningf("Could not determine fast-forward status for %s: %v", trunk, err)
		} else if !isAncestor {
			cfg.Warningf("Trunk %s has diverged from %s — skipping trunk update", trunk, remote)
			cfg.Printf("  Local and remote %s have diverged. Resolve manually.", trunk)
		} else {
			// Fast-forward the trunk branch
			if currentBranch == trunk {
				if err := git.MergeFF(remote + "/" + trunk); err != nil {
					cfg.Warningf("Failed to fast-forward %s: %v", trunk, err)
				} else {
					cfg.Successf("Trunk %s fast-forwarded to %s", trunk, short(remoteSHA))
					trunkUpdated = true
				}
			} else {
				if err := updateBranchRef(trunk, remoteSHA); err != nil {
					cfg.Warningf("Failed to fast-forward %s: %v", trunk, err)
				} else {
					cfg.Successf("Trunk %s fast-forwarded to %s", trunk, short(remoteSHA))
					trunkUpdated = true
				}
			}
		}
	}

	// --- Step 2b: Fast-forward stack branches behind their remote tracking branch ---
	updatedBranches := fastForwardBranches(cfg, s, remote, currentBranch)
	branchesUpdated := len(updatedBranches) > 0

	// --- Step 3: Cascade rebase (if trunk or any branch moved) ---
	rebased := false
	if trunkUpdated || branchesUpdated {
		cfg.Printf("")
		cfg.Printf("Rebasing stack ...")

		// Sync PR state to detect merged PRs before rebasing.
		_ = syncStackPRs(cfg, s)

		// Save original refs so we can restore on conflict.
		// Merged branches that no longer exist locally have no ref to
		// resolve. They are always skipped during rebase but we must
		// also exclude them here to avoid a rev-parse error.
		branchNames := make([]string, 0, len(s.Branches))
		for _, b := range s.Branches {
			if b.IsMerged() && !git.BranchExists(b.Branch) {
				continue
			}
			branchNames = append(branchNames, b.Branch)
		}
		originalRefs, _ := git.RevParseMap(branchNames)

		// Backfill originalRefs for merged branches that were deleted locally.
		// The rebase loop uses originalRefs[br.Branch] as ontoOldBase; without
		// a valid entry the subsequent --onto rebase would receive an empty ref.
		for _, b := range s.Branches {
			if b.IsMerged() && !git.BranchExists(b.Branch) {
				if b.Head != "" {
					if originalRefs == nil {
						originalRefs = make(map[string]string)
					}
					originalRefs[b.Branch] = b.Head
				}
			}
		}

		needsOnto := false
		var ontoOldBase string

		conflicted := false
		for i, br := range s.Branches {
			var base string
			if i == 0 {
				base = trunk
			} else {
				base = s.Branches[i-1].Branch
			}

			// Skip branches whose PRs have already been merged.
			if br.IsMerged() {
				ontoOldBase = originalRefs[br.Branch]
				needsOnto = true
				cfg.Successf("Skipping %s (PR %s merged)", br.Branch, cfg.PRLink(br.PullRequest.Number, br.PullRequest.URL))
				continue
			}

			// Skip branches whose PRs are currently in a merge queue.
			if br.IsQueued() {
				ontoOldBase = originalRefs[br.Branch]
				needsOnto = true
				cfg.Successf("Skipping %s (PR %s queued)", br.Branch, cfg.PRLink(br.PullRequest.Number, br.PullRequest.URL))
				continue
			}

			if needsOnto {
				// Find --onto target: first non-merged/queued ancestor, or trunk.
				newBase := trunk
				for j := i - 1; j >= 0; j-- {
					b := s.Branches[j]
					if !b.IsSkipped() {
						newBase = b.Branch
						break
					}
				}

				// If ontoOldBase is stale (not an ancestor of the branch), the
				// branch was already rebased past it (e.g. by a previous run).
				// Fall back to merge-base(newBase, branch) to avoid replaying
				// already-applied commits.
				actualOldBase := ontoOldBase
				if isAnc, err := git.IsAncestor(ontoOldBase, br.Branch); err == nil && !isAnc {
					if mb, err := git.MergeBase(newBase, br.Branch); err == nil {
						actualOldBase = mb
					}
				}

				if err := git.RebaseOnto(newBase, actualOldBase, br.Branch); err != nil {
					// Conflict detected — abort and restore everything
					if git.IsRebaseInProgress() {
						_ = git.RebaseAbort()
					}
					restoreErrors := restoreBranches(originalRefs)
					_ = git.CheckoutBranch(currentBranch)

					cfg.Errorf("Conflict detected rebasing %s onto %s", br.Branch, newBase)
					reportRestoreStatus(cfg, restoreErrors)
					cfg.Printf("  Run `%s` to resolve conflicts interactively.",
						cfg.ColorCyan("gh stack rebase"))
					conflicted = true
					break
				}

				cfg.Successf("Rebased %s onto %s (adjusted for merged PR)", br.Branch, newBase)
				ontoOldBase = originalRefs[br.Branch]
			} else {
				var rebaseErr error
				if i > 0 {
					// Use --onto to replay only this branch's unique commits.
					rebaseErr = git.RebaseOnto(base, originalRefs[base], br.Branch)
				} else {
					if err := git.CheckoutBranch(br.Branch); err != nil {
						cfg.Errorf("Failed to checkout %s: %v", br.Branch, err)
						conflicted = true
						break
					}
					rebaseErr = git.Rebase(base)
				}

				if rebaseErr != nil {
					// Conflict detected — abort and restore everything
					if git.IsRebaseInProgress() {
						_ = git.RebaseAbort()
					}
					restoreErrors := restoreBranches(originalRefs)
					_ = git.CheckoutBranch(currentBranch)

					cfg.Errorf("Conflict detected rebasing %s onto %s", br.Branch, base)
					reportRestoreStatus(cfg, restoreErrors)
					cfg.Printf("  Run `%s` to resolve conflicts interactively.",
						cfg.ColorCyan("gh stack rebase"))
					conflicted = true
					break
				}

				cfg.Successf("Rebased %s onto %s", br.Branch, base)
			}
		}

		if !conflicted {
			rebased = true
			_ = git.CheckoutBranch(currentBranch)
		} else {
			// Persist refreshed PR state even on conflict, then bail out
			// before pushing or reporting success.
			stack.SaveNonBlocking(gitDir, sf)
			return ErrConflict
		}
	}

	// --- Step 4: Push ---
	cfg.Printf("")
	branches := activeBranchNames(s)

	if mergedCount := len(s.MergedBranches()); mergedCount > 0 {
		cfg.Printf("Skipping %d merged %s", mergedCount, plural(mergedCount, "branch", "branches"))
	}
	if queuedCount := len(s.QueuedBranches()); queuedCount > 0 {
		cfg.Printf("Skipping %d queued %s", queuedCount, plural(queuedCount, "branch", "branches"))
	}

	if len(branches) == 0 {
		cfg.Printf("No active branches to push (all merged)")
	} else {
		// After rebase, force-with-lease is required (history rewritten).
		// Without rebase, try a normal push first.
		force := rebased
		cfg.Printf("Pushing %d %s to %s...", len(branches), plural(len(branches), "branch", "branches"), remote)
		if err := git.Push(remote, branches, force, true); err != nil {
			if !force {
				cfg.Warningf("Push failed — branches may need force push after rebase")
				cfg.Printf("  Run `%s` to push with --force-with-lease.",
					cfg.ColorCyan("gh stack push"))
			} else {
				cfg.Warningf("Push failed: %v", err)
				cfg.Printf("  Run `%s` to retry.", cfg.ColorCyan("gh stack push"))
			}
		} else {
			cfg.Successf("Pushed %d branches", len(branches))
		}
	}

	// --- Step 5: Sync PR state ---
	cfg.Printf("")
	cfg.Printf("Syncing PRs ...")
	_ = syncStackPRs(cfg, s)

	// Report PR status for each branch
	for _, b := range s.Branches {
		if b.IsMerged() {
			continue
		}
		if b.IsQueued() {
			cfg.Successf("PR %s (%s) — Queued", cfg.PRLink(b.PullRequest.Number, b.PullRequest.URL), b.Branch)
			continue
		}
		if b.PullRequest != nil {
			cfg.Successf("PR %s (%s) — Open", cfg.PRLink(b.PullRequest.Number, b.PullRequest.URL), b.Branch)
		} else {
			cfg.Warningf("%s has no PR", b.Branch)
		}
	}
	merged := s.MergedBranches()
	if len(merged) > 0 {
		names := make([]string, len(merged))
		for i, m := range merged {
			if m.PullRequest != nil {
				names[i] = fmt.Sprintf("#%d", m.PullRequest.Number)
			} else {
				names[i] = m.Branch
			}
		}
		cfg.Printf("Merged: %s", strings.Join(names, ", "))
	}

	// --- Step 6: Prune merged branches (optional) ---
	doPrune := opts.prune
	if !doPrune {
		// --prune was not provided. If interactive, prompt.
		merged := s.MergedBranches()
		var prunableCount int
		for _, b := range merged {
			if git.BranchExists(b.Branch) {
				prunableCount++
			}
		}
		if prunableCount > 0 && cfg.IsInteractive() {
			prompt := fmt.Sprintf("Prune %d merged %s?",
				prunableCount, plural(prunableCount, "branch", "branches"))
			confirmed, err := confirmPrune(cfg, prompt, true)
			if err != nil {
				if isInterruptError(err) {
					printInterrupt(cfg)
					// Save state before exiting so PR sync isn't lost.
					_ = stack.Save(gitDir, sf)
					return ErrSilent
				}
				// On any other prompt error, skip pruning silently.
			} else {
				doPrune = confirmed
			}
		}
	}

	if doPrune {
		merged := s.MergedBranches()
		var prunable []string
		for _, b := range merged {
			if git.BranchExists(b.Branch) {
				prunable = append(prunable, b.Branch)
			}
		}

		if len(prunable) > 0 {
			// If the current branch is being pruned, switch away first.
			needsSwitch := false
			for _, name := range prunable {
				if name == currentBranch {
					needsSwitch = true
					break
				}
			}
			if needsSwitch {
				switchTarget := trunk
				for _, b := range s.Branches {
					if !b.IsSkipped() {
						switchTarget = b.Branch
						break
					}
				}
				if err := git.CheckoutBranch(switchTarget); err != nil {
					cfg.Warningf("Failed to switch from %s to %s: %v", currentBranch, switchTarget, err)
				} else {
					currentBranch = switchTarget
				}
			}

			cfg.Printf("")
			pruned := 0
			for _, name := range prunable {
				if err := git.DeleteBranch(name, true); err != nil {
					cfg.Warningf("Failed to delete %s: %v", name, err)
				} else {
					cfg.Successf("Pruned %s (merged)", name)
					pruned++
				}
			}
			if pruned > 0 {
				cfg.Successf("Pruned %d merged %s", pruned, plural(pruned, "branch", "branches"))
			}
		} else if opts.prune {
			cfg.Printf("")
			cfg.Printf("No merged branches to prune")
		}

		// Clean up remote-tracking refs for all merged branches, even if
		// the local branch was already deleted. This prevents
		// `git checkout <name>` from resurrecting the branch.
		for _, b := range merged {
			_ = git.DeleteTrackingRef(remote, b.Branch)
		}
	}

	// --- Step 7: Update base SHAs and save ---
	updateBaseSHAs(s)

	if err := stack.Save(gitDir, sf); err != nil {
		return handleSaveError(cfg, err)
	}

	cfg.Printf("")
	cfg.Successf("Stack synced")
	return nil
}

// updateBranchRef updates a branch ref to point to a new SHA (for branches not checked out).
func updateBranchRef(branch, sha string) error {
	return git.UpdateBranchRef(branch, sha)
}

// restoreBranches resets each branch to its original SHA, collecting any errors.
func restoreBranches(originalRefs map[string]string) []string {
	var errors []string
	for branch, sha := range originalRefs {
		if err := git.CheckoutBranch(branch); err != nil {
			errors = append(errors, fmt.Sprintf("checkout %s: %s", branch, err))
			continue
		}
		if err := git.ResetHard(sha); err != nil {
			errors = append(errors, fmt.Sprintf("reset %s: %s", branch, err))
		}
	}
	return errors
}

// reportRestoreStatus prints whether branch restoration succeeded or partially failed.
func reportRestoreStatus(cfg *config.Config, restoreErrors []string) {
	if len(restoreErrors) > 0 {
		cfg.Warningf("Some branches could not be fully restored:")
		for _, e := range restoreErrors {
			cfg.Printf("  %s", e)
		}
	} else {
		cfg.Printf("  All branches restored to their original state.")
	}
}

// short returns the first 7 characters of a SHA.
func short(sha string) string {
	if len(sha) > 7 {
		return sha[:7]
	}
	return sha
}

// confirmPrune asks the user to confirm pruning via ConfirmFn or a terminal prompt.
func confirmPrune(cfg *config.Config, prompt string, defaultValue bool) (bool, error) {
	if cfg.ConfirmFn != nil {
		return cfg.ConfirmFn(prompt, defaultValue)
	}
	p := prompter.New(cfg.In, cfg.Out, cfg.Err)
	return p.Confirm(prompt, defaultValue)
}