// =====================================================================================================================
//  fennec, a free and open source game engine
//  Copyright © 2025 - 2026  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/>.
// =====================================================================================================================

///
/// \file fennec/containers/sequence.h
/// \brief
///
///
/// \details
/// \author Medusa Slockbower
///
/// \copyright Copyright © 2025 - 2026 Medusa Slockbower ([GPLv3](https://www.gnu.org/licenses/gpl-3.0.en.html))
///
///

#ifndef FENNEC_CONTAINERS_SEQUENCE_H
#define FENNEC_CONTAINERS_SEQUENCE_H

#include <fennec/containers/pair.h>
#include <fennec/containers/sequence.h>
#include <fennec/lang/compare.h>
#include <fennec/memory/allocator.h>

// https://en.wikipedia.org/wiki/Red%E2%80%93black_tree
// https://www.geeksforgeeks.org/dsa/insertion-in-red-black-tree/
// https://github.com/anandarao/Red-Black-Tree/blob/master/RBTree.cpp

// After rewriting the _fix_insert and _fix_erase functions the performance decreased significantly in the lower end
// but now in the higher end it remains consistent. Something I was doing was disturbing both the rb-tree and bst tree
// properties, now that is fixed. I'll see about optimizing more in the future.

// I realized that the way bintree is setup makes some insert calls O(n + log n) = O(n), so I switched to a pointer based model.
// This increased performance overall maintaining O(log n).

namespace fennec
{

///
///
/// \brief wrapper for ordered sets of elements, called sequences in mathematics
/// \details
/// This data-structure behaves like an ordered-set, but does not use pointers, instead storing the table in-array
///
/// | Property    | Value            |
/// |:-----------:|:----------------:|
/// | stable      |        ✅        |
/// | dynamic     |        ✅        |
/// | homogeneous |        ✅        |
/// | distinct    |        ✅        |
/// | ordered     |        ✅        |
/// | linear      |        ✅        |
/// | space       |   \emph{O(N)}    |
/// | access      | \emph{O(\log N)} |
/// | find        | \emph{O(\log N)} |
/// | insertion   | \emph{O(\log N)} |
/// | deletion    | \emph{O(\log N)} |
///
/// \tparam TypeT The type to contain
/// \tparam CompareT Function for comparing two values
/// \tparam AllocT An allocator class
template<typename TypeT, typename CompareT = less<TypeT>, class AllocT = allocator<pair<TypeT, bool>>>
struct sequence {

// Definitions =========================================================================================================
private:
	struct _node;

	enum color_ : bool {
		black = false,
		red   = true,
	};

	enum dir_ : bool {
		dir_left  = false,
		dir_right = true,
	};

public:

	/// \name Definitions
	/// @{

	using value_t   = TypeT;                                            //!< the value type
	using node_t    = pair<TypeT, bool>;                                //!< the node type
	using alloc_t   = allocator_traits<AllocT>::template rebind<_node>; //!< underlying alloc type
	using compare_t = CompareT;                                         //!< comparison type

	/// @}

	class iterator;
	class const_iterator;

private:
	using node = _node*;

	struct _node {
		node    parent;
		node    child[2];
		value_t key;
		bool    color;

		template<typename...ArgsT>
		constexpr _node(ArgsT&&...args)
			: parent(nullptr)
			, child { nullptr, nullptr }
			, key(fennec::forward<ArgsT>(args)...)
			, color(red) {
		}

		template<typename...ArgsT>
		constexpr _node(node p, node l, node r, ArgsT&&...args)
			: parent(p)
			, child { l, r }
			, key(fennec::forward<ArgsT>(args)...)
			, color(red) {
		}

		constexpr ~_node() {
			parent   = nullptr;
			child[0] = nullptr;
			child[1] = nullptr;
		}
	};


// Member Access Helpers

	static constexpr value_t& _key(node n) {
		return n->key;
	}

	static constexpr bool& _color(node n) {
		return n->color;
	}

	static constexpr node& _parent(node n) {
		return n->parent;
	}

	static constexpr node& _child(node n, bool dir) {
		return n->child[dir];
	}

	static constexpr node& _left(node n) {
		return n->child[dir_left];
	}

	static constexpr node& _right(node n) {
		return n->child[dir_right];
	}


// Safe Member Access Helpers


	constexpr const value_t& key(node n) const {
		return n->key;
	}

