package trie_tree

import (
	"encoding/json"
)

type TrieNode struct {
	Ch    byte
	Child []*TrieNode
}

func (r TrieNode) String() string {
	bytes, _ := json.Marshal(r)
	return string(bytes)
}

func NewTrieTree() (trie *TrieNode) {
	return &TrieNode{}
}

func newNode(c byte) *TrieNode {
	return &TrieNode{
		Ch: c,
	}
}

func (r *TrieNode) Insert(s string) {
	curr := r
	for i := range s {
		if n := curr.findChild(s[i]); n == nil {
			nn := newNode(s[i])
			curr.Child = append(curr.Child, nn)
			curr = nn
		} else {
			curr = n
		}
	}
}

func (r *TrieNode) Search(s string) []string {
	var curr = r
	for i := range s {
		if n := curr.findChild(s[i]); n != nil {
			curr = n
		} else {
			return nil
		}
	}
	lists := walkNode(curr)
	for i := range lists {
		var rlt = lists[i]
		if len(rlt) > 0 {
			rlt = rlt[1:]
		}
		lists[i] = s + rlt
	}
	return lists
}

type list []string

func walkNode(node *TrieNode) (l list) {
	if node == nil {
		return
	}
	if node.Child == nil {
		l = append(l, string(node.Ch))
		return
	}
	for i := range node.Child {
		subs := walkNode(node.Child[i])
		for j := range subs {
			l = append(l, string(node.Ch)+subs[j])
		}
	}
	return
}

func (r *TrieNode) findChild(character byte) (n *TrieNode) {
	for i := range r.Child {
		c := r.Child[i]
		if c.Ch == character {
			n = c
			break
		}
	}
	return
}