/*
  Copyright (c) 2021 Alex Diener
  
  This software is provided 'as-is', without any express or implied
  warranty. In no event will the authors be held liable for any damages
  arising from the use of this software.
  
  Permission is granted to anyone to use this software for any purpose,
  including commercial applications, and to alter it and redistribute it
  freely, subject to the following restrictions:
  
  1. The origin of this software must not be misrepresented; you must not
     claim that you wrote the original software. If you use this software
     in a product, an acknowledgment in the product documentation would be
     appreciated but is not required.
  2. Altered source versions must be plainly marked as such, and must not be
     misrepresented as being the original software.
  3. This notice may not be removed or altered from any source distribution.
  
  Alex Diener alex@ludobloom.com
*/

#include "utilities/UndoTree.h"
#include <assert.h>
#include <limits.h>

#define stemobject_implementation UndoTree

stemobject_vtable_begin();
stemobject_vtable_entry(dispose);
stemobject_vtable_end();

UndoTree * UndoTree_create(void * target) {
	stemobject_create_implementation(init, target)
}

bool UndoTree_init(UndoTree * self, void * target) {
	call_super(init, self);
	self->target = target;
	self->nodeCount = 0;
	self->nodeAllocatedCount = 0;
	self->nodes = NULL;
	self->currentNodeID = 0;
	self->nextNodeID = 1;
	self->lastTraversedRootNodeID = 0;
	self->currentDepth = 0;
	return true;
}

void UndoTree_dispose(UndoTree * self) {
	for (unsigned int nodeIndex = 0; nodeIndex < self->nodeCount; nodeIndex++) {
		call_virtual(dispose, self->nodes[nodeIndex].stateDelta);
	}
	free(self->nodes);
	call_super_virtual(dispose, self);
}

UndoTree * UndoTree_copy(UndoTree * self) {
	UndoTree * copy = UndoTree_create(self->target);
	copy->nodeCount = self->nodeCount;
	copy->nodeAllocatedCount = copy->nodeCount;
	copy->nodes = malloc(copy->nodeCount * sizeof(*copy->nodes));
	for (unsigned int nodeIndex = 0; nodeIndex < copy->nodeCount; nodeIndex++) {
		copy->nodes[nodeIndex] = self->nodes[nodeIndex];
		copy->nodes[nodeIndex].stateDelta = call_virtual(copy, self->nodes[nodeIndex].stateDelta);
	}
	copy->currentNodeID = self->currentNodeID;
	copy->nextNodeID = self->nextNodeID;
	copy->lastTraversedRootNodeID = self->lastTraversedRootNodeID;
	copy->currentDepth = self->currentDepth;
	return copy;
}

static unsigned int getNodeIndexWithID(UndoTree * self, unsigned int nodeID) {
	for (unsigned int nodeIndex = 0; nodeIndex < self->nodeCount; nodeIndex++) {
		if (self->nodes[nodeIndex].uniqueID == nodeID) {
			return nodeIndex;
		}
	}
	return UINT_MAX;
}

static void freeChildrenOfNode(UndoTree * self, unsigned int nodeID) {
	for (unsigned int nodeIndex = 0; nodeIndex < self->nodeCount; nodeIndex++) {
		if (self->nodes[nodeIndex].parentID == nodeID) {
			freeChildrenOfNode(self, self->nodes[nodeIndex].uniqueID);
			call_virtual(dispose, self->nodes[nodeIndex].stateDelta);
			self->nodes[nodeIndex].stateDelta = NULL;
		}
	}
}

void UndoTree_appendNode(UndoTree * self, compat_type(UndoStateDelta *) stateDelta, bool pruneSiblings) {
	if (pruneSiblings) {
		if (self->currentNodeID == 0) {
			for (unsigned int nodeIndex = 0; nodeIndex < self->nodeCount; nodeIndex++) {
				call_virtual(dispose, self->nodes[nodeIndex].stateDelta);
			}
			self->nodeCount = 0;
		} else {
			freeChildrenOfNode(self, self->currentNodeID);
			unsigned int offset = 0;
			for (unsigned int nodeIndex = 0; nodeIndex < self->nodeCount; nodeIndex++) {
				self->nodes[nodeIndex] = self->nodes[nodeIndex + offset];
				if (self->nodes[nodeIndex].stateDelta == NULL) {
					offset++;
					self->nodeCount--;
					nodeIndex--;
				}
			}
		}
	}
	if (self->nodeCount >= self->nodeAllocatedCount) {
		self->nodeAllocatedCount = self->nodeAllocatedCount * 2 + (self->nodeAllocatedCount == 0) * 32;
		self->nodes = realloc(self->nodes, self->nodeAllocatedCount * sizeof(*self->nodes));
	}
	self->nodes[self->nodeCount].uniqueID = self->nextNodeID++;
	self->nodes[self->nodeCount].parentID = self->currentNodeID;
	self->nodes[self->nodeCount].lastTraversedChildID = 0;
	self->nodes[self->nodeCount].depth = self->currentDepth;
	self->nodes[self->nodeCount].stateDelta = stateDelta;
	self->currentNodeID = self->nodes[self->nodeCount].uniqueID;
	self->currentDepth++;
	self->nodeCount++;
}

