#pragma once
#include <vector>
#include <unordered_map>
#include <assert.h>
#include <functional>

#include "../../core/Defines.h"

// Searchtree based class that finds all fitting nodes for a certain value. 
// A function needs to be given that hashes a value into a set of hashes (all values should give a set with the same length!)
// For each hash in this set of hashes, a value of 0 means that anything can be put there. The other hashes have to match.
template<typename K, typename V>
class NodeReplacementFinder
{
private:
	struct SearchTreeNode
	{
	private:
		std::unordered_map<K, SearchTreeNode*> children;
		union WildCardOrValue
		{
			SearchTreeNode* wildcard;
			V value;
		};
		WildCardOrValue wildCardOrValue;
		bool isLeaf;

		V GetValue() const
		{
			assert(isLeaf);
			return wildCardOrValue.value;
		}

		void SetValue(V value)
		{
			assert(isLeaf);
			wildCardOrValue.value = value;
		}

		SearchTreeNode* GetWildCard() const
		{
			assert(!isLeaf);
			return wildCardOrValue.wildcard;
		}

		void SetWildCard(SearchTreeNode* wildcard)
		{
			assert(!isLeaf);
			wildCardOrValue.wildcard = wildcard;
		}


		SearchTreeNode* GetChild(unsigned32 key) const
		{
			// If the key is a wildcard, return the wildcardnode
			if (key == 0) return GetWildCard();

			// Otherwise, find the correct child
			auto it = children.find(key);
			if (it == children.end()) return NULL;
			else return (*it).second;
		}

		SearchTreeNode* AddChild(unsigned32 key, bool isLeaf)
		{
			SearchTreeNode* child = GetChild(key);
			if (child != NULL) return child;

			SearchTreeNode* newChild = new SearchTreeNode(isLeaf);
			if (key == 0) SetWildCard(newChild);
			else children.insert(std::make_pair(key, newChild));
			return newChild;
		}

		// Recursively call add, updating the index to make sure the correct key is used
		void Add(const std::vector<unsigned32>& keys, const V& value, size_t index)
		{
			if (isLeaf) SetValue(value);
			else
			{
				unsigned32 curKey = keys[index];
				SearchTreeNode* child = AddChild(curKey, index == keys.size() - 1);
				child->Add(keys, value, index + 1);
			}
		}

		void Find(const std::vector<unsigned32>& keys, size_t index, std::vector<V>& out) const
		{
			if (isLeaf) out.push_back(GetValue());
			else
			{
				unsigned32 curKey = keys[index];
				if (curKey == 0)
				{
					// if the current key is a wildcard, explore all possible paths:
					for (auto child : children)
						child.second->Find(keys, index + 1, out);
				}
				else
				{
					// If the current key isn't a wildcard, only explore the current key and the wildcard key
					auto curKeyChild = GetChild(curKey);
					if (curKeyChild != NULL) curKeyChild->Find(keys, index + 1, out);
				}
				SearchTreeNode* wildcard = GetWildCard();
				if (wildcard != NULL) wildcard->Find(keys, index + 1, out);
			}
		}

		// Returns true if the node can be removed safely
		bool Remove(const std::vector<unsigned32>& keys, size_t index)
		{
			if (isLeaf) return true;
			else
			{
				unsigned32 curKey = keys[index];
				auto child = GetChild(curKey);
				if (child != NULL)
				{
					bool canDelete = child->Remove(keys, index + 1);
					if (canDelete)
					{
						if (curKey == 0) SetWildCard(NULL);
						else children.erase(children.find(curKey));

						delete child;
					}
				}
			}
			return children.empty();
		}


	public:
		SearchTreeNode(bool isLeaf) : isLeaf(isLeaf)
		{
			if (!isLeaf)
			{
				children = std::unordered_map<unsigned32, SearchTreeNode*>();
				SetWildCard(NULL);
			}
			else
			{
				SetValue(V());
			}
		}
		~SearchTreeNode()
		{
			if (!isLeaf)
			{
				for (auto child : children)
					delete child.second;
			}
		}

		void Add(const std::vector<unsigned32>& keys, V value)
		{
			Add(keys, value, 0);
		}

		void Remove(const std::vector<unsigned32>& keys)
		{
			Remove(keys, 0);
		}

		std::vector<V> Find(std::vector<unsigned32> keys) const
		{
			std::vector<V> res;
			Find(keys, 0, res);
			return res;
		}
	};
	std::function<std::vector<K>(const V)> GetKeys;
	SearchTreeNode* root;
public:
	NodeReplacementFinder(std::function<std::vector<K>(const V)> keyGetter)
	{
		GetKeys = keyGetter;
		root = new SearchTreeNode(false);
	}
	~NodeReplacementFinder()
	{
		delete root;
	}

	void Add(const V value)
	{
		std::vector<K> keys = GetKeys(value);
		root->Add(keys, value);
	}

	void Remove(const V value)
	{
		std::vector<K> keys = GetKeys(value);
		root->Remove(keys);
	}

	std::vector<V> Find(const V value) const
	{
		std::vector<K> keys = GetKeys(value);
		return root->Find(keys);
	}
};