package drawing

import (
	"bytes"
	"fmt"
	"io"
	"os"
	"sort"
	"strings"
)

const (
	lightVerticalRight = "\u251C"
	lightUpRight       = "\u2514"
	lightHorizontal    = "\u2500"
	lightVertical      = "\u2502"
)

type StringTree struct {
	Data     string
	Children []*StringTree
	Labels   map[string]string
}

func (t *StringTree) String() string {
	buf := &bytes.Buffer{}
	t.Fprint(buf)
	return buf.String()
}

func (t *StringTree) Print() {
	t.Fprint(os.Stdout)
}

func (t *StringTree) Fprint(w io.Writer) {
	fmt.Fprintln(w, t.Data, t.labels())
	t.printTree(w, -1, 0, 0, "", true)
}

func (t *StringTree) printTree(w io.Writer, pi, pn, depth int, indention string, sorted bool) {
	if sorted {
		sort.Slice(t.Children, func(i int, j int) bool {
			return t.Children[i].depth(0) < t.Children[j].depth(0)
		})
	}

	n := len(t.Children)
	for i, tree := range t.Children {
		var marker string
		if i == len(t.Children)-1 {
			marker = lightUpRight + lightHorizontal
		} else {
			marker = lightVerticalRight + lightHorizontal
		}
		debugStr := ""
		debug := false
		if debug {
			maxDepth := tree.depth(0)
			debugStr = fmt.Sprintf(" [depth: %d, indent: '%s', i: %d, n: %d, pi: %d, pn: %d, maxDepth: %d]", depth, indention, i, n, pi, pn, maxDepth)
		}

		fmt.Fprintf(w, "%s%s %s%s %v\n", indention, marker, tree.Data, debugStr, tree.labels())
		tree.printTree(w, i, n, depth+1, indention+indent(i, n), sorted)
	}
}

func (t *StringTree) depth(depth int) int {
	max := depth
	for _, tree := range t.Children {
		treeD := tree.depth(depth + 1)
		if treeD > max {
			max = treeD
		}
	}
	return max
}

func indent(i, n int) string {
	s := ""
	if i != n-1 { // parent is not last child
		s += lightVertical + "  "
	} else {
		s += "   "
	}
	return s
}

func (t *StringTree) labels() string {
	labelStrs := []string{}
	for k, v := range t.Labels {
		if v == "" {
			labelStrs = append(labelStrs, fmt.Sprintf("%v", k))
		} else {
			labelStrs = append(labelStrs, fmt.Sprintf("%v=%v", k, v))
		}
	}
	if len(labelStrs) > 0 {
		return "[" + strings.Join(labelStrs, ", ") + "]"
	}
	return ""
}