bool UndoTree_canUndo(UndoTree * self) {
	return self->currentNodeID != 0;
}

bool UndoTree_canRedo(UndoTree * self) {
	if (self->currentNodeID == 0) {
		return self->lastTraversedRootNodeID != 0;
	}
	unsigned int nodeIndex = getNodeIndexWithID(self, self->currentNodeID);
	assert(nodeIndex != UINT_MAX);
	return self->nodes[nodeIndex].lastTraversedChildID != 0;
}

unsigned int UndoTree_getRedoTimelineCountFromCurrentNode(UndoTree * self) {
	unsigned int parentNodeID = 0;
	if (self->currentNodeID != 0) {
		unsigned int nodeIndex = getNodeIndexWithID(self, self->currentNodeID);
		assert(nodeIndex != UINT_MAX);
		parentNodeID = self->nodes[nodeIndex].uniqueID;
	}
	unsigned int childCount = 0;
	for (unsigned int nodeIndex = 0; nodeIndex < self->nodeCount; nodeIndex++) {
		if (self->nodes[nodeIndex].parentID == parentNodeID) {
			childCount++;
		}
	}
	return childCount;
}

UndoStateDelta * UndoTree_getStateDeltaForUndo(UndoTree * self) {
	if (self->currentNodeID == 0) {
		return NULL;
	}
	unsigned int nodeIndex = getNodeIndexWithID(self, self->currentNodeID);
	if (nodeIndex == UINT_MAX) {
		return NULL;
	}
	return self->nodes[nodeIndex].stateDelta;
}

static unsigned int getNodeIndexForRedo(UndoTree * self) {
	unsigned int nodeID = self->lastTraversedRootNodeID;
	if (self->currentNodeID != 0) {
		unsigned int nodeIndex = getNodeIndexWithID(self, self->currentNodeID);
		assert(nodeIndex != UINT_MAX);
		nodeID = self->nodes[nodeIndex].lastTraversedChildID;
	}
	return getNodeIndexWithID(self, nodeID);
}

UndoStateDelta * UndoTree_getStateDeltaForRedo(UndoTree * self) {
	unsigned int nodeIndex = getNodeIndexForRedo(self);
	if (nodeIndex == UINT_MAX) {
		return NULL;
	}
	return self->nodes[nodeIndex].stateDelta;
}

static unsigned int getNodeIndexForRedoAtTimelineIndex(UndoTree * self, unsigned int timelineIndex) {
	unsigned int parentNodeID = 0;
	if (self->currentNodeID != 0) {
		unsigned int nodeIndex = getNodeIndexWithID(self, self->currentNodeID);
		assert(nodeIndex != UINT_MAX);
		parentNodeID = self->nodes[nodeIndex].uniqueID;
	}
	unsigned int childCount = 0;
	for (unsigned int nodeIndex = 0; nodeIndex < self->nodeCount; nodeIndex++) {
		if (self->nodes[nodeIndex].parentID == parentNodeID) {
			if (childCount == timelineIndex) {
				return nodeIndex;
			}
			childCount++;
		}
	}
	return UINT_MAX;
}

UndoStateDelta * UndoTree_getStateDeltaForRedoWithTimeline(UndoTree * self, unsigned int timelineIndex) {
	unsigned int nodeIndex = getNodeIndexForRedoAtTimelineIndex(self, timelineIndex);
	if (nodeIndex == UINT_MAX) {
		return NULL;
	}
	return self->nodes[nodeIndex].stateDelta;
}

bool UndoTree_undo(UndoTree * self) {
	unsigned int nodeIndex = getNodeIndexWithID(self, self->currentNodeID);
	if (nodeIndex == UINT_MAX) {
		return false;
	}
	call_virtual(revert, self->nodes[nodeIndex].stateDelta, self->target);
	self->currentNodeID = self->nodes[nodeIndex].parentID;
	self->currentDepth--;
	if (self->currentNodeID == 0) {
		self->lastTraversedRootNodeID = self->nodes[nodeIndex].uniqueID;
	} else {
		unsigned int nodeIndex2 = getNodeIndexWithID(self, self->currentNodeID);
		assert(nodeIndex2 != UINT_MAX);
		self->nodes[nodeIndex2].lastTraversedChildID = self->nodes[nodeIndex].uniqueID;
	}
	return true;
}

