package main

import "sync"

// 多叉树
type (
	//MultiTreeNode 多叉树节点
	MultiTreeNode struct {
		Key      rune
		End      bool
		Children map[rune]*MultiTreeNode
	}
	//MultiTree 多叉树
	MultiTree struct {
		Root *MultiTreeNode
		Lock sync.Mutex // 添加互斥锁
	}
	//结果
	Result struct {
		Data       string
		Start, End int
	}
)

func NewMultiTree() *MultiTree {
	return &MultiTree{
		Root: &MultiTreeNode{
			Children: make(map[rune]*MultiTreeNode),
		},
	}
}

func (mt *MultiTree) Add(segment string) {
	mt.Lock.Lock()
	defer mt.Lock.Unlock()
	path := []rune(segment)
	node := mt.Root
	pathLen := len(path)
	for index, t := range path {
		end := index == pathLen-1
		if v, ok := node.Children[t]; ok {
			node = v
			if end { //后加入的词(可能是老词的子集)，需要修正属性
				node.End = end
			}
		} else {
			leafNode := &MultiTreeNode{
				Key:      t,
				Children: make(map[rune]*MultiTreeNode),
			}
			if end { //是结束节点
				leafNode.End = true
			}
			node.Children[t] = leafNode
			node = leafNode
		}
	}
}

// Match 匹配 / 查找
func (mt *MultiTree) Match(segment string, maximum bool) []*Result {
	mt.Lock.Lock()
	defer mt.Lock.Unlock()
	var wordNode *MultiTreeNode
	node := mt.Root
	var ok bool
	path := []rune(segment)
	ret := make([]*Result, 0, 0)
	pos := 0
	//找第一层
	for pos < len(path) {
		if wordNode, ok = node.Children[path[pos]]; !ok {
			pos += 1
			continue
		}
		skipStep := 1
		for w := pos + 1; w < len(path); w++ {
			if wordNode, ok = wordNode.Children[path[w]]; !ok {
				break
			}
			if wordNode.End && !maximum {
				return []*Result{
					&Result{
						Data: string(path[pos : w+1]), Start: pos, End: w + 1},
				}
			} else if wordNode.End {
				ret = append(ret, &Result{
					Data: string(path[pos : w+1]), Start: pos, End: w + 1})
				skipStep = w - pos
			}
		}
		pos += skipStep

	}
	return ret
}