// ==== Intructions ====
// • Write a program that will have pre defined input as numbers (1,2,3,4,5)
// • Then program will first create a BST of the input and print that out
// • Then the program will balance the BST and print out the AVL tree
// (pre-order)
//
// Steps
//   • Create BST (0.5-point)
//   • Balance the tree (i.e. AVL tree) (0.5-point)
// Output:
//   • Print BST Pre-order (1-point)
//   • Print Balanced-AVL Pre-order (1-point)

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

int max(int a, int b) {
    int max = (a > b) ? a : b;
    return max;
}

typedef struct BstNode {
    int value;
    int height;
    int balance;
    struct BstNode *left;
    struct BstNode *right;
} BstNode;

BstNode *bst_new(int value) {
    BstNode *bst = (BstNode *)malloc(sizeof(BstNode));
    bst->value = value;
    bst->height = 0;
    bst->balance = 0;
    bst->left = NULL;
    bst->right = NULL;

    return bst;
}
void bst_print_preorder(BstNode *node) {
    if (node == NULL) {
        return;
    }

    printf("%d, ", node->value);
    bst_print_preorder(node->left);
    bst_print_preorder(node->right);
}

void bst_dbg(BstNode *node) {
    bst_print_preorder(node);
    printf("BstNode {\n  value: %d,\n  height: %d,\n  balance: %d\n}\n",
           node->value, node->height, node->balance);
}

BstNode *bst_search(BstNode *node, int value) {
    if (node == NULL) {
        return NULL;
    }

    if (node->value == value) {
        return node;
    }

    if (value < node->value) {
        return bst_search(node->left, value);
    } else {
        return bst_search(node->right, value);
    }
}

bool bst_contains(BstNode *node, int value) {
    return (bst_search(node, value) != NULL);
}

int bst_height(BstNode *node) {
    if (node == NULL) {
        return -1;
    }

    int left_height = bst_height(node->left);
    int right_height = bst_height(node->right);

    return max(left_height, right_height) + 1;
}

void bst_update(BstNode *node) {
    int left_height = bst_height(node->left);
    int right_height = bst_height(node->right);
    node->height = max(left_height, right_height) + 1;
    node->balance = left_height - right_height;
}

BstNode *avl_right_rotate(BstNode *node) {
    BstNode *left_child = node->left;
    node->left = left_child->right;
    left_child->right = node;
    bst_update(node);
    bst_update(left_child);
    return node;
}

BstNode *avl_left_rotate(BstNode *node) {
    BstNode *right_child = node->right;
    node->right = right_child->left;
    right_child->left = node;
    bst_update(node);
    bst_update(right_child);
    return right_child;
}

/**
 * @brief Balances a single node according to the AVL algorithm
 *
 * @param node The node to balance
 * @return The balanced node
 */
BstNode *avl_balance_node(BstNode *node) {
    if (node->balance == 2) {
        if (node->left->balance >= 0) {
            node = avl_right_rotate(node);
        } else {
            node = avl_left_rotate(node->left);
            node = avl_right_rotate(node);
        }
    } else if (node->balance == -2) {
        if (node->right->balance <= 0) {
            node = avl_left_rotate(node);
        } else {
            node = avl_right_rotate(node->right);
            node = avl_left_rotate(node);
        }
    }

    return node;
}

/**
 * @brief Insert a value into a Binary Search Tree in place
 *
 * @param node The given binary search tree to insert into
 * @param value The value to insert
 * @param balance Whether to balance the insertion with AVL
 */
void bst_insert(BstNode **node, int value, bool balance) {
    if (*node == NULL)
        *node = bst_new(value);
    else if ((*node)->value == value)
        // No duplicates allowed in our implementation
        return;
    else if ((*node)->value > value)
        bst_insert(&(*node)->left, value, balance);
    else
        bst_insert(&(*node)->right, value, balance);

    bst_update(*node);
    if (balance)
        *node = avl_balance_node(*node);
}

int main(int argc, char *argv[]) {
    printf("Price Hiller, zfp106 (Spring 2024)\n");
    printf("==================================\n");
    printf("Unbalanced BST Pre-Order:    ");
    BstNode *bst = bst_new(1);
    int input[] = {2, 3, 4, 5};
    for (int i = 0; i < sizeof(input)/sizeof(input[0]); i++) {
        bst_insert(&bst, input[i], false);
    }
    bst_print_preorder(bst);

    printf("\nBalanced AVL Tree Pre-Order: ");
    BstNode *avl = bst_new(1);
    for (int i = 0; i < sizeof(input)/sizeof(input[0]); i++) {
        bst_insert(&avl, input[i], true);
    }
    bst_print_preorder(avl);
    printf("\n");
    return EXIT_SUCCESS;
}