jesseduffield.lazygit/pkg/gui/presentation/branches.go

package presentation

import (
	"fmt"
	"regexp"
	"strings"
	"time"

	"github.com/jesseduffield/lazygit/pkg/commands/git_commands"
	"github.com/jesseduffield/lazygit/pkg/commands/models"
	"github.com/jesseduffield/lazygit/pkg/config"
	"github.com/jesseduffield/lazygit/pkg/gui/presentation/icons"
	"github.com/jesseduffield/lazygit/pkg/gui/style"
	"github.com/jesseduffield/lazygit/pkg/gui/types"
	"github.com/jesseduffield/lazygit/pkg/i18n"
	"github.com/jesseduffield/lazygit/pkg/theme"
	"github.com/jesseduffield/lazygit/pkg/utils"
	"github.com/samber/lo"
)

type colorMatcher struct {
	patterns map[string]*style.TextStyle
	isRegex  bool // NOTE: this value is needed only until the deprecated branchColors config is removed and only regex color patterns are used
}

var colorPatterns *colorMatcher

func GetBranchListDisplayStrings(
	branches []*models.Branch,
	getItemOperation func(item types.HasUrn) types.ItemOperation,
	prs map[string]*models.GithubPullRequest,
	fullDescription bool,
	diffName string,
	viewWidth int,
	tr *i18n.TranslationSet,
	userConfig *config.UserConfig,
	worktrees []*models.Worktree,
) [][]string {
	return lo.Map(branches, func(branch *models.Branch, _ int) []string {
		diffed := branch.Name == diffName
		return getBranchDisplayStrings(branch, getItemOperation(branch), fullDescription, diffed, viewWidth, tr, userConfig, worktrees, time.Now(), prs)
	})
}

// getBranchDisplayStrings returns the display string of branch
func getBranchDisplayStrings(
	b *models.Branch,
	itemOperation types.ItemOperation,
	fullDescription bool,
	diffed bool,
	viewWidth int,
	tr *i18n.TranslationSet,
	userConfig *config.UserConfig,
	worktrees []*models.Worktree,
	now time.Time,
	prs map[string]*models.GithubPullRequest,
) []string {
	checkedOutByWorkTree := git_commands.CheckedOutByOtherWorktree(b, worktrees)
	showCommitHash := fullDescription || userConfig.Gui.ShowBranchCommitHash
	branchStatus := BranchStatus(b, itemOperation, tr, now, userConfig)
	divergence := divergenceStr(b, itemOperation, tr, userConfig)

	// Recency is always three characters, plus one for the space
	availableWidth := viewWidth - 4
	if len(divergence) > 0 {
		availableWidth -= utils.StringWidth(divergence) + 1
	}
	if showCommitHash {
		availableWidth -= utils.COMMIT_HASH_SHORT_SIZE + 1
	}
	if len(prs) > 0 {
		// if we have PRs then we assume that at least one branch in the list has one
		availableWidth -= 2
	}
	paddingNeededForDivergence := availableWidth

	displayName := b.Name
	if b.DisplayName != "" {
		displayName = b.DisplayName
	}

	if len(branchStatus) > 0 {
		availableWidth -= utils.StringWidth(utils.Decolorise(branchStatus)) + 1
	}

	worktreeIcon := ""
	if checkedOutByWorkTree {
		if wt, ok := git_commands.WorktreeForBranch(b, worktrees); ok && wt.Name != b.Name {
			if icons.IsIconEnabled() {
				worktreeIcon = fmt.Sprintf("(%s %s)", icons.LINKED_WORKTREE_ICON, wt.Name)
			} else {
				worktreeIcon = fmt.Sprintf("(%s %s)", tr.LcWorktree, wt.Name)
			}

			// If the worktree name doesn't fit in the available width, omit it
			remaining := availableWidth - utils.StringWidth(worktreeIcon) - 1
			if remaining < utils.StringWidth(displayName) {
				if icons.IsIconEnabled() {
					worktreeIcon = icons.LINKED_WORKTREE_ICON
				} else {
					worktreeIcon = fmt.Sprintf("(%s)", tr.LcWorktree)
				}
			}
		} else {
			if icons.IsIconEnabled() {
				worktreeIcon = icons.LINKED_WORKTREE_ICON
			} else {
				worktreeIcon = fmt.Sprintf("(%s)", tr.LcWorktree)
			}
		}

		availableWidth -= utils.StringWidth(worktreeIcon) + 1
	}

	nameTextStyle := GetBranchTextStyle(b.Name)
	if diffed {
		nameTextStyle = theme.DiffTerminalColor
	}

	// Don't bother shortening branch names that are already 3 characters or less
	if utils.StringWidth(displayName) > max(availableWidth, 3) {
		// Never shorten the branch name to less then 3 characters
		len := max(availableWidth, 4)
		displayName = utils.TruncateWithEllipsis(displayName, len)
	}
	coloredName := nameTextStyle.Sprint(displayName)
	if checkedOutByWorkTree {
		coloredName = fmt.Sprintf("%s %s", coloredName, style.FgDefault.Sprint(worktreeIcon))
	}
	if len(branchStatus) > 0 {
		coloredName = fmt.Sprintf("%s %s", coloredName, branchStatus)
	}

	recencyColor := style.FgCyan
	if b.Recency == "  *" {
		recencyColor = style.FgGreen
	}

	res := make([]string, 0, 6)
	res = append(res, recencyColor.Sprint(b.Recency))

	var coloredPrIcon string
	pr, hasPr := prs[b.Name]
	if hasPr && ShouldShowPrForBranch(pr, b.Name, userConfig) {
		var prIcon string
		if icons.IsIconEnabled() {
			prIcon = icons.IconForRemoteUrl(pr.Url)
		} else {
			prIcon = "●"
		}
		coloredPrIcon = prColor(pr.State).Sprint(prIcon)
	}
	res = append(res, coloredPrIcon)

	if showCommitHash {
		res = append(res, utils.ShortHash(b.CommitHash))
	}

	if divergence != "" {
		if fullDescription {
			// don't right-align the divergence in half or full screen mode, since other fields
			// follow in that case and we don't know the width of our column
			paddingNeededForDivergence = 1
		} else {
			paddingNeededForDivergence -= utils.StringWidth(utils.Decolorise(coloredName)) - 1
		}
		if paddingNeededForDivergence > 0 {
			coloredName += strings.Repeat(" ", paddingNeededForDivergence)
			coloredName += style.FgCyan.Sprint(divergence)
		}
	}
	res = append(res, coloredName)

	if fullDescription {
		res = append(
			res,
			fmt.Sprintf("%s %s",
				style.FgYellow.Sprint(b.UpstreamRemote),
				style.FgYellow.Sprint(b.UpstreamBranch),
			),
			utils.TruncateWithEllipsis(b.Subject, 60),
		)
	}
	return res
}

