From 4549290925b59896e3b110b156064628f224a77e Mon Sep 17 00:00:00 2001
From: Mauricio Trajano <mauriciot1993@gmail.com>
Date: Thu, 26 Dec 2024 23:34:08 -0500
Subject: [PATCH] WIP: collapse/uncollapse all files in tree

Seems like this feature is highly requested #4095 #3554

Did a quick wip pr. I got the core functionality working with files and
had some questions on what can be improved in order to polish it more.
I.e. What keymaps if any to add to these, should the functionality be to
toggle vs collapse and uncollapse separately, and if this should be
added anywhere else (right now it's just for the file tree)

Once I figure this out I can add the tests, documentations and anything
else that's missing. I'm also pretty new to the language so if there's a
way I can structure this better please let me know
---
 pkg/gui/controllers/files_controller.go | 34 +++++++++++++++++++++++++
 pkg/gui/filetree/collapsed_paths.go     | 10 +++++---
 pkg/gui/filetree/file_tree.go           | 14 ++++++++++
 3 files changed, 55 insertions(+), 3 deletions(-)

diff --git a/pkg/gui/controllers/files_controller.go b/pkg/gui/controllers/files_controller.go
index 1ea20eeb218..20cd7775a79 100644
--- a/pkg/gui/controllers/files_controller.go
+++ b/pkg/gui/controllers/files_controller.go
@@ -186,6 +186,12 @@ func (self *FilesController) GetKeybindings(opts types.KeybindingsOpts) []*types
 			Description: self.c.Tr.Fetch,
 			Tooltip:     self.c.Tr.FetchTooltip,
 		},
+		{
+			Key:     'z',
+			Handler: self.collapseAll,
+			// GetDisabledReason: self.require(self.singleItemSelected()),
+			Description: "Collapse All directories",
+		},
 	}
 }
 
@@ -478,6 +484,34 @@ func (self *FilesController) enter() error {
 	return self.EnterFile(types.OnFocusOpts{ClickedWindowName: "", ClickedViewLineIdx: -1})
 }
 
+func (self *FilesController) collapseAll() error {
+	nodes := self.context().GetAllItems()
+
+	dirPaths := lo.FilterMap(nodes, func(file *filetree.FileNode, index int) (string, bool) {
+		return file.Path, !file.IsFile()
+	})
+
+	self.context().FileTreeViewModel.CollapseAll(dirPaths)
+
+	self.c.PostRefreshUpdate(self.c.Contexts().Files)
+
+	return nil
+}
+
+func (self *FilesController) uncollapseAll() error {
+	files := self.context().GetAllFiles()
+
+	filePaths := lo.Map(files, func(file *models.File, index int) string {
+		return file.GetPath()
+	})
+
+	self.context().FileTreeViewModel.UncollapseAll(filePaths)
+
+	self.c.PostRefreshUpdate(self.c.Contexts().Files)
+
+	return nil
+}
+
 func (self *FilesController) EnterFile(opts types.OnFocusOpts) error {
 	node := self.context().GetSelected()
 	if node == nil {
diff --git a/pkg/gui/filetree/collapsed_paths.go b/pkg/gui/filetree/collapsed_paths.go
index 903999b37d0..039c431e1ea 100644
--- a/pkg/gui/filetree/collapsed_paths.go
+++ b/pkg/gui/filetree/collapsed_paths.go
@@ -29,10 +29,14 @@ func (self *CollapsedPaths) Collapse(path string) {
 	self.collapsedPaths.Add(path)
 }
 
+func (self *CollapsedPaths) Uncollapse(path string) {
+	self.collapsedPaths.Remove(path)
+}
+
 func (self *CollapsedPaths) ToggleCollapsed(path string) {
-	if self.collapsedPaths.Includes(path) {
-		self.collapsedPaths.Remove(path)
+	if self.IsCollapsed(path) {
+		self.Uncollapse(path)
 	} else {
-		self.collapsedPaths.Add(path)
+		self.Collapse(path)
 	}
 }
diff --git a/pkg/gui/filetree/file_tree.go b/pkg/gui/filetree/file_tree.go
index 12780e3ed3d..d12988daa47 100644
--- a/pkg/gui/filetree/file_tree.go
+++ b/pkg/gui/filetree/file_tree.go
@@ -44,6 +44,8 @@ type IFileTree interface {
 	GetAllFiles() []*models.File
 	GetFilter() FileTreeDisplayFilter
 	GetRoot() *FileNode
+	CollapseAll(paths []string)
+	UncollapseAll(paths []string)
 }
 
 type FileTree struct {
@@ -171,6 +173,18 @@ func (self *FileTree) ToggleCollapsed(path string) {
 	self.collapsedPaths.ToggleCollapsed(path)
 }
 
+func (self *FileTree) CollapseAll(paths []string) {
+	for _, path := range paths {
+		self.collapsedPaths.Collapse(path)
+	}
+}
+
+func (self *FileTree) UncollapseAll(paths []string) {
+	for _, path := range paths {
+		self.collapsedPaths.ExpandToPath(path)
+	}
+}
+
 func (self *FileTree) Tree() *FileNode {
 	return NewFileNode(self.tree)
 }