package main

import (
    "fmt"
    "sort"
    "sync"
)

type KeyWordKV map[int64]string
type CharBeginKV map[string][]*KeyWordTreeNode
type PairList []Pair

func (p PairList) Len() int           { return len(p) }
func (p PairList) Less(i,j int) bool { return p[i].V > p[j].V }
func (p PairList) Swap(i,j int)      { p[i],p[j] = p[j],p[i] }

type Pair struct {
    K int64
    V int64
}

type KeyWordTreeNode struct {
    // 1,2,3
    KeyWordIds map[int64]bool
    // 百，度， 一， 下
    Char string
    // 父节点
    ParentKeyWordTreeNode *KeyWordTreeNode
    // 子节点集合
    SubKeyWordTreeNodes map[string]*KeyWordTreeNode
}

func NewKeyWordTreeNode() *KeyWordTreeNode {
    return &KeyWordTreeNode{
        KeyWordIds:            make(map[int64]bool, 0),Char:                  "",ParentKeyWordTreeNode: nil,SubKeyWordTreeNodes:   make(map[string]*KeyWordTreeNode,}
}

func NewKeyWordTreeNodeWithParams(ch string,parent *KeyWordTreeNode) *KeyWordTreeNode {
    return &KeyWordTreeNode{
        KeyWordIds:            make(map[int64]bool,Char:                  ch,ParentKeyWordTreeNode: parent,}
}

type KeyWordServer struct {
    root          *KeyWordTreeNode
    kv            KeyWordKV
    char_begin_kv CharBeginKV
    rw            sync.RWMutex
}

func NewKeyWordServer() *KeyWordServer {
    return &KeyWordServer{
        root:          NewKeyWordTreeNode(),kv:            KeyWordKV{},char_begin_kv: CharBeginKV{},}
}

func (s *KeyWordServer) Put(id int64,keyword string) {
    s.rw.Lock()
    s.kv[id] = keyword
    key_word_tmp_pt := s.root
    for _,v := range keyword {
        ch := string(v)
        if key_word_tmp_pt.SubKeyWordTreeNodes[ch] == nil {
            // root：{百:{}}
            node := NewKeyWordTreeNodeWithParams(ch,key_word_tmp_pt)
            key_word_tmp_pt.SubKeyWordTreeNodes[ch] = node
            // 记住不同位置的"百"的起点
            s.char_begin_kv[ch] = append(s.char_begin_kv[ch],node)

        }
        // {百:{度:{}}}
        key_word_tree_node := key_word_tmp_pt.SubKeyWordTreeNodes[ch]
        key_word_tree_node.KeyWordIds[id] = true
        // 更新指针
        key_word_tmp_pt = key_word_tmp_pt.SubKeyWordTreeNodes[ch]
    }
    s.rw.Unlock()
}

func (s *KeyWordServer) Sugg(keyword string,limit int) []string {
    s.rw.RLock()
    key_word_tmp_pt := s.root
    is_end := true
    for _,v := range keyword {
        ch := string(v)
        if key_word_tmp_pt.SubKeyWordTreeNodes[ch] == nil {
            is_end = false
            break
        }
        // 更新指针
        key_word_tmp_pt = key_word_tmp_pt.SubKeyWordTreeNodes[ch]
    }
    if is_end {
        // 前缀
        ret := make([]string, 0)
        ids := key_word_tmp_pt.KeyWordIds
        for id,_ := range ids {
            ret = append(ret,s.kv[id])
            limit --
            if limit == 0 {
                break
            }
        }
        return ret
    }
    s.rw.RUnlock()
    return make([]string, 0)
}

func (s *KeyWordServer) Search(keyword string,limit int) []string {
    s.rw.RLock()
    ids := make(map[int64]int64, 0)
    for pos,v := range keyword {
        ch := string(v)
        begins := s.char_begin_kv[ch]
        for _,begin := range begins {
            key_word_tmp_pt := begin
            next_pos := pos + 1
            for len(key_word_tmp_pt.SubKeyWordTreeNodes) > 0 && next_pos < len(keyword) {
                // 最大匹配
                next_ch := string(keyword[next_pos])
                if key_word_tmp_pt.SubKeyWordTreeNodes[next_ch] == nil {
                    break
                }
                key_word_tmp_pt = key_word_tmp_pt.SubKeyWordTreeNodes[next_ch]
                next_pos++
            }
            // 保存结果
            for id,_ := range key_word_tmp_pt.KeyWordIds {
                ids[id] = ids[id] + 1
            }
        }
    }
    // 排序输出果
    list := PairList{}
    for id,count := range ids {
        list = append(list,Pair{
            K: id,V: count,})
    }
    if !sort.IsSorted(list) {
        sort.Sort(list)
    }
    if len(list) > limit {
        list = list[:limit]
    }
    ret := make([]string, 0)
    for _,item := range list {
        ret = append(ret,s.kv[item.K])
    }
    s.rw.RUnlock()
    return ret
}

func (s *KeyWordServer) DebugPrint() {
    fmt.Println("s.kv =",s.kv)
    key_word_tmp_pt := s.root
    dfs(key_word_tmp_pt)
}
func dfs(root *KeyWordTreeNode) {
    if root == nil {
        return
    } else {
        fmt.Println("s.root =",root.Char)
        fmt.Println("s.KeyWordIds =",root.KeyWordIds)
        for _,v := range root.SubKeyWordTreeNodes {
            dfs(v)
        }
    }
}

func main() {
    s := NewKeyWordServer()
    s.Put(1,"ba")
    s.Put(2,"abd")
    s.Put(3,"acd")
    //s.DebugPrint()
    //fmt.Println(s.Sugg("b",2))
    //fmt.Println(s.Search("a",2))
    //fmt.Println(s.Search("b",2))
    //fmt.Println(s.Search("c",2))
    fmt.Println(s.Search("ba", 2))
}