package ds

// ds-trie.go
// 前缀树 (Trie),支持任意字符串的插入/查找/前缀匹配/删除。
// 每个节点 26 个子指针(默认按小写字母 a-z 索引),isEnd 标记单词结尾。
// 如需支持更多字符集,可改用 map[rune]*TrieNode。

// TrieNode 前缀树节点。
type TrieNode struct {
	children [26]*TrieNode
	isEnd    bool
}

// Trie 前缀树。默认只支持小写字母 a-z。
type Trie struct {
	root *TrieNode
}

// NewTrie 创建空 Trie。
func NewTrie() *Trie {
	return &Trie{root: &TrieNode{}}
}

// Insert 插入单词。仅支持 a-z,其他字符会 panic(用 rune 索引前先校验)。
func (t *Trie) Insert(word string) {
	node := t.root
	for _, ch := range word {
		idx := ch - 'a'
		if idx < 0 || idx >= 26 {
			panic("trie: 只支持小写字母 a-z,收到非法字符")
		}
		if node.children[idx] == nil {
			node.children[idx] = &TrieNode{}
		}
		node = node.children[idx]
	}
	node.isEnd = true
}

// Search 查找完整单词是否存在(必须 isEnd)。
func (t *Trie) Search(word string) bool {
	node := t.find(word)
	return node != nil && node.isEnd
}

// StartsWith 是否存在以 prefix 为前缀的单词。
func (t *Trie) StartsWith(prefix string) bool {
	return t.find(prefix) != nil
}

// find 沿前缀走到对应节点,不存在返回 nil。
func (t *Trie) find(s string) *TrieNode {
	node := t.root
	for _, ch := range s {
		idx := ch - 'a'
		if idx < 0 || idx >= 26 {
			return nil
		}
		node = node.children[idx]
		if node == nil {
			return nil
		}
	}
	return node
}

// Delete 删除单词。仅当单词存在时删除,并回收不再被任何单词共享的节点。
// 返回是否删除成功(单词存在且被删除)。
func (t *Trie) Delete(word string) bool {
	// 先确认单词存在,避免与回收语义混淆。
	node := t.find(word)
	if node == nil || !node.isEnd {
		return false
	}
	t.deleteNode(t.root, word, 0)
	return true
}

// deleteNode 递归回收:删除 word 沿途不再被需要的节点。
// 调用前需保证 word 存在。返回值:当前节点是否可被父节点回收。
func (t *Trie) deleteNode(node *TrieNode, word string, depth int) bool {
	if node == nil {
		return false
	}
	if depth == len(word) {
		node.isEnd = false
		return isEmpty(node)
	}
	idx := int(word[depth] - 'a')
	if idx < 0 || idx >= 26 {
		return false
	}
	child := node.children[idx]
	if t.deleteNode(child, word, depth+1) {
		node.children[idx] = nil
	}
	// 当前节点可回收:无子节点 且 非单词结尾
	return isEmpty(node) && !node.isEnd
}

// isEmpty 节点是否无任何子节点。
func isEmpty(node *TrieNode) bool {
	for i := 0; i < 26; i++ {
		if node.children[i] != nil {
			return false
		}
	}
	return true
}