#include <iostream>
#include <stdexcept>
#include <utility>
#include <memory>
#include <string>
#define RED true
#define BLACK false

class RBTree;
class Node {
  friend class RBTree;
  friend void left_rotate(RBTree &tree, std::shared_ptr<Node> &subtree_root);
  friend void right_rotate(RBTree &tree, std::shared_ptr<Node> &subtree_root);
 public:
  Node() {}
  Node(const std::pair<std::string, std::string> &values, bool color): values(values), color(color) {}

  std::string get_key() {
    return values.first;
  }

  std::string &get_value() {
    return values.second;
  }

 private:
  std::pair<std::string, std::string> values;
  std::shared_ptr<Node> left = nullptr;
  std::shared_ptr<Node> right = nullptr;
  std::shared_ptr<Node> parent = nullptr;
  bool color;
};


class RBTree {
  friend void left_rotate(RBTree &tree, std::shared_ptr<Node> &subtree_root);
  friend void right_rotate(RBTree &tree, std::shared_ptr<Node> &subtree_root);
 public:
  RBTree() {}
  std::string &operator[](const std::string &rhs);
  void remove(const std::string &key);

  size_t size() const {
    return current_size;
  }

  std::string at(const std::string &key);

 private:
  std::shared_ptr<Node> root_node;
  size_t current_size = 0;
  static std::shared_ptr<Node> traverse(const std::string &value, std::shared_ptr<Node> &current);
  static std::shared_ptr<Node> insert(const std::string &key,
                                      std::shared_ptr<Node> &insertion_point,
                                      RBTree &tree);
};


void left_rotate(RBTree &tree, std::shared_ptr<Node> &subtree_root) {
  std::shared_ptr<Node> new_subtree_root = subtree_root->right;
  subtree_root->right = new_subtree_root->left;

  if (new_subtree_root->left != nullptr)
    new_subtree_root->left->parent = subtree_root;

  new_subtree_root->parent = subtree_root->parent;

  if (subtree_root->parent == nullptr)
    tree.root_node = new_subtree_root;
  else if (subtree_root == subtree_root->parent->left)
    subtree_root->parent->left = new_subtree_root;
  else
    subtree_root->parent->right = new_subtree_root;

  new_subtree_root->left = subtree_root;
  subtree_root->parent = new_subtree_root;
}


void right_rotate(RBTree &tree, std::shared_ptr<Node> &subtree_root) {
  std::shared_ptr<Node> new_subtree_root = subtree_root->left;
  subtree_root->left = new_subtree_root->right;

  if (new_subtree_root->right != nullptr)
    new_subtree_root->right->parent = new_subtree_root;

  new_subtree_root->parent = subtree_root->parent;

  if (subtree_root->parent == nullptr)
    tree.root_node = new_subtree_root;
  else if (subtree_root == subtree_root->parent->left)
    subtree_root->parent->left = new_subtree_root;
  else
    subtree_root->parent->right = new_subtree_root;

  new_subtree_root->right = subtree_root;
  subtree_root->parent = new_subtree_root;
}


std::shared_ptr<Node> RBTree::insert(const std::string &key,
                                     std::shared_ptr<Node> &insertion_point,
                                     RBTree &tree) {
  Node new_node({key, ""}, RED);

  if (tree.root_node == nullptr) {
    new_node.color = BLACK;
    tree.root_node = std::make_shared<Node>(new_node);
    return tree.root_node;
  } else {
    if (key < insertion_point->get_key()) {
      insertion_point->left = std::make_shared<Node>(new_node);
      insertion_point->left->parent = insertion_point;

      if (insertion_point->color == RED)
        insertion_point->left->color = BLACK;

      if (insertion_point->left->color == BLACK && insertion_point->right == nullptr)
        right_rotate(tree, insertion_point);

      return insertion_point->left;
    } else {
      insertion_point->right = std::make_shared<Node>(new_node);
      insertion_point->right->parent = insertion_point;

      if (insertion_point->color == RED)
        insertion_point->right->color = BLACK;

      if (insertion_point->right->color == BLACK && insertion_point->right == nullptr)
        left_rotate(tree, insertion_point);

      return insertion_point->right;
    }
  }
}


std::shared_ptr<Node> RBTree::traverse(const std::string &value, std::shared_ptr<Node> &current) {
  if (value < current->get_key()) {
    if (current->left == nullptr)
      return current;

    return traverse(value, current->left);
  }

  if (value > current->get_key()) {
    if (current->right == nullptr)
      return current;

    return traverse(value, current->right);
  }

  return current;
}


void RBTree::remove(const std::string &key) {
  if (root_node == nullptr)
    throw std::out_of_range("the tree is empty");

  std::shared_ptr<Node> to_delete = traverse(key, root_node);
  bool to_the_left = false;
  bool to_the_right = false;

  if (to_delete->get_key() != key)
    throw std::out_of_range(key + " not found");

  if (to_delete == to_delete->parent->left) {
    to_delete->parent->left = nullptr;
    to_the_left = true;
  } else {
    to_delete->parent->right = nullptr;
    to_the_right = true;
  }

  std::shared_ptr<Node> deleted_node_parent = to_delete->parent;
  to_delete->parent = nullptr;

  if (to_the_left) {
    if (deleted_node_parent->color == BLACK && deleted_node_parent->right == nullptr)
      right_rotate(*this, deleted_node_parent);
  } else if (to_the_right) {
    if (deleted_node_parent->color == BLACK && deleted_node_parent->left == nullptr)
      left_rotate(*this, deleted_node_parent);
  }
}


std::string &RBTree::operator[](const std::string &rhs) {
  if (root_node == nullptr)
    return insert(rhs, root_node, *this)->get_value();

  std::shared_ptr<Node> current = traverse(rhs, root_node);

  if (current->get_key() != rhs)
    return insert(rhs, current, *this)->get_value();

  return current->get_value();
}


std::string RBTree::at(const std::string &key) {
  std::shared_ptr<Node> match = traverse(key, root_node);

  if (match->get_key() == key)
    return match->get_value();

  throw std::out_of_range(key + " not found");
}


int main() {
  RBTree test;

  test["red"] = "black";
  test["black"] = "red";
  test["green"] = "blue";
  test["mauve"] = "mauve";
  test["abc"] = "def";
  test["PanAm"] = "Boeing 747";
  test["twitter"] = "x";
  test.remove("twitter");

  try {
    std::cout << test.at("twitter") << std::endl;
  } catch (std::out_of_range err) {
    std::cout << err.what() << std::endl;
  }

  std::cout << test["mauve"] << std::endl;
  std::cout << test["PanAm"] << std::endl;

  return 0;
}