#include <assert.h>
#include <stdio.h>
#include <stdlib.h>

#define TREE_ERR 0
#define TREE_OK  1

/* Potentially use a custom allocator on a tiny MCU */
static struct {
    void *(*node_alloc)(size_t);
    void (*node_free)(void *);
} node_mm;

typedef struct node {
    struct node *left, *right;
    int          data;
} node_t;

typedef struct {
    struct node *root;
} tree_t;

node_t *node_new(int data) {
    node_t *node = node_mm.node_alloc(sizeof(node_t));
    if (!node)
        return TREE_ERR;

    node->data = data;
    node->left = node->right = NULL;
    return node;
}

int node_insert(node_t **n, int data) {
    if (!*n) {
        *n = node_new(data);
        return TREE_OK;
    }

    if ((*n)->data == data)
        return TREE_ERR; // Already exists

    return node_insert(data < (*n)->data ? &(*n)->left : &(*n)->right, data);
}

int tree_insert(tree_t *t, int data) {
    return node_insert(&t->root, data);
}

int tree_count(node_t *node) {
    return node ? 1 + tree_count(node->left) + tree_count(node->right) : 0;
}

int node_free(node_t **n) {
    if (!*n)
        return TREE_OK;

    node_free(&(*n)->left);
    node_free(&(*n)->right);

    node_mm.node_free((*n));
    (*n) = NULL;
    return TREE_OK;
}

int node_erase(node_t **n, int data) {
    if (!*n)
        return TREE_ERR; // Not found

    if ((*n)->data < data)
        return node_erase(&(*n)->left, data);
    if ((*n)->data > data)
        return node_erase(&(*n)->right, data);

    node_t *victim = *n;
    if (victim->data != data) {
        return TREE_ERR;
    }

    if (!victim->left) {
        *n = victim->right;
    } else if (!victim->right) {
        *n = victim->right;
    } else {
        /* This is ugly, seatbelt on */
        node_t **succ = &victim->right;
        while ((*succ)->left) succ = &(*succ)->left;

        victim->data = (*succ)->data;
        node_t *tmp = *succ;
        *succ = (*succ)->right;
        free(tmp);
        return TREE_OK;
    }

    node_mm.node_free(victim);
    return TREE_OK;
}

int tree_erase(tree_t *t, int data) {
    return node_erase(&t->root, data);
}

int tree_free(tree_t *t) {
    return node_free(&t->root);
}

int main(int argc, char *argv[]) {
    (void)argc;
    (void)argv;

    printf("Sizeof(node_t): %lu\n", sizeof(node_t));

    node_mm.node_free = free;
    node_mm.node_alloc = malloc;

    /* Check so this works */
    node_t *n = node_new(10);
    assert(n);

    /* Insert some data and count it */
    tree_t t = {};
    assert(tree_insert(&t, 10));
    assert(tree_insert(&t, 12));
    assert(tree_insert(&t, 10) == TREE_ERR);
    assert(tree_count(t.root) == 2);
    assert(t.root);

    /* Free it all */
    tree_free(&t);
    assert(t.root == NULL);
    assert(tree_count(t.root) == 0);

    /* Re-insert some stuff and check */
    assert(tree_insert(&t, 120));
    assert(tree_count(t.root) == 1);

    assert(tree_erase(&t, 120) == TREE_OK); // Exists
    assert(tree_erase(&t, 121) == TREE_ERR); // Does not exist
    assert(tree_count(t.root) == 0);
    assert(tree_insert(&t, 122));
    assert(tree_count(t.root) == 1);
    assert(tree_erase(&t, 122) == TREE_OK);
    assert(tree_count(t.root) == 0);

    printf("All tests passed!\n");
    return EXIT_SUCCESS;
}