	constexpr bool color(node n) {
		return n ? n->color : bool(black);
	}

	constexpr node parent(node n) {
		return n ? n->parent : nullptr;
	}

	constexpr node child(node n, bool dir) {
		return n ? n->child[dir] : nullptr;
	}

	constexpr node left(node n) {
		return n ? n->child[dir_left] : nullptr;
	}

	constexpr node right(node n) {
		return n ? n->child[dir_right] : nullptr;
	}

	constexpr node leftmost(node n) {
		if (n == nullptr) {
			return nullptr;
		}

		while (this->_left(n)) {
			n = this->_left(n);
		}
		return n;
	}

	constexpr node rightmost(node n) {
		if (n == nullptr) {
			return nullptr;
		}

		while (this->_right(n)) {
			n = this->_right(n);
		}
		return n;
	}

// Constructors & Destructors ==========================================================================================
public:
	/// \name Constructors & Destructor
	/// @{

	///
	/// \brief Default Constructor, initializes an empty sequence
	///
	/// \par Complexity
	/// \emph{O(1)}
	///
	constexpr sequence()
		: _root(nullptr), _size(0) {
	}

	///
	/// \brief Move Constructor, takes ownership of a sequence
	///
	/// \par Complexity
	/// \emph{O(1)}
	///
	constexpr sequence(sequence&&) noexcept = default;

	///
	/// \brief Copy Constructor, copies a sequence
	constexpr sequence(const sequence&) = default;
	// TODO: properly implement

	///
	/// \brief Default Destructor, destructs elements *in-order*
	///
	/// \par Complexity
	/// \emph{O(N)}
	///
	constexpr ~sequence() {
		this->clear();
	}

	/// @}

// Search ==============================================================================================================
public:
	/// \name Search
	/// @{

	///
	/// \brief Value Find Function, finds the iterator position for \emph{val}, otherwise returns \emph{end()}
	/// \param val The value to find
	/// \returns An iterator at the value
	///
	/// \par Complexity
	/// \emph{O(\log N)}
	///
	constexpr iterator find(const value_t& val) {
		node node = _root;
		while (node) {
			if (_compare(val, _key(node))) {
				node = _left(node);
			} else if (_compare(_key(node), val)) {
				node = _right(node);
			} else {
				return sequence::iterator(this, _root, node);
			}
		}
		return sequence::iterator(this, _root, node);
	}

	///
	/// \brief Value Contains Function, checks if the sequence contains a value
	/// \param val The value to find
	/// \returns \emph{true} if \emph{val} is in the sequence, \emph{false} otherwise
	bool contains(const value_t& val) {
		return find(val) != end();
	}

	/// @}


// Properties ==========================================================================================================
public:
	/// \name Properties
	/// @{

	///
	/// \returns The number of elements in the sequence
	///
	/// \par Complexity
	/// \emph{O(1)}
	///
	constexpr size_t size() const {
		return _size;
	}

	///
	/// \returns \emph{true} when there are no elements in the sequence, \emph{false} otherwise.
	///
	/// \par Complexity
	/// \emph{O(1)}
	///
	constexpr bool is_empty() const {
		return _size == 0;
	}

	/// @}

// Modifiers ===========================================================================================================
public:
	/// \name Modifiers
	/// @{

	///
	/// \brief Move Insertion, moves \emph{val} into the sequence
	/// \param val The value to insert
	///
	/// \par Complexity
	/// \emph{O(\log N)}
	///
	constexpr void insert(value_t&& val) {
		node i = _insert_bst(fennec::forward<value_t>(val));
		_fix_insert(i);
	}

	///
	/// \brief Copy Insertion, inserts a copy of \emph{val} into the sequence
	/// \param val The value to insert
	///
	/// \par Complexity
	/// \emph{O(\log N)}
	///
	constexpr void insert(const value_t& val) {
		node i = _insert_bst(val);
		_fix_insert(i);
	}

	///
	/// \brief Emplacement, constructs and adds a value into the sequence
	/// \tparam ArgsT The argument types
	/// \param args The arguments to construct with
	///
	/// \par Complexity
	/// \emph{O(\log N)}
	///
	template<typename...ArgsT>
	constexpr void emplace(ArgsT&&...args) {
		node i = _insert_bst(fennec::forward<ArgsT>(args)...);
		_fix_insert(i);
	}

