package tree

import (
	"regexp"
	"sort"
)

type Node struct {
	Name     rune
	Parents  []*Node
	Children []*Node
}

// BuildDependencyTree parses the input and returns the root node.
func BuildDependencyTree(data []string) *Node {
	depMap := make(map[rune]*Node)
	parentRegex := regexp.MustCompile("Step ([A-Z]) must")
	childRegex := regexp.MustCompile("step ([A-Z]) can begin")

	for _, line := range data {
		parent := rune(parentRegex.FindSubmatch([]byte(line))[1][0])
		child := rune(childRegex.FindSubmatch([]byte(line))[1][0])

		// Create the nodes if they don't exist
		if depMap[parent] == nil {
			depMap[parent] = &Node{Name: parent}
		}
		if depMap[child] == nil {
			depMap[child] = &Node{Name: child}
		}

		// Get each node and add the dependency info.
		pNode := depMap[parent]
		cNode := depMap[child]

		pNode.Children = append(pNode.Children, cNode)
		cNode.Parents = append(cNode.Parents, pNode)
	}

	// sort the children in each node
	for _, node := range depMap {
		sortChildren(node)
	}

	return BuildTreeRoot(depMap)
}

func sortChildren(node *Node) {
	children := node.Children
	sort.Slice(children[:], func(i, j int) bool {
		return children[i].Name < children[j].Name
	})
}

func BuildTreeRoot(depMap map[rune]*Node) *Node {
	root := &Node{Name: 0}

	for _, node := range depMap {
		if len(node.Parents) == 0 {
			root.Children = append(root.Children, node)
			node.Parents = append(node.Parents, root)
		}
	}

	sortChildren(root)
	return root
}