#include "vieter_tree_balancing.h"
#include <string.h>

bool vieter_tree_node_validate(vieter_tree_node *node,
                               uint64_t passed_black_nodes,
                               uint64_t expected_black_nodes) {
  // Path to NULL children should contain the same amount of black nodes
  if (node == NULL) {
    return passed_black_nodes == expected_black_nodes;
  }

  if (vieter_tree_node_get(node, vieter_tree_node_black)) {
    passed_black_nodes++;
  }

  // Either the node itself is black, or its children are both either NULL or
  // black.
  bool correctly_colored_children =
      vieter_tree_node_get(node, vieter_tree_node_black) ||
      ((node->children[0] == NULL ||
        vieter_tree_node_get(node->children[0], vieter_tree_node_black)) &&
       (node->children[1] == NULL ||
        vieter_tree_node_get(node->children[1], vieter_tree_node_black)));
  bool right_child_flag_set =
      node->parent == NULL ||
      (vieter_tree_node_get(node, vieter_tree_node_right) ==
       (node->parent->children[1] == node));

  return correctly_colored_children && right_child_flag_set &&
         vieter_tree_node_validate(node->children[0], passed_black_nodes,
                                   expected_black_nodes) &&
         vieter_tree_node_validate(node->children[1], passed_black_nodes,
                                   expected_black_nodes);
}