	///
	/// \brief Erase the specified value from the sequence
	/// \param val the value to erase
	///
	/// \par Complexity
	/// \emph{O(\log N)}
	///
	constexpr void erase(const value_t& val) {
		_erase(find(val)._node);
	}

	///
	/// \brief Destructs all elements, *in-order*, contained in the sequence
	///
	/// \par Complexity
	/// \emph{O(N)}
	///
	constexpr void clear() {
		list<node> visit;
		for (iterator it = begin(); it != end(); ++it) {
			visit.push_back(it._node);
		}
		for (node n : visit) {
			this->_free_node(n);
		}
		_root = nullptr;
		_size = 0;
	}

	/// @}

// Iterator ============================================================================================================

	///
	/// \returns An iterator at the smallest element in the sequence
	///
	/// \par Complexity
	/// \emph{O(\log N)}
	///
	constexpr iterator begin() {
		return sequence::iterator(this, _root);
	}

	///
	/// \brief C++ Iterator Specification \emph{begin()}
	/// \returns An iterator at the smallest element in the sequence
	///
	/// \par Complexity
	/// \emph{O(\log N)}
	///
	constexpr const_iterator begin() const {
		return sequence::const_iterator(this, _root);
	}

	///
	/// \returns An iterator after the largest element in the sequence
	///
	/// \par Complexity
	/// \emph{O(1)}
	///
	constexpr iterator end() {
		return sequence::iterator(this, _root, nullptr);
	}

	///
	/// \brief Const C++ Iterator Specification \emph{end()}
	/// \returns An iterator after the largest element in the sequence
	///
	/// \par Complexity
	/// \emph{O(1)}
	///
	constexpr const_iterator end() const {
		return sequence::const_iterator(this, _root, nullptr);
	}

	///
	/// \brief C++ Iterator Specification \emph{iterator}
	class iterator {
	public:

		///
		/// \brief prefix increment operator
		/// \param it the iterator to increment
		/// \returns \emph{it} after having moved to the next element in the list
		friend iterator& operator++(iterator& it) {
			if (it._node == nullptr) {
				return it;
			}

			node parent = it._seq->_parent(it._node);
			node pright = it._seq->right(parent);
			node next   = it._seq->leftmost(it._seq->right(it._node));

			if (it._node != pright && parent != nullptr) {
				it._visit.push_front(parent);
			}

			if (next != nullptr) {
				it._visit.push_front(next);
			}

			if (not it._visit.is_empty()) {
				it._node = it._visit.front();
				it._visit.pop_front();
			} else {
				it._node = nullptr;
			}

			return it;
		}

		///
		/// \brief postfix increment operator
		/// \param it the iterator to increment
		/// \returns \emph{it} before having moved to the next element in the list
		friend iterator operator++(iterator& it, int) {
			iterator prev = it;
			++it;
			return prev;
		}

		///
		/// \brief dereference operator
		/// \returns a reference to the value pointed by the iterator
		const value_t& operator*() const {
			return _node->key;
		}

		///
		/// \brief pointer access operator
		/// \returns a pointer to the value pointed by the iterator
		const value_t* operator->() const {
			return &_node->key;
		}

		///
		/// \brief iterator equality operator
		/// \param lhs the iterator
		/// \param rhs the iterator to compare with
		/// \returns \emph{true} if the iterators are identical, \emph{false} otherwise
		constexpr friend bool operator==(const iterator& lhs, const iterator& rhs) {
			return lhs._seq == rhs._seq and lhs._node == rhs._node;
		}

		///
		/// \brief iterator inequality operator
		/// \param lhs the iterator
		/// \param rhs the iterator to compare with
		/// \returns \emph{true} if the iterators are different, \emph{false} otherwise
		constexpr friend bool operator!=(const iterator& lhs, const iterator& rhs) {
			return lhs._seq != rhs._seq or lhs._node != rhs._node;
		}

	private:
		sequence*  _seq;
		node       _head;
		node       _node;
		list<node> _visit;

		constexpr iterator(sequence* seq, node start)
			: _seq(seq)
			, _head(start)
			, _node(seq->leftmost(start)) {
		}

		constexpr iterator(sequence* seq, node root, node start)
			: _seq(seq)
			, _head(root)
			, _node(start) {
		}

		friend struct sequence;
	};

	///
	/// \brief C++ Iterator Specification \emph{iterator}
	class const_iterator {
	public:

