#include <stdlib.h>

#include "lsm/bt_internal.h"

lsm_error lsm_bt_node_init(lsm_bt_node **ptr, const char key, void *data) {
  lsm_bt_node *node = calloc(1, sizeof(lsm_bt_node));

  if (node == NULL) {
    return lsm_error_failed_alloc;
  }

  node->key = key;
  node->data = data;
  *ptr = node;

  return lsm_error_ok;
}

void lsm_bt_node_free(lsm_bt_node *node) { free(node); }

void lsm_bt_node_free_tree(lsm_bt_node *node) {
  if (node->left != NULL) {
    lsm_bt_node_free_tree(node->left);
    lsm_bt_node_free(node->left);
  }

  if (node->right != NULL) {
    lsm_bt_node_free_tree(node->right);
    lsm_bt_node_free(node->right);
  }
}

lsm_error lsm_bt_init(lsm_bt **ptr) {
  lsm_bt *bt = calloc(1, sizeof(lsm_bt));

  if (bt == NULL) {
    return lsm_error_failed_alloc;
  }

  *ptr = bt;

  return lsm_error_ok;
}

void lsm_bt_clear(lsm_bt *bt) {
  if (bt->root != NULL) {
    lsm_bt_node_free_tree(bt->root);
    lsm_bt_node_free(bt->root);

    bt->root = NULL;
    bt->size = 0;
  }
}

void lsm_bt_free(lsm_bt *bt) {
  lsm_bt_clear(bt);
  free(bt);
}

uint64_t lsm_bt_size(const lsm_bt *bt) { return bt->size; }

lsm_error lsm_bt_insert(lsm_bt *bt, char key, void *data) {
  lsm_bt_node **dest = &bt->root;

  // Traverse down the tree until we reach the new point to insert our node
  while ((*dest != NULL) && ((*dest)->key != key)) {
    dest = key < (*dest)->key ? &(*dest)->left : &(*dest)->right;
  }

  if (*dest != NULL) {
    return lsm_error_already_present;
  }

  if (lsm_bt_node_init(dest, key, data) != lsm_error_ok) {
    return lsm_error_failed_alloc;
  }

  bt->size++;

  return lsm_error_ok;
}

lsm_error lsm_bt_search(void **out, const lsm_bt *bt, char key) {
  lsm_bt_node *node = bt->root;

  while ((node != NULL) && (node->key != key)) {
    node = key < node->key ? node->left : node->right;
  }

  if (node == NULL) {
    return lsm_error_not_found;
  }

  if (out != NULL) {
    *out = node->data;
  }

  return lsm_error_ok;
}

lsm_error lsm_bt_remove(void **out, lsm_bt *bt, char key) {
  lsm_bt_node **dest = &bt->root;

  while ((*dest != NULL) && ((*dest)->key != key)) {
    dest = key < (*dest)->key ? &(*dest)->left : &(*dest)->right;
  }

  if (*dest == NULL) {
    return lsm_error_not_found;
  }

  if (out != NULL) {
    *out = (*dest)->data;
  }
  bt->size--;

  if (((*dest)->left != NULL) && ((*dest)->right != NULL)) {
    lsm_bt_node **succ = &(*dest)->right;

    while ((*succ)->left != NULL) {
      succ = &(*succ)->left;
    }

    (*dest)->key = (*succ)->key;
    (*dest)->data = (*succ)->data;

    lsm_bt_node *succ_replacement = (*succ)->right;
    lsm_bt_node_free(*succ);
    *succ = succ_replacement;
  } else {
    lsm_bt_node *replacement =
        (*dest)->left != NULL ? (*dest)->left : (*dest)->right;
    lsm_bt_node_free(*dest);
    *dest = replacement;
  }

  return lsm_error_ok;
}

lsm_error lsm_bt_replace(void **out, lsm_bt *bt, char key, void *data) {
  lsm_bt_node *node = bt->root;

  while ((node != NULL) && (node->key != key)) {
    node = key < node->key ? node->left : node->right;
  }

  if (node == NULL) {
    return lsm_error_not_found;
  }

  if (out != NULL) {
    *out = node->data;
  }

  node->data = data;

  return lsm_error_ok;
}