zutils/cmd/tree.go

package cmd

import (
	"fmt"
	"os"
	"path/filepath"
	"strings"

	ignore "github.com/sabhiram/go-gitignore"
	"github.com/zemenawi/zutils/pkg/colors"
	"github.com/zemenawi/zutils/pkg/formatter"
)

type TreeOptions struct {
	ShowHidden  bool
	ShowIgnored bool
	MaxDepth    int
}

type TreeItem struct {
	Name     string
	Path     string
	IsDir    bool
	Size     int64
	Ignored  bool
	Depth    int
	Children []*TreeItem
}

func TreeCommand(args []string) error {
	opts := TreeOptions{
		ShowHidden:  true,
		ShowIgnored: true,
		MaxDepth:    -1,
	}

	path := "."

	for _, arg := range args {
		if arg == "-h" || arg == "--help" {
			fmt.Printf("%sUsage:%s z tree [path] [options]\n\n", colors.Cyan, colors.Reset)
			fmt.Printf("%sOptions:%s\n", colors.Bold, colors.Reset)
			fmt.Printf("  %s-h, --help%s     Show this help\n\n", colors.Green, colors.Reset)
			fmt.Printf("%sDescription:%s\n", colors.Bold, colors.Reset)
			fmt.Printf("  Displays directory tree with:\n")
			fmt.Printf("    - Hidden files (files starting with .)\n")
			fmt.Printf("    - Gitignored files (marked with [ignored by git])\n")
			fmt.Printf("    - File sizes for regular files\n")
			fmt.Printf("\n%sExamples:%s\n", colors.Bold, colors.Reset)
			fmt.Printf("  z tree                     # Show current directory\n")
			fmt.Printf("  z tree /path/to/dir        # Show specific directory\n")
			return nil
		}
		if !strings.HasPrefix(arg, "-") {
			path = arg
		}
	}

	info, err := os.Stat(path)
	if err != nil {
		return fmt.Errorf("error accessing path: %w", err)
	}
	if !info.IsDir() {
		return fmt.Errorf("'%s' is not a directory", path)
	}

	ignoreMatcher, _ := ignore.CompileIgnoreFile(filepath.Join(path, ".gitignore"))

	tree, err := buildTree(path, opts, ignoreMatcher, 0)
	if err != nil {
		return fmt.Errorf("error building tree: %w", err)
	}

	PrintBoxHeader("DIRECTORY TREE", colors.Green)
	fmt.Printf("%s%s📁 %s%s\n\n", colors.Bold, colors.Cyan, path, colors.Reset)

	totalSize := printTreeRecursive(tree, "", true, true)

	fmt.Println()
	PrintSectionHeader("SUMMARY")
	fmt.Printf("%s%s📊 Total Size:%s %s\n", colors.Green, colors.Bold, colors.Reset, formatter.FormatSize(totalSize))

	return nil
}

func buildTree(rootPath string, opts TreeOptions, ignoreMatcher *ignore.GitIgnore, depth int) (*TreeItem, error) {
	if opts.MaxDepth >= 0 && depth > opts.MaxDepth {
		return nil, nil
	}

	info, err := os.Stat(rootPath)
	if err != nil {
		return nil, err
	}

	relPath, _ := filepath.Rel(filepath.Dir(rootPath), rootPath)
	ignored := false
	if ignoreMatcher != nil {
		ignored = ignoreMatcher.MatchesPath(relPath)
	}

	if ignored && !opts.ShowIgnored {
		return nil, nil
	}

	baseName := filepath.Base(rootPath)
	isHiddenDir := info.IsDir() && len(baseName) > 1 && strings.HasPrefix(baseName, ".")

	item := &TreeItem{
		Name:    baseName,
		Path:    rootPath,
		IsDir:   info.IsDir(),
		Size:    info.Size(),
		Ignored: ignored,
		Depth:   depth,
	}

	if !info.IsDir() || isHiddenDir {
		if isHiddenDir {
			item.Size = info.Size()
			item.Children = nil
		}
		return item, nil
	}

	entries, err := os.ReadDir(rootPath)
	if err != nil {
		return item, nil
	}

	var children []*TreeItem
	for _, entry := range entries {
		if !opts.ShowHidden && strings.HasPrefix(entry.Name(), ".") {
			continue
		}

		childPath := filepath.Join(rootPath, entry.Name())
		childItem, err := buildTree(childPath, opts, ignoreMatcher, depth+1)
		if err != nil || childItem == nil {
			continue
		}
		children = append(children, childItem)
	}

	item.Children = children
	return item, nil
}

func printTreeRecursive(item *TreeItem, prefix string, isLast bool, isRoot bool) int64 {
	var totalSize int64

	if !isRoot {
		branch := "└── "
		if !isLast {
			branch = "├── "
		}

		name := item.Name
		if item.IsDir {
			name = colors.Cyan + item.Name + colors.Reset
		} else {
			name = colors.White + item.Name + colors.Reset
		}

		if item.Ignored {
			name += colors.Gray + " [ignored by git]" + colors.Reset
		}

		if !item.IsDir || (item.IsDir && item.Children == nil) {
			name += colors.Gray + fmt.Sprintf(" (%s)", formatter.FormatSize(item.Size)) + colors.Reset
		}

		fmt.Printf("%s%s%s%s\n", prefix, branch, name, colors.Reset)
		totalSize += item.Size
	}

	if item.Children != nil {
		branch := ""
		if !isRoot {
			if isLast {
				branch = "    "
			} else {
				branch = "│   "
			}
		}

		for i, child := range item.Children {
			isLastChild := i == len(item.Children)-1
			totalSize += printTreeRecursive(child, prefix+branch, isLastChild, false)
		}
	}

	return totalSize
}