#include <fx/ds/tree.h>
#include <stdlib.h>
#include <string.h>

#define ITERATOR_RECURSIVE           0x01u

#define ITERATOR_IS_RECURSIVE(it)    (((it)->_f01 & ITERATOR_RECURSIVE) != 0)
#define ITERATOR_UNSET_RECURSIVE(it) ((it)->_f01 &= ~ITERATOR_RECURSIVE)
#define ITERATOR_SET_RECURSIVE(it)   ((it)->_f01 |= ITERATOR_RECURSIVE)

#define NODE_PARENT(n)               ((n)->__p01)
#define NODE_FIRST_CHILD(n)          ((n)->__p02)
#define NODE_NEXT_SIBLING(n)         ((n)->__p03)

/*** PRIVATE DATA *************************************************************/

struct fx_tree_p {
	struct fx_tree_node *t_root;
};

struct fx_tree_iterator_p {
	size_t i, depth;
	fx_tree_node *node;

	unsigned char _f01;
};

/*** PRIVATE FUNCTIONS ********************************************************/

static void tree_set_root(struct fx_tree_p *tree, struct fx_tree_node *node)
{
	tree->t_root = node;
}

static const struct fx_tree_node *next_node(
	const struct fx_tree_node *node, bool recursive, int *depth_diff)
{
	if (!node) {
		return NULL;
	}

	if (!recursive) {
		node = NODE_NEXT_SIBLING(node);
		return node;
	}

	int d = 0;
	struct fx_tree_node *next = NODE_FIRST_CHILD(node);
	if (next) {
		d = 1;
		*depth_diff = d;
		return next;
	}

	const struct fx_tree_node *n = node;
	next = NODE_NEXT_SIBLING(n);
	while (!next) {
		n = NODE_PARENT(n);
		if (!n) {
			break;
		}

		d--;
		next = NODE_NEXT_SIBLING(n);
	}

	*depth_diff = d;
	return next;
}

static void remove_node(struct fx_tree_node *node)
{
	struct fx_tree_node *parent = NODE_PARENT(node);
	if (!parent) {
		return;
	}

	struct fx_tree_node *n0 = NULL, *n1 = NULL;
	n0 = NODE_FIRST_CHILD(parent);

	while (n0) {
		if (n0 == node) {
			break;
		}

		n1 = n0;
		n0 = NODE_NEXT_SIBLING(n0);
	}

	if (!n0) {
		return;
	}

	if (n1) {
		NODE_NEXT_SIBLING(n1) = NODE_NEXT_SIBLING(n0);
	} else {
		NODE_FIRST_CHILD(parent) = NODE_NEXT_SIBLING(n0);
	}

	NODE_PARENT(n0) = NODE_NEXT_SIBLING(n0) = NULL;
}

static void reparent_children(
	struct fx_tree_node *old_parent, struct fx_tree_node *new_parent)
{
	struct fx_tree_node *last = NODE_FIRST_CHILD(new_parent);
	while (last && NODE_NEXT_SIBLING(last)) {
		last = NODE_NEXT_SIBLING(last);
	}

	struct fx_tree_node *cur = NODE_FIRST_CHILD(old_parent);
	while (cur) {
		struct fx_tree_node *next = NODE_NEXT_SIBLING(cur);
		NODE_PARENT(cur) = new_parent;
		NODE_NEXT_SIBLING(cur) = NULL;

		if (last) {
			NODE_NEXT_SIBLING(last) = cur;
		} else {
			NODE_FIRST_CHILD(new_parent) = cur;
		}

		last = cur;
		cur = next;
	}
}

/*** PUBLIC FUNCTIONS *********************************************************/

void fx_tree_set_root(fx_tree *tree, struct fx_tree_node *node)
{
	FX_CLASS_DISPATCH_STATIC(FX_TYPE_TREE, tree_set_root, tree, node);
}

void fx_tree_node_add_child(struct fx_tree_node *parent, struct fx_tree_node *child)
{
	if (NODE_PARENT(child)) {
		return;
	}

	NODE_PARENT(child) = parent;
	if (!NODE_FIRST_CHILD(parent)) {
		NODE_FIRST_CHILD(parent) = child;
		return;
	}

	struct fx_tree_node *cur = NODE_FIRST_CHILD(parent);
	while (NODE_NEXT_SIBLING(cur)) {
		cur = NODE_NEXT_SIBLING(cur);
	}

	NODE_NEXT_SIBLING(cur) = child;
}

