From 60872c91e6fa054d0c7fa4b09df8d9aa572702ff Mon Sep 17 00:00:00 2001
From: Joel Baranick <joel.baranick@ensighten.com>
Date: Thu, 1 Sep 2022 23:50:34 -0700
Subject: [PATCH] Update status to differentiate the main vs linked worktrees

---
 pkg/gui/controllers.go                        | 15 ++++++-
 pkg/gui/controllers/helpers/helpers.go        |  2 +
 pkg/gui/controllers/helpers/refresh_helper.go | 15 ++++++-
 .../controllers/helpers/worktree_helper.go    | 43 +++++++++++++++++++
 4 files changed, 71 insertions(+), 4 deletions(-)
 create mode 100644 pkg/gui/controllers/helpers/worktree_helper.go

diff --git a/pkg/gui/controllers.go b/pkg/gui/controllers.go
index 328620d48..c883c3cc5 100644
--- a/pkg/gui/controllers.go
+++ b/pkg/gui/controllers.go
@@ -45,7 +45,17 @@ func (gui *Gui) resetHelpersAndControllers() {
 	patchBuildingHelper := helpers.NewPatchBuildingHelper(helperCommon)
 	stagingHelper := helpers.NewStagingHelper(helperCommon)
 	mergeConflictsHelper := helpers.NewMergeConflictsHelper(helperCommon)
-	refreshHelper := helpers.NewRefreshHelper(helperCommon, refsHelper, rebaseHelper, patchBuildingHelper, stagingHelper, mergeConflictsHelper, gui.fileWatcher)
+	worktreeHelper := helpers.NewWorktreeHelper(helperCommon)
+	refreshHelper := helpers.NewRefreshHelper(
+		helperCommon,
+		refsHelper,
+		rebaseHelper,
+		patchBuildingHelper,
+		stagingHelper,
+		mergeConflictsHelper,
+		worktreeHelper,
+		gui.fileWatcher,
+	)
 	diffHelper := helpers.NewDiffHelper(helperCommon)
 	cherryPickHelper := helpers.NewCherryPickHelper(
 		helperCommon,
@@ -99,7 +109,8 @@ func (gui *Gui) resetHelpersAndControllers() {
 			modeHelper,
 			appStatusHelper,
 		),
-		Search: helpers.NewSearchHelper(helperCommon),
+		Search:   helpers.NewSearchHelper(helperCommon),
+		Worktree: worktreeHelper,
 	}
 
 	gui.CustomCommandsClient = custom_commands.NewClient(
diff --git a/pkg/gui/controllers/helpers/helpers.go b/pkg/gui/controllers/helpers/helpers.go
index 846638249..e87b57eb0 100644
--- a/pkg/gui/controllers/helpers/helpers.go
+++ b/pkg/gui/controllers/helpers/helpers.go
@@ -47,6 +47,7 @@ type Helpers struct {
 	AppStatus         *AppStatusHelper
 	WindowArrangement *WindowArrangementHelper
 	Search            *SearchHelper
+	Worktree          *WorktreeHelper
 }
 
 func NewStubHelpers() *Helpers {
@@ -80,5 +81,6 @@ func NewStubHelpers() *Helpers {
 		AppStatus:         &AppStatusHelper{},
 		WindowArrangement: &WindowArrangementHelper{},
 		Search:            &SearchHelper{},
+		Worktree:          &WorktreeHelper{},
 	}
 }
diff --git a/pkg/gui/controllers/helpers/refresh_helper.go b/pkg/gui/controllers/helpers/refresh_helper.go
index 8431ae4cf..9745cf58e 100644
--- a/pkg/gui/controllers/helpers/refresh_helper.go
+++ b/pkg/gui/controllers/helpers/refresh_helper.go
@@ -28,6 +28,7 @@ type RefreshHelper struct {
 	patchBuildingHelper  *PatchBuildingHelper
 	stagingHelper        *StagingHelper
 	mergeConflictsHelper *MergeConflictsHelper
+	worktreeHelper       *WorktreeHelper
 	fileWatcher          types.IFileWatcher
 }
 
@@ -38,6 +39,7 @@ func NewRefreshHelper(
 	patchBuildingHelper *PatchBuildingHelper,
 	stagingHelper *StagingHelper,
 	mergeConflictsHelper *MergeConflictsHelper,
+	worktreeHelper *WorktreeHelper,
 	fileWatcher types.IFileWatcher,
 ) *RefreshHelper {
 	return &RefreshHelper{
@@ -47,6 +49,7 @@ func NewRefreshHelper(
 		patchBuildingHelper:  patchBuildingHelper,
 		stagingHelper:        stagingHelper,
 		mergeConflictsHelper: mergeConflictsHelper,
+		worktreeHelper:       worktreeHelper,
 		fileWatcher:          fileWatcher,
 	}
 }
@@ -635,8 +638,16 @@ func (self *RefreshHelper) refreshStatus() {
 	}
 
 	name := presentation.GetBranchTextStyle(currentBranch.Name).Sprint(currentBranch.Name)
-	repoName := utils.GetCurrentRepoName()
-	status += fmt.Sprintf("%s → %s ", repoName, name)
+
+	var repoName string
+	worktreeName := self.worktreeHelper.GetCurrentWorktreeName()
+	if len(worktreeName) > 0 {
+		worktreeName = fmt.Sprintf("[%s]", worktreeName)
+		repoName = self.worktreeHelper.GetMainWorktreeName()
+	} else {
+		repoName = utils.GetCurrentRepoName()
+	}
+	status += fmt.Sprintf("%s%s → %s ", repoName, worktreeName, name)
 
 	self.c.SetViewContent(self.c.Views().Status, status)
 }
diff --git a/pkg/gui/controllers/helpers/worktree_helper.go b/pkg/gui/controllers/helpers/worktree_helper.go
new file mode 100644
index 000000000..9d22f0db1
--- /dev/null
+++ b/pkg/gui/controllers/helpers/worktree_helper.go
@@ -0,0 +1,43 @@
+package helpers
+
+import (
+	"path/filepath"
+)
+
+type IWorktreeHelper interface {
+	GetMainWorktreeName() string
+	GetCurrentWorktreeName() string
+}
+
+type WorktreeHelper struct {
+	c *HelperCommon
+}
+
+func NewWorktreeHelper(c *HelperCommon) *WorktreeHelper {
+	return &WorktreeHelper{
+		c: c,
+	}
+}
+
+func (self *WorktreeHelper) GetMainWorktreeName() string {
+	for _, worktree := range self.c.Model().Worktrees {
+		if worktree.Main {
+			return filepath.Base(worktree.Path)
+		}
+	}
+
+	return ""
+}
+
+func (self *WorktreeHelper) GetCurrentWorktreeName() string {
+	for _, worktree := range self.c.Model().Worktrees {
+		if worktree.Current {
+			if worktree.Main {
+				return ""
+			}
+			return worktree.Name
+		}
+	}
+
+	return ""
+}