#include <stack>

#include <boost/make_shared.hpp>

#include <openssl/rand.h>

#include "SHAMap.h"

void SHAMap::getMissingNodes(std::vector<SHAMapNode>& nodeIDs, std::vector<uint256>& hashes, int max)
{
	boost::recursive_mutex::scoped_lock sl(mLock);

	assert(root->isValid());
	
	if(root->isFullBelow())
	{
#ifdef GMN_DEBUG
		std::cerr << "getMissingNodes: root is full below" << std::endl;
#endif
		clearSynching();
		return;
	}

	if(!root->isInner())
	{
		std::cerr << "synching empty tree" << std::endl;
		return;
	}

	std::stack<SHAMapTreeNode::pointer> stack;
	stack.push(root);

	while ((max > 0) && (!stack.empty()))
	{
		SHAMapTreeNode::pointer node = stack.top();
		stack.pop();

#ifdef GMN_DEBUG
		std::cerr << "gMN: popped " << node->getString() << std::endl;
#endif

		for (int i = 0; i < 16; ++i)
			if( !node->isEmptyBranch(i))
			{
#ifdef GMN_DEBUG
				std::cerr << "gMN: " << node->getString() << " has non-empty branch " << i << std::endl;
#endif
				SHAMapTreeNode::pointer desc = getNode(node->getChildNodeID(i), node->getChildHash(i), false);
				if(desc)
				{
					if (desc->isInner() && !desc->isFullBelow())
						stack.push(desc);
				}
				else if (max-- > 0)
				{
#ifdef GMN_DEBUG
					std::cerr << "gMN: need " << node->getChildNodeID(i).getString() << std::endl;
#endif
					nodeIDs.push_back(node->getChildNodeID(i));
				}
			}
	}
}

bool SHAMap::getNodeFat(const SHAMapNode& wanted, std::vector<SHAMapNode>& nodeIDs,
	std::list<std::vector<unsigned char> >& rawNodes)
{ // Gets a node and some of its children
	boost::recursive_mutex::scoped_lock sl(mLock);

	SHAMapTreeNode::pointer node = getNode(wanted);
	if (!node)
	{
		assert(false); // Remove for release, this can happen if we get a bogus request
		return false;
	}

	nodeIDs.push_back(*node);
	Serializer s;
	node->addRaw(s);
	rawNodes.push_back(s.peekData());

	if (node->isRoot() || node->isLeaf()) // don't get a fat root, can't get a fat leaf
		return true;

	for (int i = 0; i < 16; ++i)
		if (!node->isEmptyBranch(i))
		{
			SHAMapTreeNode::pointer nextNode = getNode(node->getChildNodeID(i), node->getChildHash(i), false);
			assert(nextNode);
			if(nextNode)
			{
				nodeIDs.push_back(*nextNode);
				Serializer s;
				nextNode->addRaw(s);
				rawNodes.push_back(s.peekData());
		 	}
		}

		return true;
}

bool SHAMap::addRootNode(const std::vector<unsigned char>& rootNode)
{
	boost::recursive_mutex::scoped_lock sl(mLock);

	// we already have a root node
	if (root->getNodeHash().isNonZero())
	{
#ifdef DEBUG
		std::cerr << "got root node, already have one" << std::endl;
#endif
		return true;
	}

	SHAMapTreeNode::pointer node = boost::make_shared<SHAMapTreeNode>(SHAMapNode(), rootNode, 0);
	if (!node) return false;

#ifdef DEBUG
	node->dump();
#endif

	returnNode(root, true);

	root = node;
	mTNByID[*root] = root;
	if (!root->getNodeHash())
	{
		root->setFullBelow();
		clearSynching();
	}

	return true;
}

bool SHAMap::addRootNode(const uint256& hash, const std::vector<unsigned char>& rootNode)
{
	boost::recursive_mutex::scoped_lock sl(mLock);

	// we already have a root node
	if (root->getNodeHash().isNonZero())
	{
#ifdef DEBUG
		std::cerr << "got root node, already have one" << std::endl;
#endif
		assert(root->getNodeHash() == hash);
		return true;
	}

	SHAMapTreeNode::pointer node = boost::make_shared<SHAMapTreeNode>(SHAMapNode(), rootNode, 0);
	if (!node) return false;
	if (node->getNodeHash() != hash) return false;

	returnNode(root, true);
	root = node;
	mTNByID[*root] = root;
	if (!root->getNodeHash())
	{
		root->setFullBelow();
		clearSynching();
	}

	return true;
}

