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

#include <string.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;
  }

  bool side = key > parent->key;
  parent->children[side] = child;

  if (child != NULL) {
    child->parent = parent;
    vieter_tree_node_set(child, vieter_tree_node_right, side);
  }
}

void vieter_tree_node_set_children(vieter_tree_node *parent,
                                   vieter_tree_node **children) {
  memcpy(parent->children, children, 2 * sizeof(vieter_tree_node *));

  for (uint8_t i = 0; i < 2; i++) {
    if (parent->children[i] != NULL) {
      parent->children[i]->parent = parent;
      vieter_tree_node_set(parent->children[i], vieter_tree_node_right, i);
    }
  }
}

void vieter_tree_node_set_child(vieter_tree_node *parent,
                                vieter_tree_node *child, bool right) {
  parent->children[right] = child;

  if (child != NULL) {
    child->parent = parent;
    vieter_tree_node_set(child, vieter_tree_node_right, right);
  }
}

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;
    node = node->children[key > parent->key];
  }

  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;
    }

    node = node->children[key > node->key];
  }

  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_ptr, uint64_t key) {
  vieter_tree_node *target;
  vieter_tree_error res = vieter_tree_node_search_node(&target, *root_ptr, key);

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

  *out = target->data;
  vieter_tree_node *possible_new_root;

  if (target->children[0] != NULL && target->children[1] != NULL) {
    vieter_tree_node *replacement = target->children[1];

    while (replacement->children[0] != NULL) {
      replacement = replacement->children[0];
    }

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

    possible_new_root = vieter_tree_node_remove_balanced(replacement);
  } else {
    possible_new_root = vieter_tree_node_remove_balanced(target);
  }

  if (possible_new_root == NULL) {
    *root_ptr = NULL;
  } else if (possible_new_root->parent == NULL) {
    *root_ptr = possible_new_root;
  }

  return vieter_tree_ok;
}

void vieter_tree_node_set(vieter_tree_node *node, vieter_tree_node_flag flag,
                          bool set) {
  if (set) {
    node->flags |= flag;
  } else {
    node->flags &= ~flag;
  }
}

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

vieter_tree_node *vieter_tree_node_next(vieter_tree_node *node) {
  if (node->children[1] != NULL) {
    node = node->children[1];

    while (node->children[0] != NULL) {
      node = node->children[0];
    }

    return node;
  }

  while (node->parent != NULL &&
         vieter_tree_node_get(node, vieter_tree_node_right)) {
    node = node->parent;
  }

  return node->parent;
}