/*
 * This function should be rewritten to use tree rotations instead.
 */
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 *root = grand_parent;
  vieter_tree_node *children[2];
  vieter_tree_node *grand_children[4];

  uint64_t key_root = root->key;
  void *data_root = root->data;

  if (vieter_tree_node_get(node, vieter_tree_node_right)) {
    children[0] = parent;
    children[1] = node;

    if (vieter_tree_node_get(parent, vieter_tree_node_right)) {
      root->key = parent->key;
      root->data = parent->data;

      parent->key = key_root;
      parent->data = data_root;

      grand_children[0] = grand_parent->children[0];
      grand_children[1] = parent->children[0];
      grand_children[2] = node->children[0];
      grand_children[3] = node->children[1];
    } else {
      root->key = node->key;
      root->data = node->data;

      node->key = key_root;
      node->data = data_root;

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

    if (vieter_tree_node_get(parent, vieter_tree_node_right)) {
      root->key = node->key;
      root->data = node->data;

      node->key = key_root;
      node->data = data_root;

      grand_children[0] = grand_parent->children[0];
      grand_children[1] = node->children[0];
      grand_children[2] = node->children[1];
      grand_children[3] = parent->children[1];
    } else {
      root->key = parent->key;
      root->data = parent->data;

      parent->key = key_root;
      parent->data = data_root;

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

  vieter_tree_node_set(
      root, vieter_tree_node_right,
      vieter_tree_node_get(grand_parent, vieter_tree_node_right));

  vieter_tree_node_set(children[0], vieter_tree_node_black, true);
  vieter_tree_node_set(children[1], vieter_tree_node_black, true);

  vieter_tree_node_set(root, vieter_tree_node_black, false);

  vieter_tree_node_set_children(children[0], grand_children);
  vieter_tree_node_set_children(children[1], grand_children + 2);
  vieter_tree_node_set_children(root, children);

  return root;
}

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

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

vieter_tree_node *vieter_tree_node_rotate(vieter_tree_node *old_root,
                                          bool dir) {
  vieter_tree_node *new_root = old_root->children[1 - dir];

  if (old_root->parent != NULL) {
    vieter_tree_node_set_child(
        old_root->parent, new_root,
        vieter_tree_node_get(old_root, vieter_tree_node_right));
  } else {
    new_root->parent = NULL;
  }

  // Right rotation
  if (dir) {
    vieter_tree_node_set_child(old_root, new_root->children[1], false);
    vieter_tree_node_set_child(new_root, old_root, true);
  }
  // Left rotation
  else {
    vieter_tree_node_set_child(old_root, new_root->children[0], true);
    vieter_tree_node_set_child(new_root, old_root, false);
  }

  return new_root;
}

/*
 * This function is currently implemented by very literally following the
 * Wikipedia pseudocode. It's honestly not too bad, and I couldn't be bothered
 * to properly research how to implement red-black removal (how is this so hard
 * to find?)
 *
 * https://en.wikipedia.org/wiki/Red%E2%80%93black_tree#Removal_of_a_black_non-root_leaf
 */
vieter_tree_node *vieter_tree_node_remove_balanced(vieter_tree_node *node) {
  vieter_tree_node *out;

  if (node->parent == NULL) {
    out = node->children[0] != NULL ? node->children[0] : node->children[1];

    vieter_tree_node_free(node);

    // This only happens when the root was the only element in the tree
    if (out == NULL)
      return out;

    vieter_tree_node_set(out, vieter_tree_node_black, true);
    out->parent = NULL;

    return out;
  }

  // A red node can only have 0 or 2 children. The node we receive only has
  // one child at most, so we know if it's red that it doesn't have any
  // children. A black node that has a single (right) child can be replaced by
  // this child, granted it becomes black as well. Either way, the node can be
  // replaced by its right child (even if it's NULL).
  if (!vieter_tree_node_get(node, vieter_tree_node_black) ||
      node->children[1] != NULL) {
    out = node->children[1];

    vieter_tree_node_set_child(
        node->parent, node->children[1],
        vieter_tree_node_get(node, vieter_tree_node_right));

    if (node->children[1] != NULL)
      vieter_tree_node_set(node->children[1], vieter_tree_node_black, true);

    vieter_tree_node_free(node);

    return out;
  }

  // The complicated case is when we want to remove a black leaf
  // https://en.wikipedia.org/wiki/Red%E2%80%93black_tree#Removal_of_a_black_non-root_leaf

  out = node;
  vieter_tree_node *parent = node->parent;
  vieter_tree_node *sibling, *close_nephew, *distant_nephew;
  bool dir = vieter_tree_node_get(node, vieter_tree_node_right);
  parent->children[dir] = NULL;

  vieter_tree_node_free(node);

  goto start;

  do {
    dir = vieter_tree_node_get(node, vieter_tree_node_right);

  start:
    sibling = parent->children[1 - dir];
    distant_nephew = sibling->children[1 - dir];
    close_nephew = sibling->children[dir];

    if (!vieter_tree_node_get(sibling, vieter_tree_node_black))
      goto case3;

    if (distant_nephew != NULL &&
        !vieter_tree_node_get(distant_nephew, vieter_tree_node_black))
      goto case6;

    if (close_nephew != NULL &&
        !vieter_tree_node_get(close_nephew, vieter_tree_node_black))
      goto case5;

    if (!vieter_tree_node_get(parent, vieter_tree_node_black))
      goto case4;

    // Case 2
    vieter_tree_node_set(sibling, vieter_tree_node_black, false);
    node = parent;
    out = node;
  } while ((parent = node->parent) != NULL);

  // Case 1
  return out;

case3:
  out = vieter_tree_node_rotate(parent, dir);
  vieter_tree_node_set(parent, vieter_tree_node_black, false);
  vieter_tree_node_set(sibling, vieter_tree_node_black, true);
  sibling = close_nephew;
  distant_nephew = sibling->children[1 - dir];

  if (distant_nephew != NULL &&
      !vieter_tree_node_get(distant_nephew, vieter_tree_node_black))
    goto case6;

  close_nephew = sibling->children[dir];

  if (close_nephew != NULL &&
      !vieter_tree_node_get(close_nephew, vieter_tree_node_black))
    goto case5;

case4:
  vieter_tree_node_set(sibling, vieter_tree_node_black, false);
  vieter_tree_node_set(parent, vieter_tree_node_black, true);

  return out;

case5:
  vieter_tree_node_rotate(sibling, 1 - dir);
  vieter_tree_node_set(sibling, vieter_tree_node_black, false);
  vieter_tree_node_set(close_nephew, vieter_tree_node_black, true);
  distant_nephew = sibling;
  sibling = close_nephew;

case6:
  out = vieter_tree_node_rotate(parent, dir);
  vieter_tree_node_set(sibling, vieter_tree_node_black,
                       vieter_tree_node_get(parent, vieter_tree_node_black));
  vieter_tree_node_set(parent, vieter_tree_node_black, true);
  vieter_tree_node_set(distant_nephew, vieter_tree_node_black, true);

  return out;
}