// GetBranchTextStyle branch color
func GetBranchTextStyle(name string) style.TextStyle {
	if style, ok := colorPatterns.match(name); ok {
		return *style
	}

	return theme.DefaultTextColor
}

func (m *colorMatcher) match(name string) (*style.TextStyle, bool) {
	if m.isRegex {
		for pattern, style := range m.patterns {
			if matched, _ := regexp.MatchString(pattern, name); matched {
				return style, true
			}
		}
	} else {
		// old behavior using the deprecated branchColors behavior matching on branch type
		branchType := strings.Split(name, "/")[0]
		if value, ok := m.patterns[branchType]; ok {
			return value, true
		}
	}

	return nil, false
}

func BranchStatus(
	branch *models.Branch,
	itemOperation types.ItemOperation,
	tr *i18n.TranslationSet,
	now time.Time,
	userConfig *config.UserConfig,
) string {
	itemOperationStr := ItemOperationToString(itemOperation, tr)
	if itemOperationStr != "" {
		return style.FgCyan.Sprintf("%s %s", itemOperationStr, Loader(now, userConfig.Gui.Spinner))
	}

	result := ""
	if branch.IsTrackingRemote() {
		if branch.UpstreamGone {
			result = style.FgRed.Sprint(tr.UpstreamGone)
		} else if branch.MatchesUpstream() {
			result = style.FgGreen.Sprint("✓")
		} else if branch.RemoteBranchNotStoredLocally() {
			result = style.FgMagenta.Sprint("?")
		} else if branch.IsBehindForPull() && branch.IsAheadForPull() {
			result = style.FgYellow.Sprintf("↓%s↑%s", branch.BehindForPull, branch.AheadForPull)
		} else if branch.IsBehindForPull() {
			result = style.FgYellow.Sprintf("↓%s", branch.BehindForPull)
		} else if branch.IsAheadForPull() {
			result = style.FgYellow.Sprintf("↑%s", branch.AheadForPull)
		}
	}

	return result
}

func divergenceStr(
	branch *models.Branch,
	itemOperation types.ItemOperation,
	tr *i18n.TranslationSet,
	userConfig *config.UserConfig,
) string {
	result := ""
	if ItemOperationToString(itemOperation, tr) == "" && userConfig.Gui.ShowDivergenceFromBaseBranch != "none" {
		behind := branch.BehindBaseBranch.Load()
		if behind != 0 {
			if userConfig.Gui.ShowDivergenceFromBaseBranch == "arrowAndNumber" {
				result += fmt.Sprintf("↓%d", behind)
			} else {
				result += "↓"
			}
		}
	}

	return result
}

func SetCustomBranches(customBranchColors map[string]string, isRegex bool) {
	colorPatterns = &colorMatcher{
		patterns: utils.SetCustomColors(customBranchColors),
		isRegex:  isRegex,
	}
}

func prColor(state string) style.TextStyle {
	switch state {
	case "OPEN":
		return style.FgGreen
	case "CLOSED":
		return style.FgRed
	case "MERGED":
		return style.FgMagenta
	case "DRAFT":
		return style.FgBlackLighter
	default:
		return style.FgDefault
	}
}

func ShouldShowPrForBranch(pr *models.GithubPullRequest, branchName string, userConfig *config.UserConfig) bool {
	if !lo.Contains(userConfig.Git.MainBranches, branchName) {
		return true
	}

	// For main branches we only want to show the PR if it's open (or draft), on the assumption that a
	// closed PR for a main branch is always a mistake.
	return pr.State != "CLOSED" && pr.State != "MERGED"
}