		///
		/// \brief prefix increment operator
		/// \param it the iterator to increment
		/// \returns \emph{it} after having moved to the next element in the list
		friend const_iterator& operator++(const_iterator& it) {
			if (it._node == nullptr) {
				return it;
			}

			node parent = it._seq->_parent(it._node);
			node pright = it._seq->right(parent);
			node next   = it._seq->leftmost(it._seq->right(it._node));

			if (it._node != pright && parent != nullptr) {
				it._visit.push_front(parent);
			}

			if (next != nullptr) {
				it._visit.push_front(next);
			}

			if (not it._visit.is_empty()) {
				it._node = it._visit.front();
				it._visit.pop_front();
			} else {
				it._node = nullptr;
			}

			return it;
		}

		///
		/// \brief postfix increment operator
		/// \param it the iterator to increment
		/// \returns \emph{it} before having moved to the next element in the list
		friend const_iterator operator++(const_iterator& it, int) {
			const_iterator prev = it;
			++it;
			return prev;
		}

		///
		/// \brief dereference operator
		/// \returns a reference to the value pointed by the iterator
		const value_t& operator*() const {
			return _node->key;
		}

		///
		/// \brief pointer access operator
		/// \returns a pointer to the value pointed by the iterator
		const value_t* operator->() const {
			return &_node->key;
		}

		///
		/// \brief iterator equality operator
		/// \param lhs the iterator
		/// \param rhs the iterator to compare with
		/// \returns \emph{true} if the iterators are identical, \emph{false} otherwise
		constexpr friend bool operator==(const const_iterator& lhs, const const_iterator& rhs) {
			return lhs._seq == rhs._seq and lhs._node == rhs._node;
		}

		///
		/// \brief iterator inequality operator
		/// \param lhs the iterator
		/// \param rhs the iterator to compare with
		/// \returns \emph{true} if the iterators are different, \emph{false} otherwise
		constexpr friend bool operator!=(const const_iterator& lhs, const const_iterator& rhs) {
			return lhs._seq != rhs._seq or lhs._node != rhs._node;
		}

	private:
		const sequence*  _seq;
		node             _head;
		node             _node;
		list<node>       _visit;

		constexpr const_iterator(const sequence* seq, node start)
			: _seq(seq)
			, _head(start)
			, _node(seq->leftmost(start)) {
		}

		constexpr const_iterator(const sequence* seq, node root, node start)
			: _seq(seq)
			, _head(root)
			, _node(start) {
		}

		friend struct sequence;
	};


// Private Member Variables ============================================================================================
private:
	alloc_t   _alloc;
	node      _root;
	compare_t _compare;
	size_t    _size;


// Private Helpers =====================================================================================================
private:

	template<typename...ArgsT>
	constexpr node _make_node(ArgsT&&...args) {
		node res = _alloc.allocate(1);
		fennec::construct(res, fennec::forward<ArgsT>(args)...);
		return res;
	}

	constexpr void _free_node(node n) {
		fennec::destruct(n);
		_alloc.deallocate(n);
	}

	constexpr node _rotate(node sub, bool dir) {
		if (sub == nullptr) {
			return nullptr;
		}
		node sub_parent = _parent(sub);
		node new_root   = _child(sub, not dir);
		node new_child  = _child(new_root, dir);

		_child(sub, not dir) = new_child;
		if (new_child != nullptr) {
			_parent(new_child) = sub;
		}
		_child(new_root, dir) = sub;

		_parent(new_root) = sub_parent;
		_parent(sub) = new_root;
		if (sub_parent != nullptr) {
			_child(sub_parent, sub == _right(sub_parent)) = new_root;
		} else {
			_root = new_root;
		}
		return new_root;
	}

	constexpr void _recolor(node n) {
		bool c = color(n) == black;
		if (n == _root) { // Only recolor if not the root node
			_color(n) = c;
		}
		_color(_left(n))  = !_color(_left(n)) ;
		_color(_right(n)) = !_color(_right(n));
	}

