// =====================================================================================================================
//  fennec, a free and open source game engine
//  Copyright © 2025  Medusa Slockbower
//
//  This program is free software: you can redistribute it and/or modify
//  it under the terms of the GNU General Public License as published by
//  the Free Software Foundation, either version 3 of the License, or
//  (at your option) any later version.
//
//  This program is distributed in the hope that it will be useful,
//  but WITHOUT ANY WARRANTY; without even the implied warranty of
//  MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
//  GNU General Public License for more details.
//
//  You should have received a copy of the GNU General Public License
//  along with this program.  If not, see <https://www.gnu.org/licenses/>.
// =====================================================================================================================

#ifndef FENNEC_CONTAINERS_RDTREE_H
#define FENNEC_CONTAINERS_RDTREE_H

#include <fennec/containers/list.h>
#include <fennec/containers/optional.h>
#include <fennec/containers/traversal.h>
#include <fennec/memory/allocator.h>

namespace fennec
{

///
/// \brief Rooted-Directed Tree
/// \tparam TypeT Data type
/// \tparam AllocT Allocator Type
template<typename TypeT, typename AllocT = allocator<TypeT>>
struct rdtree {

// Definitions =========================================================================================================
protected:
	struct node;

public:
	using value_t = TypeT;
	using alloc_t = typename allocator_traits<AllocT>::template rebind<node>;
	static constexpr size_t root = 0;
	static constexpr size_t npos = -1;

	enum control_ : uint8_t {
		control_continue = 0,
		control_jump_over,
		control_break
	};

protected:
	struct node {
		optional<TypeT> value;
		size_t parent, child, prev, next;

		template<typename...ArgsT>
		constexpr node(size_t p, size_t c, size_t v, size_t n, ArgsT&&...args)
			: value(fennec::forward<ArgsT>(args)...)
			, parent(p), child(c), prev(v), next(n) {
		}

		constexpr ~node() {
			parent = npos;
			child  = npos;
			prev   = npos;
			next   = npos;
		}
	};

public:

// Constructors ========================================================================================================

	template<typename...ArgsT>
	explicit constexpr rdtree(ArgsT&&...args)
		: _table(), _freed(), _size(1) {
		_table.callocate(8);
		fennec::construct(&_table[0], npos, npos, npos, npos, fennec::forward<ArgsT>(args)...);
	}

	constexpr rdtree(const rdtree& tree)
		: _table(tree._table), _freed(tree._freed), _size(tree._size) {
	}

	constexpr rdtree(rdtree&& tree) noexcept
		: _table(fennec::move(tree._table)), _freed(fennec::move(tree._freed)), _size(tree._size) {
	}


// Assignment ==========================================================================================================

	constexpr rdtree& operator=(const rdtree& rhs) {
		for (value_t* it : this->_table) {
			fennec::destruct(it);
		}
		_table  = rhs._table;
		_freed = rhs._freed;
		_size  = rhs._size;
		return *this;
	}

	constexpr rdtree& operator=(rdtree&& rhs) noexcept {
		for (value_t* it : _table) {
			fennec::destruct(it);
		}
		_table  = fennec::move(rhs._table);
		_freed = fennec::move(rhs._freed);
		_size  = rhs._size;
		return *this;
	}


// Properties ==========================================================================================================

	constexpr size_t size() const {
		return _size;
	}

	constexpr size_t capacity() const {
		return _table.capacity();
	}

	constexpr bool empty() const {
		return _size == 0;
	}


// Access ==============================================================================================================

	///
	/// \param i The id of the node to check
	/// \returns The id of the parent node
	constexpr size_t parent(size_t i) const {
		return i == npos ? npos : _table[i].parent;
	}

	///
	/// \param i The id of the node to check
	/// \returns The id of the child node
	constexpr size_t child(size_t i) const {
		return i == npos ? npos : _table[i].child;
	}

	///
	/// \param i The id of the node to check
	/// \returns The id of the next node
	constexpr size_t next(size_t i) const {
		return i == npos ? npos : _table[i].next;
	}

	///
	/// \param i The id of the node to check
	/// \returns The id of the previous node
	constexpr size_t prev(size_t i) const {
		return i == npos ? npos : _table[i].prev;
	}

	///
	/// \param i the node to start at
	/// \returns the left-most child of node `i`
	constexpr size_t left_most(size_t i) const {
		size_t n = i;
		if ((n = child(n)) == npos) {
			return i;
		}
		while (true) {
			size_t p = n;
			if ((n = child(n)) == npos) {
				return p;
			}
		}
	}

	///
	/// \param i the node to start at
	/// \returns the right-most child of node `i`
	constexpr size_t right_most(size_t i) const {
		if ((i = child(i)) == npos) {
			return npos;
		}
		while (true) {
			size_t n;
			while ((n = next(i)) != npos) {
				i = n;
			}
			n = i;
			if ((i = child(i)) == npos) {
				return n;
			}
		}
	}

	///
	/// \param i The id of the node to access
	/// \returns A reference to the value of the node wrapped in an optional
	constexpr optional<value_t>& operator[](size_t i) {
		return _table[i].value;
	}

