using System.Collections.Generic;

namespace GFGGame
{
    //支持String模糊查询的字典
    public class StrTrieDictionary<TValue>
    {
        private readonly Dictionary<string, TValue> dictionary;
        private readonly Trie trie;

        public StrTrieDictionary()
        {
            dictionary = new Dictionary<string, TValue>();
            trie = new Trie();
        }

        public void Add(string key, TValue value)
        {
            dictionary.Add(key, value);
            trie.Insert(key);
        }

        public List<string> SearchKeys(string pattern, int n)
        {
            return trie.Search(pattern, n);
        }

        public List<TValue> SearchByKeys(string pattern, int n)
        {
            var keys = trie.Search(pattern, n);
            var rs = new List<TValue>();
            foreach (var key in keys)
            {
                dictionary.TryGetValue(key, out var value);
                if (value != null) rs.Add(value);
            }

            return rs;
        }

        public TValue this[string key]
        {
            get => dictionary[key];
            set => Add(key, value);
        }

        public bool ContainsKey(string key)
        {
            return dictionary.ContainsKey(key);
        }

        public bool Remove(string key)
        {
            if (!dictionary.Remove(key)) return false;
            trie.Remove(key);
            return true;
        }
    }

    internal class TrieNode
    {
        public char Key { get; set; }
        public bool IsEndOfWord { get; set; }
        public Dictionary<char, TrieNode> Children { get; set; }

        public TrieNode(char key)
        {
            Key = key;
            IsEndOfWord = false;
            Children = new Dictionary<char, TrieNode>();
        }
    }

    internal class Trie
    {
        private readonly TrieNode root;

        public Trie()
        {
            root = new TrieNode('\0');
        }

        public void Insert(string word)
        {
            var node = root;
            foreach (var c in word)
            {
                if (!node.Children.ContainsKey(c))
                {
                    node.Children[c] = new TrieNode(c);
                }

                node = node.Children[c];
            }

            node.IsEndOfWord = true;
        }

        public List<string> Search(string pattern, int n)
        {
            var result = new List<string>();
            var node = root;
            foreach (var c in pattern)
            {
                if (!node.Children.ContainsKey(c))
                {
                    return result;
                }

                node = node.Children[c];
            }

            SearchHelper(node, pattern, result, n);
            return result;
        }

        private static void SearchHelper(TrieNode node, string prefix, ICollection<string> result, int n)
        {
            if (node.IsEndOfWord)
            {
                result.Add(prefix);
            }

            if (result.Count >= n)
            {
                return;
            }

            foreach (var child in node.Children.Values)
            {
                SearchHelper(child, prefix + child.Key, result, n);
                if (result.Count >= n)
                {
                    return;
                }
            }
        }

        public bool Remove(string word)
        {
            var node = root;
            var stack = new Stack<TrieNode>();
            foreach (var c in word)
            {
                if (!node.Children.ContainsKey(c))
                {
                    return false;
                }

                stack.Push(node);
                node = node.Children[c];
            }

            if (!node.IsEndOfWord)
            {
                return false;
            }

            node.IsEndOfWord = false;
            if (node.Children.Count != 0) return true;
            while (stack.Count > 0 && node.Children.Count == 0 && !node.IsEndOfWord)
            {
                var parent = stack.Pop();
                parent.Children.Remove(node.Key);
                node = parent;
            }

            return true;
        }
    }
}