void fx_tree_node_add_sibling(struct fx_tree_node *node, struct fx_tree_node *to_add)
{
	if (NODE_PARENT(to_add) || !NODE_PARENT(node)) {
		return;
	}

	fx_tree_node_add_child(NODE_PARENT(node), to_add);
}

struct fx_tree_node *fx_tree_node_get_child(struct fx_tree_node *node, size_t at)
{
	size_t i = 0;
	struct fx_tree_node *cur = NODE_FIRST_CHILD(node);

	while (i < at) {
		if (!cur) {
			return NULL;
		}

		cur = NODE_NEXT_SIBLING(cur);
		i++;
	}

	return cur;
}

fx_iterator *fx_tree_begin(fx_tree *tree)
{
	struct fx_tree_p *p = fx_object_get_private(tree, FX_TYPE_TREE);
	return fx_tree_node_begin(p->t_root);
}

const fx_iterator *fx_tree_cbegin(const fx_tree *tree)
{
	struct fx_tree_p *p = fx_object_get_private(tree, FX_TYPE_TREE);
	return fx_tree_node_begin(p->t_root);
}

fx_iterator *fx_tree_node_begin(struct fx_tree_node *node)
{
	fx_tree_iterator *it_obj = fx_object_create(FX_TYPE_TREE_ITERATOR);
	struct fx_tree_iterator_p *it
		= fx_object_get_private(it_obj, FX_TYPE_TREE_ITERATOR);

	it->node = NODE_FIRST_CHILD(node);
	it->i = 0;
	it->depth = 0;

	ITERATOR_UNSET_RECURSIVE(it);

	if (!it->node) {
		fx_iterator_set_status(it_obj, FX_ERR_NO_DATA);
	}

	return it_obj;
}

const fx_iterator *fx_tree_node_cbegin(const struct fx_tree_node *node)
{
	return fx_tree_node_begin((struct fx_tree_node *)node);
}

fx_iterator *fx_tree_node_begin_recursive(struct fx_tree_node *node)
{
	fx_tree_iterator *it_obj = fx_object_create(FX_TYPE_TREE_ITERATOR);
	struct fx_tree_iterator_p *it
		= fx_object_get_private(it_obj, FX_TYPE_TREE_ITERATOR);

	it->node = node;
	it->i = 0;
	it->depth = 0;

	ITERATOR_SET_RECURSIVE(it);

	if (!it->node) {
		fx_iterator_set_status(it_obj, FX_ERR_NO_DATA);
	}

	return it_obj;
}

const fx_iterator *fx_tree_node_cbegin_recursive(const struct fx_tree_node *node)
{
	return fx_tree_node_begin_recursive((struct fx_tree_node *)node);
}

/*** VIRTUAL FUNCTIONS ********************************************************/

static void tree_init(fx_object *obj, void *priv)
{
	struct fx_tree_p *tree = priv;
}

static void tree_fini(fx_object *obj, void *priv)
{
	struct fx_tree_p *tree = priv;
}

/*** ITERATOR FUNCTIONS *******************************************************/

static enum fx_status iterator_move_next(const fx_iterator *obj)
{
	struct fx_tree_iterator_p *it
		= fx_object_get_private(obj, FX_TYPE_TREE_ITERATOR);

	int depth_diff = 0;
	const struct fx_tree_node *next
		= next_node(it->node, ITERATOR_IS_RECURSIVE(it), &depth_diff);

	if (next) {
		it->depth += depth_diff;
		it->i++;
	} else {
		it->depth = 0;
		it->i = 0;
	}

	it->node = (struct fx_tree_node *)next;
	return (it->node != NULL) ? FX_SUCCESS : FX_ERR_NO_DATA;
}

static enum fx_status iterator_erase(fx_iterator *obj)
{
	struct fx_tree_iterator_p *it
		= fx_object_get_private(obj, FX_TYPE_TREE_ITERATOR);

	if (!it->node) {
		return FX_ERR_OUT_OF_BOUNDS;
	}

	struct fx_tree_node *parent = NODE_PARENT(it->node);
	if (!parent) {
		return FX_ERR_NOT_SUPPORTED;
	}

	int d = 0;
	struct fx_tree_node *n = it->node;
	struct fx_tree_node *next = NODE_NEXT_SIBLING(n);