	///
	/// \param i The id of the node to access
	/// \returns A const-qualified reference to the value of the node wrapped in an optional
	constexpr const optional<value_t>& operator[](size_t i) const {
		return _table[i].value;
	}


// Insertion & Deletion ================================================================================================

	///
	/// \brief Insertion, creates a node in the tree with parent `parent`
	/// \param parent the parent node, if `npos` sets the value of the root node
	/// \param val the value to insert
	/// \returns the index of the created node
	constexpr size_t insert(size_t parent, const value_t& val) {
		return this->_insert(parent, val);
	}

	///
	/// \brief Insertion, creates a node in the tree with parent `parent`
	/// \param parent the parent node, if `npos` sets the value of the root node
	/// \param val the value to insert
	/// \returns the index of the created node
	constexpr size_t insert(size_t parent, value_t&& val) {
		return this->_insert(parent, fennec::forward<value_t>(val));
	}

	///
	/// \brief Erase a node in the tree and all of it's children
	/// \param i the index of the node
	constexpr void erase(size_t i) {
		_erase(i);
	}


// Traversal ===========================================================================================================

	struct pre_order {
		list<size_t> visit;

		size_t operator()(const rdtree&, size_t start) {
			return start;
		}

		size_t operator[](const rdtree& tree, size_t node, uint8_t mode) {
			if (not visit.empty()) {
				node = visit.front();
				visit.pop_front();
			} else {
				return npos;
			}

			size_t next = tree.next(node);
			if (next != npos) {
				visit.push_front(next);
			}

			size_t child = tree.child(node);
			if (child != npos && mode != traversal_control_jump_over) {
				visit.push_front(child);
			}

			return node;
		}
	};

	struct in_order {
		list<size_t> visit;
		size_t       head;

		size_t operator()(const rdtree& tree, size_t start) {
			head = start;
			return tree.left_most(start);
		}

		size_t operator[](const rdtree& tree, size_t node, uint8_t) {
			if (not visit.empty()) {
				node = visit.front();
				visit.pop_front();
			} else {
				return npos;
			}

			size_t parent = tree.parent(node);
			size_t next   = tree.next(node);
			size_t nnext  = next == npos ? npos : tree.next(next);

			if (next != npos && node != head) {
				if (nnext == npos && parent != npos) {
					visit.push_back(parent);
				}

				visit.push_back(tree.left_most(next));
			}

			return node;
		}
	};

	struct post_order {
		list<size_t> visit;
		size_t       head;

		size_t operator()(const rdtree& tree, size_t start) {
			head = start;
			return tree.left_most(start);
		}

		size_t operator[](const rdtree& tree, size_t node, uint8_t) {
			if (not visit.empty()) {
				node = visit.front();
				visit.pop_front();
			} else {
				return npos;
			}

			if (node == head) {
				return node;
			}

			size_t next = tree.next(node);
			visit.push_back(next != npos ? tree.left_most(next) : tree.parent(node));

			return node;
		}
	};

	template<typename OrderT, typename VisitorT>
	void traverse(VisitorT&& visit, size_t i = root) {
		OrderT order;
		i = order(*this, i);
		while (i != npos) {
			uint8_t mode = visit(*_table[i].value);
			if (mode == traversal_control_break) {
				break;
			}
			i = order[*this, i, mode];
		}
	}


protected:
	allocation<node, alloc_t> _table;
	list<size_t>              _freed;
	size_t                    _size;

	void _expand() {
		_table.creallocate(_table.capacity() * 2);
	}

	size_t _next_free() {
		size_t next = _size;
		if (not _freed.empty()) {
			next = _freed.back();
			_freed.pop_back();
		}
		if (_size >= capacity()) {
			_expand();
		}
		++_size;
		return next;
	}

	template<typename...ArgsT>
	constexpr size_t _insert(size_t p, ArgsT&&...args) {
		if (_size == 0) {
			fennec::construct(&_table[root], npos, npos, npos, npos, fennec::forward<ArgsT>(args)...);
			_size = 1;
			return root;
		}

		if (p == npos) {
			_table[root].value = value_t(fennec::forward<ArgsT>(args)...);
			_size = _size == 0 ? 1 : _size;
			return root;
		}

		size_t i = _next_free();
		size_t n = child(p);
		_table[p].child = i;
		if (n != npos) _table[n].prev  = i;
		fennec::construct(&_table[i], p, npos, npos, n, fennec::forward<ArgsT>(args)...);
		return i;
	}

	constexpr void _erase(size_t i) {
		list<size_t> queue;
		size_t j = 0;
		queue.push_back(child(i));
		while (queue.empty() == false) {
			size_t n = queue.front(); queue.pop_front();
			if (n == npos) continue;
			queue.push_back(next(n));
			queue.push_back(child(n));
			fennec::destruct(&_table[n]);
			_freed.push_back(n);
			--_size;
			++j;
		}

		fennec::destruct(&_table[i]);
		if (i != root) _freed.push_back(i);
		--_size;
	}
};

}

#endif // FENNEC_CONTAINERS_RDTREE_H
