#include "TracingHeap.h"

#include "../pointers/pointers.h"

int TracingHeap_Create(TracingHeap* self, allocator_t* allocator)
{
	if (self == NULL) {
		return EDESTADDRREQ;
	}
	if (1==0
			|| self->config.trace_data == NULL
			|| self->config.destructor == NULL
			|| self->config.size_query == NULL) {
		return EBADSTATE;
	}

	self->allocator = allocator;

	self->objects = NULL;
	self->black_objects = NULL;
	self->white_objects = NULL;
	self->grey_objects  = NULL;

	self->reachable_color   = TRACINGCOLOR_WHITE;
	self->unreachable_color = TRACINGCOLOR_BLACK;

	return EXIT_SUCCESS;
}

static void _TracingHeap_ForceDestroyObject(TracingHeap* self, TracingObject* object)
{
	self->config.destructor(self->config.destructor_context, object->data);
	size_t allocated_size = self->config.size_query(self->config.size_query_context, object->data);
	Allocator_Free(self->allocator, object, allocated_size);
}

void TracingHeap_Destroy(TracingHeap* self)
{
	if (self == NULL) {
		return;
	}

	for (TracingObject* object = self->objects; object != NULL; object = object->global_next) {
		_TracingHeap_ForceDestroyObject(self, object);
	}

	memset(self, 0, sizeof(*self));
}

static TracingObject* _TracingObject_FromDataPointer(void* datap)
{
	return rewindp(datap, sizeof(TracingObject));
}

static void _TracingObject_ChangeObjectColor(TracingObject* object, enum TracingColor new_color, TracingObject** list_start)
{
	if (object->color != new_color) {
		object->color_prev->color_next = object->color_next;
		object->color_next->color_prev = object->color_prev;

		object->color = new_color;
		object->color_prev = NULL;
		object->color_next = *list_start;
		(*list_start)->color_prev = object;
		*list_start = object;
	}
}

static TracingObject** _TracingHeap_ColorListStart(TracingHeap* self, enum TracingColor color)
{
	TracingObject** list_start;
	switch (color) {
		case TRACINGCOLOR_BLACK:
			list_start = &self->black_objects;
			break;
		case TRACINGCOLOR_GREY:
			list_start = &self->grey_objects;
			break;
		case TRACINGCOLOR_WHITE:
			list_start = &self->white_objects;
			break;
        }

	return list_start;
}

static void _TracingHeap_MakeReachable(TracingHeap* self, TracingObject* object)
{
	TracingObject** reachable_list_start = _TracingHeap_ColorListStart(self, self->reachable_color);
	_TracingObject_ChangeObjectColor(object, self->reachable_color, reachable_list_start);
}

static TracingObject* _TracingHeap_PopColorList(TracingHeap* self, enum TracingColor color)
{
	TracingObject** list_start = _TracingHeap_ColorListStart(self, color);

	TracingObject* popped = *list_start;
	popped->color_next->color_prev = NULL; // unlink next element from popped
	*list_start = popped->color_next; // link list to next element
	popped->color_next = NULL; // unlink popped from next element

	return popped;
}

static int _TracingHeap_TraceReferenceCallback(void* context, void* reference)
{
	TracingHeap*   self   = context;
	TracingObject* object = _TracingObject_FromDataPointer(reference);

	if (object->color == self->unreachable_color) {
		_TracingObject_ChangeObjectColor(object, TRACINGCOLOR_GREY, &self->grey_objects);
	}

	return EXIT_SUCCESS;
}

int TracingHeap_AddTracingRoot(TracingHeap* self, void* data)
{
	if (self == NULL) {
		return EDESTADDRREQ;
	}
	
	TracingObject* root_object = _TracingObject_FromDataPointer(data);
	_TracingObject_ChangeObjectColor(root_object, TRACINGCOLOR_GREY, &self->grey_objects);

	return EXIT_SUCCESS;
}

int TracingHeap_TraceNext(TracingHeap* self)
{
	if (self == NULL) {
		return EDESTADDRREQ;
	}

	if (self->grey_objects == NULL) {
		return EBADSTATE;
	}

	TracingObject* trace_object = _TracingHeap_PopColorList(self, TRACINGCOLOR_GREY);

	int tracing_code = self->config.trace_data(
		self->config.trace_context,
		trace_object->data,
		_TracingHeap_TraceReferenceCallback,
		self
	);
	if (tracing_code) {
		return tracing_code;
	}

	_TracingHeap_MakeReachable(self, trace_object);

	return EXIT_SUCCESS;
}

size_t TracingHeap_TraceNextN(TracingHeap* self, size_t n)
{
	for (size_t i = 0; i < n; i++) {
		if (TracingHeap_TraceNext(self)) {
			return i;
		}
	}

	return n;
}

static int _TracingHeap_DestroyObject(TracingHeap* self, TracingObject* unreachable_object)
{
	int destructor_code = self->config.destructor(self->config.destructor_context, unreachable_object->data);
	if (destructor_code) {
		return destructor_code;
	}

	size_t allocated_size = self->config.size_query(self->config.size_query_context, unreachable_object->data);

	Allocator_Free(self->allocator, unreachable_object, allocated_size);

	return EXIT_SUCCESS;
}

int TracingHeap_BeginTrace(TracingHeap* self)
{
	if (self == NULL) {
		return EDESTADDRREQ;
	}
	if (self->grey_objects != NULL) {
		return EBADSTATE;
	}
	if (*_TracingHeap_ColorListStart(self, self->unreachable_color) == NULL) {
		return EBADSTATE;
	}

	// swap reachable and unreachable color
	// xor swap algorithm!!
	self->reachable_color   = self->reachable_color ^ self->unreachable_color;
	self->unreachable_color = self->unreachable_color ^ self->reachable_color;
	self->reachable_color   = self->reachable_color ^ self->unreachable_color;

	return EXIT_SUCCESS;
}

int TracingHeap_EndTrace(TracingHeap* self)
{
	if (self == NULL) {
		return EDESTADDRREQ;
	}
	if (self->grey_objects != NULL) {
		return EBADSTATE;
	}

	TracingObject** unreachable_list = _TracingHeap_ColorListStart(self, self->unreachable_color);
	TracingObject* unreachable_object = *unreachable_list;

	while (unreachable_object != NULL) {
		TracingObject* next_unreachable = unreachable_object->global_next;

		if (_TracingHeap_DestroyObject(self, unreachable_object)) {
			*unreachable_list = unreachable_object;
			unreachable_object->global_prev = NULL;
			return ECANCELED;
		}

		unreachable_object = next_unreachable;
	}

	*unreachable_list = NULL;

	return EXIT_SUCCESS;
}

void* TracingHeap_Allocate(TracingHeap* self, size_t bytes)
{
	if (self == NULL) {
		return NULL;
	}

	TracingObject* object = Allocator_Allocate(self->allocator, sizeof(TracingObject) + bytes);
	if (object == NULL) {
		return NULL;
	}

	object->color = self->reachable_color;

	TracingObject** reachable_list_start = _TracingHeap_ColorListStart(self, self->reachable_color);
	object->color_next = *reachable_list_start;
	*reachable_list_start = object;
	object->color_prev = NULL;


	object->global_next = self->objects; // link before global object list
	object->global_prev = NULL; // no previous object
	self->objects->global_prev = object; // link first object via previous link
	self->objects = object; // insert into linked list

	return object->data;
}