bool SHAMap::addKnownNode(const SHAMapNode& node, const std::vector<unsigned char>& rawNode)
{ // return value: true=okay, false=error
	assert(!node.isRoot());
	if (!isSynching()) return false;

	boost::recursive_mutex::scoped_lock sl(mLock);

	if (checkCacheNode(node)) return true;

	std::stack<SHAMapTreeNode::pointer> stack = getStack(node.getNodeID(), true);
	if (stack.empty()) return false;

	SHAMapTreeNode::pointer iNode=stack.top();
	if (!iNode)
	{	// we should always have a root
		assert(false);
		return true;
	}

	if (iNode->isLeaf() || (iNode->getDepth() == node.getDepth()))
	{
#ifdef DEBUG
		std::cerr << "got inner node, already had it (late)" << std::endl;
#endif
		return true;
	}

	if (iNode->getDepth() != (node.getDepth() - 1))
	{ // Either this node is broken or we didn't request it (yet)
#ifdef DEBUG
		std::cerr << "unable to hook node " << node.getString() << std::endl;
		std::cerr << " stuck at " << iNode->getString() << std::endl;
		std::cerr << "got depth=" << node.getDepth() << ", walked to= " << iNode->getDepth() << std::endl;
#endif
		return false;
	}

	int branch = iNode->selectBranch(node.getNodeID());
	if (branch < 0)
	{
		assert(false);
		return false;
	}
	uint256 hash = iNode->getChildHash(branch);
	if (!hash) return false;

	SHAMapTreeNode::pointer newNode = boost::make_shared<SHAMapTreeNode>(node, rawNode, mSeq);
	if (hash != newNode->getNodeHash()) // these aren't the droids we're looking for
		return false;

	mTNByID[*newNode] = newNode;
	if (!newNode->isLeaf())
		return true; // only a leaf can fill a branch

	// did this new leaf cause its parents to fill up
	do
	{
		iNode = stack.top();
		stack.pop();
		assert(iNode->isInner());
		for(int i = 0; i < 16; ++i)
			if (!iNode->isEmptyBranch(i))
			{
				SHAMapTreeNode::pointer nextNode = getNode(iNode->getChildNodeID(i), iNode->getChildHash(i), false);
				if (!nextNode) return true;
				if (nextNode->isInner() && !nextNode->isFullBelow()) return true;
			}
		iNode->setFullBelow();
	} while (!stack.empty());
	if (root->isFullBelow()) clearSynching();
	return true;
}

bool SHAMap::deepCompare(SHAMap& other)
{ // Intended for debug/test only
	std::stack<SHAMapTreeNode::pointer> stack;
	boost::recursive_mutex::scoped_lock sl(mLock);

	stack.push(root);
	while (!stack.empty())
	{
		SHAMapTreeNode::pointer node = stack.top();
		stack.pop();

		SHAMapTreeNode::pointer otherNode;
		if(node->isRoot()) otherNode = other.root;
		else otherNode = other.getNode(*node, node->getNodeHash(), false);

		if (!otherNode)
		{
			std::cerr << "unable to fetch node" << std::endl;
			return false;
		}
		else if (otherNode->getNodeHash() != node->getNodeHash())
		{
			std::cerr << "node hash mismatch" << std::endl;
			return false;
		}

#ifdef DC_DEBUG
		std::cerr << "Comparing inner nodes " << node->getString() << std::endl;
#endif

		if (node->getNodeHash() != otherNode->getNodeHash())
			return false;
		if (node->isLeaf())
		{
			if (!otherNode->isLeaf()) return false;
			if (node->peekItem()->getTag() != otherNode->peekItem()->getTag()) return false;
			if (node->peekItem()->getData() != otherNode->peekItem()->getData()) return false;
		}
		else if (node->isInner())
		{
			if (!otherNode->isInner())
				return false;
			for(int i=0; i<16; i++)
			{
				if(node->isEmptyBranch(i))
				{
					if(!otherNode->isEmptyBranch(i)) return false;
				}
				else
				{
					SHAMapTreeNode::pointer next = getNode(node->getChildNodeID(i), node->getChildHash(i), false);
					if (!next)
					{
						std::cerr << "unable to fetch inner node" << std::endl;
						return false;
					}
					stack.push(next);
				}
			}
		}
	}
	return true;
}

#ifdef DEBUG
#define SMS_DEBUG
#endif

static SHAMapItem::pointer makeRandomAS()
{
		Serializer s;
		for(int d = 0; d < 3; ++d) s.add32(rand());
		return boost::make_shared<SHAMapItem>(s.getRIPEMD160().to256(), s.peekData());
}

