#include "vieter_tree_node.h"
#include "vieter_tree_balancing.h"

vieter_tree_node *vieter_tree_node_init() {
  return calloc(1, sizeof(vieter_tree_node));
}

void vieter_tree_node_free(vieter_tree_node *node) { free(node); }

void vieter_tree_node_add_child(vieter_tree_node *parent, uint64_t key,
                                vieter_tree_node *child) {
  if (parent == NULL) {
    return;
  }

  if (key < parent->key) {
    parent->left = child;
  } else {
    parent->right = child;
  }

  if (child != NULL) {
    child->parent = parent;
  }
}

void vieter_tree_node_replace_with_child(vieter_tree_node *to_replace,
                                         vieter_tree_node *replacement) {
  to_replace->key = replacement->key;
  to_replace->data = replacement->data;
  to_replace->left = replacement->left;
  to_replace->right = replacement->right;

  if (to_replace->left != NULL)
    to_replace->left->parent = to_replace;

  if (to_replace->right != NULL)
    to_replace->right->parent = to_replace;
}

vieter_tree_error vieter_tree_node_insert(vieter_tree_node *root, uint64_t key,
                                          void *data) {
  vieter_tree_node *node = root;
  vieter_tree_node *parent = root;

  while (node != NULL) {
    if (node->key == key) {
      return vieter_tree_already_present;
    }

    parent = node;

    if (key < parent->key) {
      node = parent->left;
    } else {
      node = parent->right;
    }
  }

  vieter_tree_node *new_node = vieter_tree_node_init();
  new_node->key = key;
  new_node->data = data;

  vieter_tree_node_add_child(parent, key, new_node);
  vieter_tree_node_balance_after_insert(new_node);

  return vieter_tree_ok;
}

vieter_tree_error vieter_tree_node_search_node(vieter_tree_node **out,
                                               vieter_tree_node *root,
                                               uint64_t key) {
  vieter_tree_node *node = root;

  while (node != NULL) {
    if (node->key == key) {
      *out = node;

      return vieter_tree_ok;
    }

    if (key < node->key) {
      node = node->left;
    } else {
      node = node->right;
    }
  }

  return vieter_tree_not_present;
}

vieter_tree_error vieter_tree_node_search(void **out, vieter_tree_node *root,
                                          uint64_t key) {
  vieter_tree_node *target;
  vieter_tree_error res = vieter_tree_node_search_node(&target, root, key);

  if (res != vieter_tree_ok) {
    return res;
  }

  *out = target->data;

  return vieter_tree_ok;
}

vieter_tree_error vieter_tree_node_remove(void **out, vieter_tree_node *root,
                                          uint64_t key) {
  vieter_tree_node *target;
  vieter_tree_error res = vieter_tree_node_search_node(&target, root, key);

  if (res != vieter_tree_ok) {
    return res;
  }

  *out = target->data;

  if (target->left == NULL && target->right == NULL) {
    vieter_tree_node_add_child(target->parent, target->key, NULL);

    vieter_tree_node_free(target);
  } else if ((target->left == NULL) ^ (target->right == NULL)) {
    vieter_tree_node *child =
        target->left != NULL ? target->left : target->right;

    if (target->parent != NULL) {
      vieter_tree_node_add_child(target->parent, child->key, child);
      vieter_tree_node_free(target);
    } else {
      vieter_tree_node_replace_with_child(target, child);
      vieter_tree_node_free(child);
    }

  } else {
    vieter_tree_node *replacement = target->right;

    while (replacement->left != NULL) {
      replacement = replacement->left;
    }

    // We use replacement->key here because the right child can be NULL, so
    // reading its key isn't safe. Using replacement->key however, the child
    // will still get placed into the right location because of how binary
    // trees work.
    vieter_tree_node_add_child(replacement->parent, replacement->key,
                               replacement->right);

    target->key = replacement->key;
    target->data = replacement->data;

    vieter_tree_node_free(replacement);
  }

  return vieter_tree_ok;
}

void vieter_tree_node_set(vieter_tree_node *node, vieter_node_flag flag) {
  node->flags |= flag;
}

void vieter_tree_node_unset(vieter_tree_node *node, vieter_node_flag flag) {
  node->flags &= ~flag;
}

bool vieter_tree_node_get(vieter_tree_node *node, vieter_node_flag flag) {
  return (node->flags & flag) != 0;
}