#include "kpl.h"

extern inline void kpl_tree_init(kpl_tree *tree, kpl_tree_node_compare_fn node_compare_fn,
    kpl_tree_node_free_fn node_free_fn);

static void kpl_tree_free_helper(kpl_tree_node *node, kpl_tree_node_free_fn node_free_fn) {
    if (node->tree_left)
        kpl_tree_free_helper(node->tree_left, node_free_fn);
    if (node->tree_right)
        kpl_tree_free_helper(node->tree_right, node_free_fn);
    node_free_fn(node);
}

void kpl_tree_free(kpl_tree *tree) {
    kpl_tree_free_helper(tree->root_node, tree->node_free_fn);
}

static void *kpl_tree_find_helper(kpl_tree_node *root_node, const void *find_node,
    kpl_tree_node_compare_fn node_compare_fn) {
    if (!root_node)
        return NULL;
    int32_t compare = node_compare_fn(root_node, find_node);
    if (compare < 0)
        return kpl_tree_find_helper(root_node->tree_left, find_node, node_compare_fn);
    if (compare > 0)
        return kpl_tree_find_helper(root_node->tree_right, find_node, node_compare_fn);
    return root_node;
}

void *kpl_tree_find(kpl_tree *tree, const void *find_node) {
    return kpl_tree_find_helper(tree->root_node, find_node, tree->node_compare_fn);
}

static int32_t kpl_tree_node_height_get(const kpl_tree_node *node) {
    return node ? node->tree_height : 0;
}

static void kpl_tree_node_height_update(kpl_tree_node *node) {
    int32_t left_height = kpl_tree_node_height_get(node->tree_left);
    int32_t right_height = kpl_tree_node_height_get(node->tree_right);
    node->tree_height = MAX(left_height, right_height) + 1;
}

static kpl_tree_node *kpl_tree_node_rotate_right(kpl_tree_node *right_node) {
   kpl_tree_node *right_left = right_node->tree_left, *right_left_right = right_left->tree_right;
   right_left->tree_right = right_node;
   right_node->tree_left = right_left_right;
   kpl_tree_node_height_update(right_node);
   kpl_tree_node_height_update(right_left);
   return right_left;
}

static kpl_tree_node *kpl_tree_node_rotate_left(kpl_tree_node *left_node) {
    kpl_tree_node *left_right = left_node->tree_right, *left_right_left = left_right->tree_left;
    left_right->tree_left = left_node;
    left_node->tree_right = left_right_left;
    kpl_tree_node_height_update(left_node);
    kpl_tree_node_height_update(left_right);
    return left_right;
}

static kpl_tree_node *kpl_tree_insert_helper(kpl_tree_node *restrict root_node,
    kpl_tree_node *restrict insert_node, kpl_tree_node_compare_fn node_compare_fn) {
    if (!root_node)
        return insert_node;
    int32_t compare = node_compare_fn(root_node, insert_node);
    if (compare < 0)
        root_node->tree_left = kpl_tree_insert_helper(root_node->tree_left, insert_node, node_compare_fn);
    else if (compare > 0)
        root_node->tree_right = kpl_tree_insert_helper(root_node->tree_right, insert_node, node_compare_fn);
    else {
        insert_node->tree_duplicates = root_node->tree_duplicates;
        root_node->tree_duplicates = insert_node;
        return root_node;
    }
    kpl_tree_node_height_update(root_node);
    int32_t balance = kpl_tree_node_height_get(root_node->tree_left) -
        kpl_tree_node_height_get(root_node->tree_right);
    if (balance > 1 && compare < 0)
        return kpl_tree_node_rotate_right(root_node);
    if (balance < -1 && compare > 0)
        return kpl_tree_node_rotate_left(root_node);
    if (balance > 1 && compare > 0) {
        root_node->tree_left = kpl_tree_node_rotate_left(root_node->tree_left);
        return kpl_tree_node_rotate_right(root_node);
    }
    if (balance < -1 && compare < 0) {
        root_node->tree_right = kpl_tree_node_rotate_right(root_node->tree_right);
        return kpl_tree_node_rotate_left(root_node);
    }
    return root_node;
}

void kpl_tree_insert(kpl_tree *tree, void *insert_node) {
    tree->root_node = kpl_tree_insert_helper(tree->root_node, insert_node, tree->node_compare_fn);
}