#include <lib/seg/mask_data.h>

#include <stdio.h>

// Allocate Mask Data for Label
struct MaskData* create_mask_data(uint16_t label)
{
  struct MaskData *data = (struct MaskData*)malloc(sizeof(struct MaskData));
  data->label = label;
  data->area = 0;
  data->perimeter = 0;
  return data;
}

// Compare mask data labels
bool_t compare_labels(struct MaskData* left, struct MaskData* right)
{
  return left->label < right->label;
}

// Create AVL Mask node
struct AVLNode* create_avl_mask_node(struct MaskData* data)
{
  struct AVLNode* node = (struct AVLNode*)malloc(sizeof(struct AVLNode));
  if (node == NULL) {
    return NULL;
  }
  node->data = data;
  node->compare = (bool_t (*)(void*,void*))&compare_labels;
  node->left = NULL;
  node->right = NULL;
  node->height = 1; // Leaf initially
  return node;
}

// Insert MaskData into the AVL Tree
struct Result insert_mask(struct AVLNode* node, struct MaskData* data)
{
  struct Result result;
  // 1. Standard BST insertion
  if (node == NULL) {
    return (struct Result) {create_avl_mask_node(data), TRUE};
  }

  struct MaskData *node_data = (struct MaskData*)node->data;
  if (node->compare(data, node_data)) {
    result = insert_mask(node->left, data);
    if (!result.success) {
      fprintf(stderr, "Failed to insert!");
      return result;
    }
    node->left = (struct AVLNode*)result.data;
  } else if (node->compare(node->data, data)) {
    result = insert_mask(node->right, data);
    if (!result.success) {
      fprintf(stderr, "Failed to insert!");
      return result;
    }
    node->right = (struct AVLNode*)result.data;
  } else {
    return (struct Result) {node, FALSE};
  }

  // 2. Update height of the ancestor node
  node->height = 1 + max_height(get_height(node->left), get_height(node->right));

  ssize_t balance = get_balance_factor(node);

  // 4. If the node becomes unbalanced

  // LeftLeft
  if ((balance > 1) && node->compare(data, node->left->data)) {
    return (struct Result) {right_rotate(node), TRUE};
  }
  // RightRight
  if ((balance < -1) && node->compare(node->right->data, data)) {
    return (struct Result) {left_rotate(node), TRUE};
  }
  // LeftRight
  if ((balance > 1) && node->compare(node->left->data, data)) {
    return (struct Result) {right_rotate(node), TRUE};
  }
  // RightLeft
  if ((balance < -1) && node->compare(data,node->right->data)) {
    return (struct Result) {left_rotate(node), TRUE};
  }
  return (struct Result) {node, TRUE};
}

// Allocate a label's Mask data in a tree
//  If it already exists, skip the allocation
struct AVLNode* insert_mask_alloc(struct AVLNode* node, uint16_t label)
{
  struct MaskData* data = create_mask_data(label);
  struct Result result = insert_mask(node, data);
  if (!result.success) {
    free(data);
  }
  return (struct AVLNode*)result.data;
}

// Print AVL Node Mask Data Label
void print_label(struct AVLNode* root)
{
  if (root != NULL) {
    print_label(root->left);
    struct MaskData* data = root->data;
    printf("%d: (%zu, %zu) ", data->label, data->area, data->perimeter);
    print_label(root->right);
  }
}

// Increase the label's area
bool_t increase_label_area(struct AVLNode* root, uint16_t label)
{
  if (root == NULL) {
    return FALSE;
  }
  struct MaskData* data = (struct MaskData*)root->data;
  if (data->label == label) {
    data->area++;
  }
  else if (data->label > label) {
    return increase_label_area(root->left, label);
  }
  else if (data->label < label) {
    return increase_label_area(root->right, label);
  }
  return TRUE;
}

// Increase the label's perimeter
bool_t increase_label_perimeter(struct AVLNode* root, uint16_t label)
{
  if (root == NULL) {
    return FALSE;
  }
  struct MaskData* data = (struct MaskData*)root->data;
  if (data->label == label) {
    data->perimeter++;
  }
  else if (data->label > label) {
    return increase_label_perimeter(root->left, label);
  }
  else if (data->label < label) {
    return increase_label_perimeter(root->right, label);
  }
  return TRUE;
}

// Increase the label's area
//  Create an AVL node if it doesn't exist
struct AVLNode* increase_label_area_alloc(struct AVLNode* root, uint16_t label)
{
  struct AVLNode* new_root = root;
  bool_t success = increase_label_area(new_root, label);
  if (success == FALSE) {
    new_root = insert_mask_alloc(new_root, label);
    increase_label_area(new_root, label);
  }
  return new_root;
}

// Increase the label's perimeter
//  Create an AVL node if it doesn't exist
struct AVLNode* increase_label_perimeter_alloc(struct AVLNode* root, uint16_t label)
{
  struct AVLNode* new_root = root;
  bool_t success = increase_label_perimeter(new_root, label);
  if (success == FALSE) {
    new_root = insert_mask_alloc(new_root, label);
    increase_label_perimeter(new_root, label);
  }
  return new_root;
}