/* ScummVM - Graphic Adventure Engine
 *
 * ScummVM is the legal property of its developers, whose names
 * are too numerous to list here. Please refer to the COPYRIGHT
 * file distributed with this source distribution.
 *
 * 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 <http://www.gnu.org/licenses/>.
 *
 */

#ifndef COMMON_RB_TREE_H
#define COMMON_RB_TREE_H

#include "common/func.h"
#include "common/util.h"
#include "common/scummsys.h"

namespace Common {

template<typename Key, typename Val>
struct PairFirst {
	const Key &operator()(const Pair<Key, Val> &p) {
		return p.first;
	}
};

template<typename T>
struct Identity {
	const T &operator()(const T &t) {
		return t;
	}
};

/**
 * @defgroup common_rb_tree Red-black tree
 * @ingroup common
 *
 * @brief API for operating on a red black tree.
 *
 * @{
 */

/**
 * Red-black tree implementation with insertion and deletion algorithms based on the ones
 * found in Introduction to Algorithms by Cormen, Leiserson, Rivest and Stein.
 * Used in the implementation of Common::StableMap and Common::MultiMap
 */
template<class ValueType, class Key, class KeyProj, class KeyComp = Common::Less<Key> >
class RBTree {
public:
	enum class Color {
		kRed,
		kBlack,
	};

	struct Node {
		Node *parent;
		Node *left;
		Node *right;
		Color color;
		ValueType value;
	};

	template<typename Ref, typename Ptr>
	class Iterator {
	public:
		friend RBTree;

		Iterator() = default;
		Iterator(const Iterator &) = default;
		Iterator &operator=(const Iterator &) = default;

		Ref operator*() const { return _current->value; }

		Ptr operator->() const { return &_current->value; }

		Iterator &operator++() {
			if (_current->right) {
				_current = leftmost(_current->right);
			} else {
				auto p = _current->parent;
				while (p && p->right == _current) {
					_current = p;
					p = p->parent;
				}
				_current = p;
			}
			return *this;
		}

		Iterator operator++(int) {
			auto temp = *this;
			++(*this);
			return temp;
		}

		bool operator==(const Iterator &rhs) const {
			return _current == rhs._current;
		}

		bool operator!=(const Iterator &rhs) const {
			return _current != rhs._current;
		}

	private:
		explicit Iterator(Node *n) : _current(n) {}

		Node *_current = nullptr;
	};

	using BasicIterator = Iterator<ValueType &, ValueType *>;             /*!< RBTree iterator. */
	using ConstIterator = Iterator<const ValueType &, const ValueType *>; /*!< Const-qualified RBTree iterator. */

	RBTree(KeyComp comp = {}) : _comp(Common::move(comp)) {
	}

	/** Construct an RBTree as a copy of the given tree @p other . */
	RBTree(const RBTree &other) : _comp(other._comp) {
		for (const auto &val : other)
			insert(val);
	}

	/** Construct an RBTree by moving the contents of the given tree (using the C++11 move semantics). */
	RBTree(RBTree &&other) : _root(other._root), _leftmost(other._leftmost), _size(other._size), _comp(Common::move(other._comp)) {
	}

	/** Assign the given tree to this tree. */
	RBTree &operator=(const RBTree &rhs) {
		*this = RBTree(rhs);
		return *this;
	}

	/** Moves the contents of the given tree into this tree (using the C++11 move semantics). */
	RBTree &operator=(RBTree &&rhs) {
		clear();
		_root = rhs._root;
		_leftmost = rhs._leftmost;
		_size = rhs._size;
		_comp = Common::move(rhs._comp);
		rhs._root = nullptr;
		rhs._leftmost = nullptr;
		rhs._size = 0;
		return *this;
	}

	/** Clears the contents of the tree. */
	void clear() {
		erase(begin(), end());
		_size = 0;
		_root = nullptr;
		_leftmost = nullptr;
	}

	/** Return an iterator pointing to the first element in the tree. */
	BasicIterator begin() { return BasicIterator{_leftmost}; }

	/** Return a const iterator pointing to the first element of the tree. */
	ConstIterator begin() const { return ConstIterator{_leftmost}; }

	/** Return an iterator pointing to the last element in the tree. */
	BasicIterator end() { return BasicIterator{nullptr}; }

	/** Return a const iterator pointing to the last element of the tree. */
	ConstIterator end() const { return ConstIterator{nullptr}; }

	/**
	 * Returns an iterator to the first item thas is not less than @p key
	 * in the tree (or end() if this cannot be found).
	 */
	BasicIterator lowerBound(const Key &key) {
		Node *it = _root;
		Node *res = nullptr;
		while (it) {
			if (!_comp(KeyProj()(it->value), key)) {
				res = it;
				it = it->left;
			} else {
				it = it->right;
			}
		}
		return BasicIterator{res};
	}

