1package tree
2
3import (
4	"bytes"
5	"fmt"
6	"path/filepath"
7	"strings"
8
9	"github.com/fatih/color"
10)
11
12const (
13	symEmpty  = "    "
14	symBranch = "├── "
15	symLeaf   = "└── "
16	symVert   = "│   "
17)
18
19var (
20	colMount = color.New(color.FgCyan, color.Bold).SprintfFunc()
21	colDir   = color.New(color.FgBlue, color.Bold).SprintfFunc()
22	colTpl   = color.New(color.FgGreen, color.Bold).SprintfFunc()
23	// sep is intentionally NOT platform-agnostic. This is used for the CLI output
24	// and should always be a regular slash.
25	sep = "/"
26)
27
28// Root is the root of a tree
29type Root struct {
30	Name    string
31	Subtree *Tree
32	Prefix  string
33}
34
35// New creates a new tree
36func New(name string) *Root {
37	return &Root{
38		Name:    name,
39		Subtree: NewTree(),
40	}
41}
42
43// AddFile adds a new file to the tree
44func (r *Root) AddFile(path string, _ string) error {
45	return r.insert(path, false, "")
46}
47
48// AddMount adds a new mount point to the tree
49func (r *Root) AddMount(path, dest string) error {
50	return r.insert(path, false, dest)
51}
52
53// AddTemplate adds a template to the tree
54func (r *Root) AddTemplate(path string) error {
55	return r.insert(path, true, "")
56}
57
58func (r *Root) insert(path string, template bool, mountPath string) error {
59	t := r.Subtree
60	p := strings.Split(path, "/")
61	for i, e := range p {
62		n := &Node{
63			Name:    e,
64			Type:    "dir",
65			Subtree: NewTree(),
66		}
67		if i == len(p)-1 {
68			n.Type = "file"
69			n.Subtree = nil
70			n.Template = template
71			if mountPath != "" {
72				n.Mount = true
73				n.Path = mountPath
74			}
75		}
76		node, _ := t.Insert(n)
77		// do we need to extend an existing subtree?
78		if i < len(p)-1 && node.Subtree == nil {
79			node.Subtree = NewTree()
80			node.Type = "dir"
81		}
82		t = node.Subtree
83	}
84	return nil
85}
86
87// Format returns a pretty printed string of all nodes in and below
88// this node, e.g. ├── baz
89func (r *Root) Format(maxDepth int) string {
90	out := &bytes.Buffer{}
91
92	// any mount will be colored and include the on-disk path
93	_, _ = out.WriteString(colDir(r.Name))
94
95	// finish this folders output
96	_, _ = out.WriteString("\n")
97
98	// let our children format themselves
99	for i, node := range r.Subtree.Nodes {
100		last := i == len(r.Subtree.Nodes)-1
101		_, _ = out.WriteString(node.format("", last, maxDepth, 1))
102	}
103	return out.String()
104}
105
106// List returns a flat list of all files in this tree
107func (r *Root) List(maxDepth int) []string {
108	out := make([]string, 0, r.Len())
109	for _, t := range r.Subtree.Nodes {
110		out = append(out, t.list(r.Prefix, maxDepth, 0, true)...)
111	}
112	return out
113}
114
115// ListFolders returns a flat list of all folders in this tree
116func (r *Root) ListFolders(maxDepth int) []string {
117	out := make([]string, 0, r.Len())
118	for _, t := range r.Subtree.Nodes {
119		out = append(out, t.list(r.Prefix, maxDepth, 0, false)...)
120	}
121	return out
122}
123
124// String returns the name of this tree
125func (r *Root) String() string {
126	return r.Name
127}
128
129// FindFolder returns the subtree rooted at path
130func (r *Root) FindFolder(path string) (*Root, error) {
131	path = strings.TrimSuffix(path, "/")
132	t := r.Subtree
133	p := strings.Split(path, "/")
134	prefix := ""
135	for _, e := range p {
136		_, node := t.find(e)
137		if node == nil || node.Type == "file" || node.Subtree == nil {
138			return nil, fmt.Errorf("not found")
139		}
140		t = node.Subtree
141		prefix = filepath.Join(prefix, e)
142	}
143	return &Root{Name: r.Name, Subtree: t, Prefix: prefix}, nil
144}
145
146// SetName changes the name of this tree
147func (r *Root) SetName(n string) {
148	r.Name = n
149}
150
151// Len returns the number of entries in this folder and all subfolder including
152// this folder itself
153func (r *Root) Len() int {
154	var l int
155	for _, t := range r.Subtree.Nodes {
156		l += t.Len()
157	}
158	return l
159}
160