	// run-of-the-mill bst insert
	template<typename...ArgsT>
	constexpr node _insert_bst(ArgsT&&...args) {
		node res = _make_node(fennec::forward<ArgsT>(args)...);

		if (_root == nullptr) {
			++_size;
			_color(res) = black;
			return _root = res;
		}

		node i = _root;
		node p = nullptr;
		bool d = dir_left;
		while (i != nullptr) {
			p = i;

			if (_compare(_key(res), _key(i))) {
				i = _left(i);
				d = dir_left;
			} else if (_compare(_key(i), _key(res))) {
				i = _right(i);
				d = dir_right;
			} else {
				_free_node(res);
				return nullptr;
			}
		}

		++_size;
		_child(p, d) = res;
		_parent(res) = p;
		return res;
	}

	// This makes some cheats given that the structure is modified only by internal functions
	// If such is the case, ONLY LL, LR, RL, and RR will show up
	// Then we just need to handle splitting a 4-node
	constexpr void _fix_insert(node n) {
		if (n == nullptr) {
			return;
		}

		node p = _parent(n);
		while (n != _root && color(n) == red && color(p) == red) {
			node g = _parent(p);
			bool d = n == _right(p);
			bool r = p == _right(g);
			node u = _child(g, !r);

			if (color(u) == red) {
				_recolor(g);
				n = g;
				p = _parent(n);
				continue;
			}

			if (d != r) {
				_rotate(p, r);
				n = p;
				p = _parent(n);
			}

			_rotate(g, not r);
			fennec::swap(_color(p), _color(g));
			n = p;
			p = _parent(n);
		}
		_color(_root) = black;
	}

	constexpr void _swap_val(node a, node b) {
		fennec::swap(_key(a), _key(b));
	}

	constexpr node _red_child(node x) {
		node l = _left(x);
		node r = _right(x);

		if (color(l) == red) {
			return l;
		}

		if (color(r) == red) {
			return r;
		}

		return nullptr;
	}

	constexpr void _fix_erase(node n) {
		if (n == nullptr) {
			return;
		}

		if (n == _root) {
			_root = nullptr;
			return;
		}

		node o = n;
		node p = _parent(n);
		if (p == nullptr) {
			_root = nullptr;
			return;
		}

		bool d = n == _right(p);
		node c = _red_child(n);
		node s = nullptr;
		if (_color(n) == red || c != nullptr) {
			_child(p, d) = c;
			if (c != nullptr) {
				_parent(c) = p;
			}
			_color(c) = black;
			return;
		}

		while (n != _root) {
			p = _parent(n);
			d = n == _right(p);
			s = _child(p, !d);

			if (s == nullptr) {
				break;
			}

			if (_color(s) == red) {
				_color(s) = black;
				_color(p) = red;
				_rotate(p, d);
				continue;
			}

			node nc = _child(s, d);
			node nf = _child(s, !d);

			if (color(nc) == black && color(nf) == black) {
				_color(s) = red;
				if (_color(p) == red) {
					_color(p) = black;
					break;
				}
				n = p;
				continue;
			}

			if (color(nf) == black) {
				_color(nc) = black;
				_color(s) = red;
				_rotate(s, !d);
				s  = nc;
				nf = s;
			}
			_color(s) = _color(p);
			_color(p) = black;
			_color(nf) = black;
			_rotate(p, d);
			break;
		}

		p = parent(o);
		if (p != nullptr) {
			if (o == _left(p)) {
				_left(p) = nullptr;
			} else {
				_right(p) = nullptr;
			}
			_color(_root) = black;
		} else {
			_root = nullptr;
		}
	}

	constexpr void _erase(node n) {
		if (n == nullptr) {
			return;
		}

		node l = _left(n);
		node r = _right(n);

		// 2 children
		if (l != nullptr && r != nullptr) {
			node s = leftmost(r);
			_swap_val(n, s);
			n = s;
			l = _left(n);
			r = _right(n);
		}

		node p = _parent(n);
		bool d = n == right(p);
		node c = l != nullptr ? l : r;

		// Single child
		if (c != nullptr) {
			_parent(c) = p;
		}

		// Handles root cases
		if (p == nullptr) {
			_root = c;
			if (c == nullptr) {
				_free_node(n);
				--_size;
				return;
			} else {
				_color(c) = black;
			}
		}

		// Single Child, Red, and Root cases
		if (p == nullptr || c != nullptr || _color(n) == red) {
			if (p != nullptr) {
				_child(p, d) = c;
			}
			_free_node(n);
			--_size;
			return;
		}

		_fix_erase(n);
		_free_node(n);
		--_size;
	}
};

}

#endif // FENNEC_CONTAINERS_SEQUENCE_H