	/**
	 * Returns a const iterator to the first item thas is not less than @p key
	 * in the tree (or end() if this cannot be found).
	 */
	ConstIterator lowerBound(const Key &key) const {
		Node *it = _root;
		Node *res = nullptr;
		while (it) {
			if (!_comp(KeyProj()(it->value), key)) {
				res = it;
				it = it->left;
			} else {
				it = it->right;
			}
		}
		return ConstIterator{res};
	}

	/**
	 * Returns an iterator to the first item bigger than @p key
	 * in the tree (or end() if this cannot be found).
	 */
	BasicIterator upperBound(const Key &key) {
		Node *it = _root;
		Node *res = nullptr;
		while (it) {
			if (!_comp(key, KeyProj()(it->value))) {
				it = it->right;
			} else {
				res = it;
				it = it->left;
			}
		}
		return BasicIterator{res};
	}

	/**
	 * Return a const iterator to the first item bigger than @p key
	 * in the tree (or end() if this cannot be found).
	 */
	ConstIterator upperBound(const Key &key) const {
		Node *it = _root;
		Node *res = nullptr;
		while (it) {
			if (!_comp(key, KeyProj()(it->value))) {
				it = it->right;
			} else {
				res = it;
				it = it->left;
			}
		}
		return ConstIterator{res};
	}

	/** Erases the item in the tree pointed by @p it .*/
	BasicIterator erase(BasicIterator it) {
		Node *const z = it._current;
		Node *y = z;
		assert(y);
		const auto ret = it++;
		Color y_prev_color = y->color;
		Node *x;
		Node *xp = nullptr;
		if (!y->left) {
			x = y->right;
			xp = y->parent; // since x is put in y's place by the next call
			transplant(y, y->right);
		} else if (!y->right) {
			x = y->left;
			xp = y->parent;
			transplant(y, y->left);
		} else {
			y = leftmost(z->right);
			y_prev_color = y->color;
			x = y->right;
			if (y != z->right) {
				xp = y->parent;
				transplant(y, y->right);
				y->right = z->right;
				y->right->parent = y;
			} else {
				xp = y;
			}
			transplant(z, y);
			y->left = z->left;
			y->left->parent = y;
			y->color = z->color;
		}
		if (y_prev_color == Color::kBlack)
			fixDelete(x, xp);
		delete z;
		--_size;
		return ret;
	}

	/**
	 * Erase the elements from @p first to @p last and return an iterator pointing to the next element in the tree.
	 * @note
	 * If [first, last) is not a valid range for the tree, the behaviour is undefined.
	 */
	BasicIterator erase(BasicIterator first, BasicIterator last) {
		while (first != last)
			erase(first++);
		return last;
	}

	/** Inserts a value @p val into the tree returning an iterator for the inserted value. */
	BasicIterator insert(const ValueType &val) {
		return internalInsert(&_root, val);
	}

	/**
	 * Inserts the element @p val starting from @p start instead of the tree's root.
	 * This operations is meant for efficient insertions after a call to lowerBound.
	 * For example:
	 * @code
	 * auto it = tree.lowerBound(value);
	 * if (it == tree.end() || *it != value)
	 *    tree.insert(it, value);
	 * @endcode
	 * inserts 'value' if its not contained in 'tree' (assumes that key and value type are the same)
	 * @note
	 * If @p start is not the lower bound, or it's not end(),
	 * the resulting tree could be unsorted.
	 */
	BasicIterator insert(BasicIterator start, const ValueType &val) {
		if (start == end())
			return insert(val);
		return internalInsert(&start._current, val);
	}

	/** Returns the number of values in the tree. */
	size_t size() const { return _size; }

	/**
	 * Returns true if the tree is empty.
	 * Shorthand for:
	 * @code
	 * tree.size() == 0
	 * @endcode
	 */
	bool isEmpty() const { return _size == 0; }

	~RBTree() { clear(); }

private:
	KeyComp _comp;
	Node *_root = nullptr;
	Node *_leftmost = nullptr;
	size_t _size = 0;

	BasicIterator internalInsert(Node **starting_point, const ValueType &val) {
		auto it = starting_point;
		Node *parent = nullptr;
		while (*it) {
			parent = *it;
			if (_comp(KeyProj()((*it)->value), KeyProj()(val))) {
				it = &(*it)->right;
			} else {
				it = &(*it)->left;
			}
		}
		*it = new Node{
			parent,
			nullptr,
			nullptr,
			Color::kRed,
			val,
		};
		if (!_leftmost || (parent == _leftmost && _leftmost->left == *it))
			_leftmost = *it;
		++_size;
		auto ret = BasicIterator{*it};
		fixInsert(*it);
		return ret;
	}

