using System;
using System.Collections.Generic;

namespace WildRide.Utils
{
    public class PrefixTree<T>
    {
        private readonly Node _root;

        public PrefixTree()
        {
            _root = new Node();
        }

        public void Insert(string key, T data)
        {
            Node.Add(_root, new Substring(key), data);
        }

        public int Count()
        {
            int count = 0;
            Node.Count(_root, ref count);
            return count;
        }

        public IEnumerable<T> Find(string key)
        {
            Node node;
            if (Node.Find(_root, new Substring(key), out node))
            {
                return Node.GetData(node);
            }
            return null;
        }

        private struct Substring
        {
            public Substring(string str) :
                this(str, 0, str.Length)
            {
            }

            public Substring(string str, int start, int len)
            {
                String = str;
                Start = start;
                Length = len;
            }

            public readonly int Start;
            public readonly int Length;
            public readonly string String;

            public int StartsWith(Substring otherStr)
            {
                int pos;
                int len = Math.Min(this.Length, otherStr.Length);
                for (pos = 0; pos < len; ++pos)
                    if (otherStr.String[pos + otherStr.Start] != this.String[pos + this.Start])
                        break;
                return pos;
            }

            public override string ToString()
            {
                return String.Substring(Start, Length);
            }
        }

        private class Node
        {
            public Node() :
                this(default(Substring))
            {
            }

            private Node(Substring key, T data = default(T))
            {
                _key = key;
                _data = data;
                _children = new Dictionary<char, Node>();
            }

            private readonly Dictionary<char, Node> _children;
            private T _data;
            private Substring _key;

            internal static void Add(Node node, Substring key, T data)
            {
                if (key.Length == 0)
                {
                    node._data = data;
                    return;
                }

                Node child;
                if (node._children.TryGetValue(key.String[key.Start], out child))
                {
                    int keyPos = child._key.StartsWith(key);
                    if (keyPos == child._key.Length)
                    {
                        Add(child, new Substring(key.String, key.Start + keyPos, key.Length - keyPos), data);
                        return;
                    }
                    if (keyPos > 0)
                    {
                        Node newRoot = new Node(new Substring(child._key.String, child._key.Start, keyPos));
                        child._key = new Substring(child._key.String, child._key.Start + keyPos,
                            child._key.Length - keyPos);
                        newRoot._children.Add(child._key.String[child._key.Start], child);
                        node._children[newRoot._key.String[newRoot._key.Start]] = newRoot;
                        Add(newRoot, new Substring(key.String, key.Start + keyPos, key.Length - keyPos), data);
                    }
                }
                else
                {
                    Node n = new Node(key, data);
                    node._children.Add(key.String[key.Start], n);
                }
            }

            internal static bool Find(Node node, Substring key, out Node result)
            {
                if (key.Length == 0)
                {
                    result = node;
                    return true;
                }

                Node child;
                if (node._children.TryGetValue(key.String[key.Start], out child))
                {
                    int keyPos = child._key.StartsWith(key);
                    if (keyPos == child._key.Length)
                    {
                        return Find(child, new Substring(key.String, key.Start + keyPos, key.Length - keyPos), out result);
                    }
                    if (keyPos > 0 && (key.Length - keyPos) == 0)
                    {
                        result = child;
                        return true;
                    }
                }

                result = default(Node);
                return false;
            }

            internal static IEnumerable<T> GetData(Node node)
            {
                if (node._data != null)
                    yield return node._data;
                foreach (var child in node._children)
                {
                    foreach (T data in GetData(child.Value))
                        yield return data;
                }
            }

            internal static void Count(Node node, ref int count)
            {
                count++;
                foreach (var child in node._children)
                {
                    Count(child.Value, ref count);
                }
            }

            public override string ToString()
            {
                return _key.ToString();
            }
        }
    }
}