static bool confuseMap(SHAMap &map, int count)
{
	// add a bunch of random states to a map, then remove them
	// map should be the same
	uint256 beforeHash = map.getHash();

	std::list<uint256> items;

	for (int i = 0; i < count; ++i)
	{
		SHAMapItem::pointer item = makeRandomAS();
		items.push_back(item->getTag());
		if (!map.addItem(*item, false))
		{
			std::cerr << "Unable to add item to map" << std::endl;
			return false;
		}
	}

	for (std::list<uint256>::iterator it = items.begin(); it != items.end(); ++it)
	{
		if (!map.delItem(*it))
		{
			std::cerr << "Unable to remove item from map" << std::endl;
			return false;
		}
	}

	if (beforeHash != map.getHash())
	{
		std::cerr << "Hashes do not match" << std::endl;
		return false;
	}

	return true;
}

bool SHAMap::syncTest()
{
	unsigned int seed;
	RAND_pseudo_bytes(reinterpret_cast<unsigned char *>(&seed), sizeof(seed));
	srand(seed);

	SHAMap source, destination;


	// add random data to the source map
	int items = 10000;
	for (int i = 0; i < items; ++i)
		source.addItem(*makeRandomAS(), false);

#ifdef DEBUG
	std::cerr << "Adding items, then removing them" << std::endl;
#endif
	if(!confuseMap(source, 500)) return false;

	source.setImmutable();

#ifdef DEBUG
	std::cerr << "SOURCE COMPLETE, SYNCHING" << std::endl;
#endif

	std::vector<SHAMapNode> nodeIDs, gotNodeIDs;
	std::list<std::vector<unsigned char> > gotNodes;
	std::vector<uint256> hashes;

	std::vector<SHAMapNode>::iterator nodeIDIterator;
	std::list<std::vector<unsigned char> >::iterator rawNodeIterator;

	int passes = 0;
	int nodes = 0;

	destination.setSynching();

	if (!source.getNodeFat(SHAMapNode(), nodeIDs, gotNodes))
	{
		std::cerr << "GetNodeFat(root) fails" << std::endl;
		assert(false);
		return false;
	}
	if (gotNodes.size() != 1)
	{
		std::cerr << "Didn't get root node " << gotNodes.size() << std::endl;
		assert(false);
		return false;
	}
	if (!destination.addRootNode(*gotNodes.begin()))
	{
		std::cerr << "AddRootNode fails" << std::endl;
		assert(false);
		return false;
	}
	nodeIDs.clear();
	gotNodes.clear();

#ifdef DEBUG
	std::cerr << "ROOT COMPLETE, INNER SYNCHING" << std::endl;
#endif
#ifdef SMS_DEBUG
	int bytes = 0;
#endif

	do
	{
		++passes;
		hashes.clear();

		// get the list of nodes we know we need
		destination.getMissingNodes(nodeIDs, hashes, 2048);
		if(nodeIDs.empty()) break;

#ifdef SMS_DEBUG
		std::cerr << nodeIDs.size() << " needed nodes" << std::endl;
#endif
		
		// get as many nodes as possible based on this information
		for (nodeIDIterator = nodeIDs.begin(); nodeIDIterator != nodeIDs.end(); ++nodeIDIterator)
			if (!source.getNodeFat(*nodeIDIterator, gotNodeIDs, gotNodes))
			{
				std::cerr << "GetNodeFat fails" << std::endl;
				assert(false);
				return false;
			}
		assert(gotNodeIDs.size() == gotNodes.size());
		nodeIDs.clear();
		hashes.clear();

		if (gotNodeIDs.empty())
		{
			std::cerr << "No nodes gotten" << std::endl;
			assert(false);
			return false;
		}

#ifdef SMS_DEBUG
		std::cerr << gotNodeIDs.size() << " found nodes" << std::endl;
#endif
		for (nodeIDIterator = gotNodeIDs.begin(), rawNodeIterator = gotNodes.begin();
				nodeIDIterator != gotNodeIDs.end(); ++nodeIDIterator, ++rawNodeIterator)
		{
			++nodes;
#ifdef SMS_DEBUG
			bytes += rawNodeIterator->size();
#endif
			if (!destination.addKnownNode(*nodeIDIterator, *rawNodeIterator))
			{
				std::cerr << "AddKnownNode fails" << std::endl;
				assert(false);
				return false;
			}
		}
		gotNodeIDs.clear();
		gotNodes.clear();


	} while (1);
	destination.clearSynching();

#ifdef SMS_DEBUG
	std::cerr << "SYNCHING COMPLETE " << items << " items, " << nodes << " nodes, " <<
		bytes / 1024 << " KB" << std::endl;
#endif

	if (!source.deepCompare(destination))
	{
		std::cerr << "DeepCompare fails" << std::endl;
		assert(false);
		return false;
	}

#ifdef SMS_DEBUG
	std::cerr << "SHAMapSync test passed: " << items << " items, " <<
		passes << " passes, " << nodes << " nodes" << std::endl;
#endif
	
	return true;
}