	static Node *leftmost(Node *n) {
		while (n->left)
			n = n->left;
		return n;
	}

	// neither rotate changes _leftmost
	void rotateLeft(Node *t) {
		assert(t);
		assert(t->right);
		Node *r = t->right;
		Node *p = t->parent;
		// set r->left as t->right
		t->right = r->left;
		if (r->left)
			r->left->parent = t;
		// set the parent of r
		r->parent = p;
		if (!p)
			_root = r;
		else if (p->right == t)
			p->right = r;
		else
			p->left = r;
		// set the parent of t
		t->parent = r;
		r->left = t;
	}

	void rotateRight(Node *t) {
		assert(t);
		assert(t->left);
		Node *l = t->left;
		Node *p = t->parent;
		assert(p != l);
		// set l->right as t->left
		t->left = l->right;
		if (l->right)
			l->right->parent = t;
		// set the parent of l
		l->parent = p;
		if (!p)
			_root = l;
		else if (p->right == t)
			p->right = l;
		else
			p->left = l;
		// set the parent of t
		l->right = t;
		t->parent = l;
	}

	void transplant(Node *t, Node *u) {
		if (!t->parent) {
			_root = u;
		} else if (t == t->parent->left) {
			t->parent->left = u;
		} else {
			t->parent->right = u;
		}
		if (u) {
			u->parent = t->parent;
		}
		if (t == _leftmost)
			_leftmost = u ? leftmost(u) : t->parent;
	}

	void fixInsert(Node *t) {
		while (t->parent && t->parent->color == Color::kRed) {
			Node *p = t->parent;
			Node *g = p->parent;
			// cannot be null since p is not _root
			assert(g);
			if (p == g->left) {
				Node *const u = g->right;
				if (u && u->color == Color::kRed) {
					p->color = u->color = Color::kBlack;
					g->color = Color::kRed;
					t = g;
				} else {
					if (t == p->right) {
						rotateLeft(p);
						t = p;
						p = t->parent;
					}
					p->color = Color::kBlack;
					// g is not changed by the previous rotation
					g->color = Color::kRed;
					rotateRight(g);
				}
			} else {
				Node *const u = g->left;
				if (u && u->color == Color::kRed) {
					p->color = u->color = Color::kBlack;
					g->color = Color::kRed;
					t = g;
				} else {
					if (t == p->left) {
						rotateRight(p);
						t = p;
						p = t->parent;
					}
					p->color = Color::kBlack;
					g->color = Color::kRed;
					rotateLeft(g);
				}
			}
		}
		_root->color = Color::kBlack;
	}

	void fixDelete(Node *t, Node *p) {
		while (t != _root && (!t || t->color == Color::kBlack)) {
			if (t == p->left) {
				// since the deleted node was black and t is in its
				// place, it has to have a sibling
				Node *b = p->right;
				assert(b);
				if (b->color == Color::kRed) {
					b->color = Color::kBlack;
					p->color = Color::kRed;
					rotateLeft(p);
					p = b->left;
					// since b was red, it had two black children,
					// the right one now being p->right
					b = p->right;
				}
				if ((!b->left || b->left->color == Color::kBlack) &&
					(!b->right || b->right->color == Color::kBlack)) {
					b->color = Color::kRed;
					t = p;
					p = p->parent;
				} else {
					if (!b->right || b->right->color == Color::kBlack) {
						// b->left exists, since b->right is black, and it is red
						// because of one of the above checks
						b->left->color = Color::kBlack;
						b->color = Color::kRed;
						rotateRight(b);
						b = p->right; // p is not changed by the rotation
					}
					b->color = p->color;
					p->color = Color::kBlack;
					if (b->right)
						b->right->color = Color::kBlack;
					rotateLeft(p);
					break;
				}
			} else { // same as above case but left and right swapped
				Node *b = p->left;
				assert(b);
				if (b->color == Color::kRed) {
					b->color = Color::kBlack;
					p->color = Color::kRed;
					rotateRight(p);
					p = b->right;
					b = p->left;
				}
				if ((!b->left || b->left->color == Color::kBlack) &&
					(!b->right || b->right->color == Color::kBlack)) {
					b->color = Color::kRed;
					t = p;
					p = p->parent;
				} else {
					if (!b->left || b->left->color == Color::kBlack) {
						b->right->color = Color::kBlack;
						b->color = Color::kRed;
						rotateLeft(b);
						b = p->left;
					}
					b->color = p->color;
					p->color = Color::kBlack;
					if (b->left)
						b->left->color = Color::kBlack;
					rotateRight(p);
					break;
				}
			}
		}
		if (t)
			t->color = Color::kBlack;
	}
};

/** @} */

} // End of namespace Common

#endif