	if (!next) {
		next = NODE_FIRST_CHILD(n);
	}

	while (!next) {
		n = NODE_PARENT(n);
		if (!n) {
			break;
		}

		d--;
		next = NODE_NEXT_SIBLING(n);
	}

	remove_node(it->node);
	reparent_children(it->node, parent);

	return FX_SUCCESS;
}

static fx_iterator_value iterator_get_value(fx_iterator *obj)
{
	struct fx_tree_iterator_p *it
		= fx_object_get_private(obj, FX_TYPE_TREE_ITERATOR);

	if (!it->node) {
		return FX_ITERATOR_VALUE_NULL;
	}

	return FX_ITERATOR_VALUE_PTR(it->node);
}

static const fx_iterator_value iterator_get_cvalue(const fx_iterator *obj)
{
	struct fx_tree_iterator_p *it
		= fx_object_get_private(obj, FX_TYPE_TREE_ITERATOR);

	if (!it->node) {
		return FX_ITERATOR_VALUE_NULL;
	}

	return FX_ITERATOR_VALUE_CPTR(it->node);
}

/*** CLASS DEFINITION *********************************************************/

// ---- fx_tree DEFINITION
FX_TYPE_CLASS_DEFINITION_BEGIN(fx_tree)
	FX_TYPE_CLASS_INTERFACE_BEGIN(fx_object, FX_TYPE_OBJECT)
		FX_INTERFACE_ENTRY(to_string) = NULL;
	FX_TYPE_CLASS_INTERFACE_END(fx_object, FX_TYPE_OBJECT)

	FX_TYPE_CLASS_INTERFACE_BEGIN(fx_iterable, FX_TYPE_ITERABLE)
		FX_INTERFACE_ENTRY(it_begin) = fx_tree_begin;
		FX_INTERFACE_ENTRY(it_cbegin) = fx_tree_cbegin;
	FX_TYPE_CLASS_INTERFACE_END(fx_iterable, FX_TYPE_ITERABLE)
FX_TYPE_CLASS_DEFINITION_END(fx_tree)

FX_TYPE_DEFINITION_BEGIN(fx_tree)
	FX_TYPE_ID(0x8d8fa36b, 0xc515, 0x4803, 0x8124, 0xfd704f01b8ae);
	FX_TYPE_CLASS(fx_tree_class);
	FX_TYPE_IMPLEMENTS(FX_TYPE_ITERABLE);
	FX_TYPE_INSTANCE_PRIVATE(struct fx_tree_p);
	FX_TYPE_INSTANCE_INIT(tree_init);
	FX_TYPE_INSTANCE_FINI(tree_fini);
FX_TYPE_DEFINITION_END(fx_tree)

// ---- fx_tree_iterator DEFINITION
FX_TYPE_CLASS_DEFINITION_BEGIN(fx_tree_iterator)
	FX_TYPE_CLASS_INTERFACE_BEGIN(fx_object, FX_TYPE_OBJECT)
		FX_INTERFACE_ENTRY(to_string) = NULL;
	FX_TYPE_CLASS_INTERFACE_END(fx_object, FX_TYPE_OBJECT)

	FX_TYPE_CLASS_INTERFACE_BEGIN(fx_iterator, FX_TYPE_ITERATOR)
		FX_INTERFACE_ENTRY(it_move_next) = iterator_move_next;
		FX_INTERFACE_ENTRY(it_erase) = iterator_erase;
		FX_INTERFACE_ENTRY(it_get_value) = iterator_get_value;
		FX_INTERFACE_ENTRY(it_get_cvalue) = iterator_get_cvalue;
	FX_TYPE_CLASS_INTERFACE_END(fx_iterator, FX_TYPE_ITERATOR)
FX_TYPE_CLASS_DEFINITION_END(fx_tree_iterator)

FX_TYPE_DEFINITION_BEGIN(fx_tree_iterator)
	FX_TYPE_ID(0xb896e671, 0x84b2, 0x4892, 0xaf09, 0x407f305f4bf8);
	FX_TYPE_EXTENDS(FX_TYPE_ITERATOR);
	FX_TYPE_CLASS(fx_tree_iterator_class);
	FX_TYPE_INSTANCE_PRIVATE(struct fx_tree_iterator_p);
FX_TYPE_DEFINITION_END(fx_tree_iterator)