#include "vieter_tree.h"
#include "vieter_tree_node.h"

struct vieter_tree {
  uint64_t size;
  vieter_tree_node *root;
};

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_node_black);
    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->left != NULL) {
      stack[size] = node->left;
      size++;
    }

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

    vieter_tree_node_free(node);
  }

  free(stack);

end:
  free(tree);
}