#include "vieter_tree_balancing.h"
#include "vieter_tree_internal.h"

vieter_tree *vieter_tree_init() { return calloc(1, sizeof(vieter_tree)); }

uint64_t vieter_tree_size(vieter_tree *tree) { return tree->size; }

vieter_tree_error vieter_tree_insert(vieter_tree *tree, uint64_t key,
                                     void *data) {
  if (tree->size == 0) {
    tree->root = vieter_tree_node_init();
    tree->root->key = key;
    tree->root->data = data;
    vieter_tree_node_set(tree->root, vieter_tree_node_black, true);
    tree->size = 1;

    return vieter_tree_ok;
  }

  vieter_tree_error res = vieter_tree_node_insert(tree->root, key, data);

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

  tree->size++;

  return vieter_tree_ok;
}

vieter_tree_error vieter_tree_search(void **out, vieter_tree *tree,
                                     uint64_t key) {
  if (tree->size == 0) {
    return vieter_tree_not_present;
  }

  return vieter_tree_node_search(out, tree->root, key);
}

vieter_tree_error vieter_tree_remove(void **out, vieter_tree *tree,
                                     uint64_t key) {
  if (tree->size == 0) {
    return vieter_tree_not_present;
  }

  vieter_tree_error res = vieter_tree_node_remove(out, &tree->root, key);

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

  tree->size--;

  return vieter_tree_ok;
}

void vieter_tree_free(vieter_tree *tree) {
  if (tree->size == 0) {
    goto end;
  }

  uint64_t capacity = 4;
  uint64_t size = 1;
  vieter_tree_node **stack = malloc(capacity * sizeof(vieter_tree_node *));
  stack[0] = tree->root;

  vieter_tree_node *node;

  while (size > 0) {
    node = stack[size - 1];
    size--;

    if (size + 2 > capacity) {
      capacity *= 2;
      stack = realloc(stack, capacity * sizeof(vieter_tree_node *));
    }

    if (node->children[0] != NULL) {
      stack[size] = node->children[0];
      size++;
    }

    if (node->children[1] != NULL) {
      stack[size] = node->children[1];
      size++;
    }

    vieter_tree_node_free(node);
  }

  free(stack);

end:
  free(tree);
}

bool vieter_tree_validate(vieter_tree *tree) {
  if (tree->size == 0) {
    return true;
  }

  // DFS to get expected black nodes
  uint64_t expected_black_nodes = 0;
  vieter_tree_node *node = tree->root;

  while (node != NULL) {
    if (vieter_tree_node_get(node, vieter_tree_node_black)) {
      expected_black_nodes++;
    }

    node = node->children[0];
  }

  return vieter_tree_node_get(tree->root, vieter_tree_node_black) &&
         vieter_tree_node_validate(tree->root, 0, expected_black_nodes);
}