bool UndoTree_redo(UndoTree * self) {
	unsigned int nodeIndex = getNodeIndexForRedo(self);
	if (nodeIndex == UINT_MAX) {
		return false;
	}
	call_virtual(apply, self->nodes[nodeIndex].stateDelta, self->target);
	self->currentNodeID = self->nodes[nodeIndex].uniqueID;
	self->currentDepth++;
	return true;
}

bool UndoTree_redoWithTimeline(UndoTree * self, unsigned int timelineIndex) {
	unsigned int nodeIndex = getNodeIndexForRedoAtTimelineIndex(self, timelineIndex);
	if (nodeIndex == UINT_MAX) {
		return false;
	}
	call_virtual(apply, self->nodes[nodeIndex].stateDelta, self->target);
	self->currentNodeID = self->nodes[nodeIndex].uniqueID;
	self->currentDepth++;
	return true;
}

static unsigned int getCommonParentNodeIndex(UndoTree * self, unsigned int nodeIndex1, unsigned int nodeIndex2) {
	if (nodeIndex1 == UINT_MAX || nodeIndex2 == UINT_MAX) {
		return UINT_MAX;
	}
	while (self->nodes[nodeIndex1].depth > self->nodes[nodeIndex2].depth) {
		nodeIndex1 = getNodeIndexWithID(self, self->nodes[nodeIndex1].parentID);
	}
	while (self->nodes[nodeIndex2].depth > self->nodes[nodeIndex1].depth) {
		nodeIndex2 = getNodeIndexWithID(self, self->nodes[nodeIndex2].parentID);
	}
	while (nodeIndex1 != nodeIndex2) {
		nodeIndex1 = getNodeIndexWithID(self, self->nodes[nodeIndex1].parentID);
		nodeIndex2 = getNodeIndexWithID(self, self->nodes[nodeIndex2].parentID);
	}
	return nodeIndex1;
}

bool UndoTree_traverseToNode(UndoTree * self, unsigned int nodeID) {
	unsigned int targetNodeIndex = getNodeIndexWithID(self, nodeID);
	if (targetNodeIndex == UINT_MAX && nodeID != 0) {
		return false;
	}
	unsigned int commonParentNodeIndex = UINT_MAX;
	unsigned int undoCount = 0;
	unsigned int redoCount = 0;
	if (self->currentNodeID != 0) {
		unsigned int currentNodeIndex = getNodeIndexWithID(self, self->currentNodeID);
		assert(currentNodeIndex != UINT_MAX);
		commonParentNodeIndex = getCommonParentNodeIndex(self, currentNodeIndex, targetNodeIndex);
		if (commonParentNodeIndex == UINT_MAX) {
			undoCount = self->nodes[currentNodeIndex].depth + 1;
		} else {
			undoCount = self->nodes[currentNodeIndex].depth - self->nodes[commonParentNodeIndex].depth;
		}
	}
	for (unsigned int undoIndex = 0; undoIndex < undoCount; undoIndex++) {
		UndoTree_undo(self);
	}
	if (targetNodeIndex != UINT_MAX) {
		unsigned int intermediateNodeIndex = targetNodeIndex;
		while (intermediateNodeIndex != commonParentNodeIndex) {
			unsigned int parentNodeIndex = getNodeIndexWithID(self, self->nodes[intermediateNodeIndex].parentID);
			if (parentNodeIndex == UINT_MAX) {
				self->lastTraversedRootNodeID = self->nodes[intermediateNodeIndex].uniqueID;
			} else {
				self->nodes[parentNodeIndex].lastTraversedChildID = self->nodes[intermediateNodeIndex].uniqueID;
			}
			intermediateNodeIndex = parentNodeIndex;
		}
		if (commonParentNodeIndex == UINT_MAX) {
			redoCount = self->nodes[targetNodeIndex].depth + 1;
		} else {
			redoCount = self->nodes[targetNodeIndex].depth - self->nodes[commonParentNodeIndex].depth;
		}
	}
	for (unsigned int redoIndex = 0; redoIndex < redoCount; redoIndex++) {
		UndoTree_redo(self);
	}
	return undoCount != 0 || redoCount != 0;
}

void UndoTree_removeAllNodes(UndoTree * self) {
	for (unsigned int nodeIndex = 0; nodeIndex < self->nodeCount; nodeIndex++) {
		call_virtual(dispose, self->nodes[nodeIndex].stateDelta);
	}
	self->nodeCount = 0;
	self->currentNodeID = 0;
	self->lastTraversedRootNodeID = 0;
	self->currentDepth = 0;
}