package trie

import (
	"sort"
)

// Node represents a node in the trie
type Node struct {
	children map[rune]*Node
	isEnd    bool
}

// New creates a new empty trie root node
func New() *Node {
	return &Node{
		children: make(map[rune]*Node),
		isEnd:    false,
	}
}

// Add adds a word to the trie
func (n *Node) Add(word string) {
	current := n
	for _, char := range word {
		if _, exists := current.children[char]; !exists {
			current.children[char] = &Node{
				children: make(map[rune]*Node),
				isEnd:    false,
			}
		}
		current = current.children[char]
	}
	current.isEnd = true
}

// countWords recursively counts the number of words in the subtree
func (n *Node) countWords() int {
	count := 0
	if n.isEnd {
		count++
	}
	for _, child := range n.children {
		count += child.countWords()
	}
	return count
}

// Count returns the number of words that have the given prefix
func (n *Node) Count(prefix string) int {
	current := n
	for _, char := range prefix {
		if _, exists := current.children[char]; !exists {
			return 0
		}
		current = current.children[char]
	}
	return current.countWords()
}

// PrefixCount represents a prefix and its count
type PrefixCount struct {
	Prefix string
	Count  int
}

// collectPrefixes recursively collects all prefixes and their counts
func collectPrefixes(n *Node, currentPrefix string, results *[]PrefixCount) {
	count := n.countWords()
	if count > 0 {
		*results = append(*results, PrefixCount{
			Prefix: currentPrefix,
			Count:  count,
		})
	}
	// Sort children by rune to ensure consistent traversal order
	chars := make([]rune, 0, len(n.children))
	for char := range n.children {
		chars = append(chars, char)
	}
	sort.Slice(chars, func(i, j int) bool {
		return chars[i] < chars[j]
	})
	for _, char := range chars {
		collectPrefixes(n.children[char], currentPrefix+string(char), results)
	}
}

// TopPrefixes returns a sorted list of prefixes by their counts
func TopPrefixes(n *Node) []PrefixCount {
	var results []PrefixCount
	collectPrefixes(n, "", &results)

	sort.Slice(results, func(i, j int) bool {
		if results[i].Count == results[j].Count {
			return results[i].Prefix < results[j].Prefix
		}
		return results[i].Count > results[j].Count
	})

	return results
}