#include "balancing.h"

vieter_tree_node *vieter_tree_node_balance(vieter_tree_node *node) {
  vieter_tree_node *parent = node->parent;
  vieter_tree_node *grand_parent = parent->parent;
  vieter_tree_node *grand_grand_parent = grand_parent->parent;
  vieter_tree_node *root, *left, *right;
  vieter_tree_node *children[4];

  if (node->key < parent->key && parent->key < grand_parent->key) {
    root = parent;
    left = node;
    right = grand_parent;

    children[0] = node->left;
    children[1] = node->right;
    children[2] = parent->right;
    children[3] = grand_parent->right;
  } else if (node->key < parent->key && parent->key > grand_parent->key) {
    root = node;
    left = grand_parent;
    right = parent;

    children[0] = grand_parent->left;
    children[1] = node->left;
    children[2] = node->right;
    children[3] = parent->right;
  } else if (node->key > parent->key && parent->key < grand_parent->key) {
    root = node;
    left = parent;
    right = grand_parent;

    children[0] = parent->left;
    children[1] = node->left;
    children[2] = node->right;
    children[3] = grand_parent->right;
  } else {
    root = parent;
    left = grand_parent;
    right = node;

    children[0] = grand_parent->left;
    children[1] = parent->left;
    children[2] = node->left;
    children[3] = node->right;
  }

  vieter_tree_node_set(left, vieter_node_black);
  vieter_tree_node_set(right, vieter_node_black);
  vieter_tree_node_unset(root, vieter_node_black);

  left->left = children[0];
  left->right = children[1];

  if (children[0] != NULL)
    children[0]->parent = left;

  if (children[1] != NULL)
    children[1]->parent = left;

  right->left = children[2];
  right->right = children[3];

  if (children[2] != NULL)
    children[2]->parent = right;

  if (children[3] != NULL)
    children[3]->parent = right;

  root->left = left;
  root->right = right;
  left->parent = root;
  right->parent = root;

  if (grand_grand_parent != NULL) {
    vieter_tree_node_add_child(grand_grand_parent, root->key, root);
  } else {
    root->parent = NULL;
  }

  return root;
}

void vieter_tree_node_balance_after_insert(vieter_tree_node *node) {
  while (!vieter_tree_node_get(node, vieter_node_black) &&
         node->parent != NULL &&
         !vieter_tree_node_get(node->parent, vieter_node_black)) {
    node = vieter_tree_node_balance(node);
  }

  // The root must always be black
  if (node->parent == NULL) {
    vieter_tree_node_set(node, vieter_node_black);
  }
}