From e7b3a691d2fb8354ad14a8f3307f56ab46a4828a Mon Sep 17 00:00:00 2001
From: Conflux <wanliqunn@163.com>
Date: Wed, 11 Sep 2024 15:05:24 +0800
Subject: [PATCH] feat: Add folder diff command for comparing with 0gStorage

---
 cmd/diff_dir.go      | 59 +++++++++++++++++++++++++++++++++++
 go.mod               |  2 +-
 transfer/dir/diff.go | 74 ++++++++++++++++++++++++++++++++++++++++++++
 3 files changed, 134 insertions(+), 1 deletion(-)
 create mode 100644 cmd/diff_dir.go

diff --git a/cmd/diff_dir.go b/cmd/diff_dir.go
new file mode 100644
index 0000000..1c3b2fb
--- /dev/null
+++ b/cmd/diff_dir.go
@@ -0,0 +1,59 @@
+package cmd
+
+import (
+	"context"
+
+	"github.com/0glabs/0g-storage-client/transfer"
+	"github.com/0glabs/0g-storage-client/transfer/dir"
+	"github.com/sirupsen/logrus"
+	"github.com/spf13/cobra"
+)
+
+var (
+	diffDirArgs downloadArgument
+
+	diffDirCmd = &cobra.Command{
+		Use:   "diff-dir",
+		Short: "Diff directory from ZeroGStorage network",
+		Run:   diffDir,
+	}
+)
+
+func init() {
+	bindDownloadFlags(diffDirCmd, &diffDirArgs)
+
+	rootCmd.AddCommand(diffDirCmd)
+}
+
+func diffDir(*cobra.Command, []string) {
+	ctx := context.Background()
+	var cancel context.CancelFunc
+	if diffDirArgs.timeout > 0 {
+		ctx, cancel = context.WithTimeout(ctx, diffDirArgs.timeout)
+		defer cancel()
+	}
+
+	localRoot, err := dir.BuildFileTree(diffDirArgs.file)
+	if err != nil {
+		logrus.WithError(err).Fatal("Failed to build local file tree")
+	}
+
+	downloader, closer, err := newDownloader(diffDirArgs)
+	if err != nil {
+		logrus.WithError(err).Fatal("Failed to initialize downloader")
+	}
+	defer closer()
+
+	zgRoot, err := transfer.BuildFileTree(ctx, downloader, diffDirArgs.root, diffDirArgs.proof)
+	if err != nil {
+		logrus.WithError(err).Fatal("Failed to build file tree from ZeroGStorage network")
+	}
+
+	diffRoot, err := dir.Diff(zgRoot, localRoot)
+	if err != nil {
+		logrus.WithError(err).Fatal("Failed to diff directory")
+	}
+
+	// Print the diff result
+	dir.PrettyPrint(diffRoot)
+}
diff --git a/go.mod b/go.mod
index 1a1d2d5..8aa482b 100644
--- a/go.mod
+++ b/go.mod
@@ -4,6 +4,7 @@ go 1.21.5
 
 require (
 	github.com/ethereum/go-ethereum v1.14.7
+	github.com/fatih/color v1.16.0
 	github.com/gin-contrib/cors v1.7.2
 	github.com/gin-gonic/gin v1.10.0
 	github.com/go-playground/validator/v10 v10.22.0
@@ -54,7 +55,6 @@ require (
 	github.com/decred/dcrd/dcrec/secp256k1/v4 v4.3.0 // indirect
 	github.com/ethereum/c-kzg-4844 v1.0.0 // indirect
 	github.com/ethereum/go-verkle v0.1.1-0.20240306133620-7d920df305f0 // indirect
-	github.com/fatih/color v1.16.0 // indirect
 	github.com/fjl/memsize v0.0.2 // indirect
 	github.com/fsnotify/fsnotify v1.6.0 // indirect
 	github.com/gabriel-vasile/mimetype v1.4.4 // indirect
diff --git a/transfer/dir/diff.go b/transfer/dir/diff.go
index cad5d51..5a3ba8b 100644
--- a/transfer/dir/diff.go
+++ b/transfer/dir/diff.go
@@ -1,6 +1,10 @@
 package dir
 
 import (
+	"fmt"
+	"strings"
+
+	"github.com/fatih/color"
 	"github.com/google/btree"
 	"github.com/pkg/errors"
 )
@@ -84,3 +88,73 @@ func diff(current, next *FsNode) *DiffNode {
 
 	return root
 }
+
+// PrettyPrint prints the DiffNode tree in a human-readable format with a tree skeleton structure.
+func PrettyPrint(root *DiffNode) {
+	prettyPrint(root, 0, false, nil)
+}
+
+func prettyPrint(node *DiffNode, depth int, isLast bool, prefixes []string) {
+	if depth == 0 { // Root directory case
+		fmt.Println(node.Node.Name)
+		printChildEntries(node, 1, nil)
+		return
+	}
+
+	// Create prefix and branch for current node
+	prefix := strings.Join(prefixes, "")
+	branch := "├─"
+	if isLast {
+		branch = "└─"
+	}
+
+	// Print node with optional status
+	coloredStatus := colorStringForDiffStatus(node.Status)
+	fmt.Printf("%s%s %s %s\n", prefix, branch, node.Node.Name, coloredStatus)
+
+	// Update prefixes for children
+	newPrefixes := append(prefixes, "│  ")
+	if isLast {
+		newPrefixes[len(newPrefixes)-1] = "    "
+	}
+
+	// Recursively print children if it's a directory
+	if node.Node.Type == FileTypeDirectory {
+		printChildEntries(node, depth+1, newPrefixes)
+	}
+}
+
+func printChildEntries(node *DiffNode, depth int, prefixes []string) {
+	if node.Entries == nil || node.Entries.Len() == 0 {
+		return
+	}
+
+	// Get the last node in a BTree efficiently by descending to the last element.
+	var last *DiffNode
+	node.Entries.Descend(func(item *DiffNode) bool {
+		last = item
+		return false // Stop the iteration as soon as we get the last node.
+	})
+
+	// Print all child entries of the current node
+	node.Entries.Ascend(func(item *DiffNode) bool {
+		// Check if this is the last entry in the directory
+		isLast := item == last
+		prettyPrint(item, depth, isLast, prefixes)
+		return true
+	})
+}
+
+// colorStringForDiffStatus returns a color string based on the diff status.
+func colorStringForDiffStatus(status DiffStatus) string {
+	switch status {
+	case DiffStatusAdded:
+		return color.GreenString("[+]")
+	case DiffStatusRemoved:
+		return color.RedString("[-]")
+	case DiffStatusModified:
+		return color.YellowString("[*]")
+	default:
+		return ""
+	}
+}