/*
 * eset_rtp (ESET Real-time file system protection module)
 * Copyright (C) 1992-2025 ESET, spol. s r.o.
 *
 * This program is free software: you can redistribute it and/or modify
 * it under the terms of the GNU General Public License as published by
 * the Free Software Foundation, either version 3 of the License, or
 * (at your option) any later version.
 *
 * This program is distributed in the hope that it will be useful,
 * but WITHOUT ANY WARRANTY; without even the implied warranty of
 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 * GNU General Public License for more details.
 *
 * You should have received a copy of the GNU General Public License
 * along with this program.  If not, see <https://www.gnu.org/licenses/>.
 *
 * In case of any questions, you can contact us at ESET, spol. s r.o., Einsteinova 24, 851 01 Bratislava, Slovakia.
 */

#include "ertp_heap.h"

#include <linux/compiler.h>
#include <linux/kernel.h>
#include <linux/module.h>

#include "ertp_logs.h"

#define ERTP_HEAP_NODE_KEY_INVALID 0

void ertp_heap_node_init(struct ertp_heap_node *node) {
  node->parent = NULL;
  node->left = NULL;
  node->right = NULL;
  node->key = ERTP_HEAP_NODE_KEY_INVALID;
}

void ertp_heap_init(struct ertp_heap *heap,
                    void (*deinit_node)(struct ertp_heap_node *node)) {
  heap->root = NULL;
  heap->last = NULL;
  heap->size = 0;
  heap->deinit_node = deinit_node;
}

void ertp_heap_deinit(struct ertp_heap *heap) {
  struct ertp_heap_node *node;

  while (!ertp_heap_empty(heap)) {
    node = heap->last;
    if (ertp_heap_remove(heap, node)) {
      heap->deinit_node(node);
    }
  }
}

void ertp_heap_clear(struct ertp_heap *heap) {
  ertp_heap_deinit(heap);
  ertp_heap_init(heap, heap->deinit_node);
}

void ertp_heap_swap(struct ertp_heap *a, struct ertp_heap *b) {
  struct ertp_heap tmp;

  tmp.root = a->root;
  tmp.last = a->last;
  tmp.size = a->size;
  tmp.deinit_node = a->deinit_node;

  a->root = b->root;
  a->last = b->last;
  a->size = b->size;
  a->deinit_node = b->deinit_node;

  b->root = tmp.root;
  b->last = tmp.last;
  b->size = tmp.size;
  b->deinit_node = tmp.deinit_node;
};

static inline bool has_heap_property(struct ertp_heap_node *node) {
  return node->parent->key <= node->key;
}

static inline bool has_children(struct ertp_heap_node *node) {
  return (node->left != NULL || node->right != NULL);
}

static void swap_node_with_parent(struct ertp_heap *heap,
                                  struct ertp_heap_node *node) {
  struct ertp_heap_node *parent = node->parent;
  struct ertp_heap_node *node_left;
  struct ertp_heap_node *node_right;

  if (unlikely(parent == NULL)) {
    BUG();
  }

  if (parent->parent) {
    if (parent->parent->left == parent) {
      parent->parent->left = node;
    } else if (parent->parent->right == parent) {
      parent->parent->right = node;
    } else {
      BUG();
    }
  } else {
    heap->root = node;
  }

  node_left = node->left;
  node_right = node->right;

  node->parent = parent->parent;

  if (node == parent->left) {
    node->left = parent;
    node->right = parent->right;

    if (node->right) {
      node->right->parent = node;
    }

  } else if (node == parent->right) {
    node->left = parent->left;
    node->right = parent;

    if (node->left) {
      node->left->parent = node;
    }
  } else {
    BUG();
  }

  parent->parent = node;
  parent->left = node_left;
  parent->right = node_right;

  if (node_left) {
    node_left->parent = parent;
  }

  if (node_right) {
    node_right->parent = parent;
  }

  if (node == heap->last) {
    ertp_pr_trace(ERTP_LOG_HEAP, "new last: %lu, previous last: %lu",
                  parent->key, heap->last->key);
    heap->last = parent;
  }
}

static struct ertp_heap_node *find_place_for_new_node(
    const struct ertp_heap *heap) {
  struct ertp_heap_node *current_node = heap->last;

  if (unlikely(current_node == NULL)) {
    BUG();
  }

  while (current_node->parent != NULL &&
         current_node == current_node->parent->right) {
    current_node = current_node->parent;
  }

  if (current_node->parent == NULL) {
    while (current_node->left != NULL) {
      current_node = current_node->left;
    }
    return current_node;
  }

  if (current_node->right != NULL) {
    current_node = current_node->parent->right;
    while (current_node->left != NULL) {
      current_node = current_node->left;
    }
    return current_node;
  }

  return current_node->parent;
}

void ertp_heap_insert(struct ertp_heap *heap, struct ertp_heap_node *node,
                      unsigned long key) {
  struct ertp_heap_node *parent_for_new_node;

  if (unlikely(key == ERTP_HEAP_NODE_KEY_INVALID)) {
    ertp_pr_debug("tried to insert node with invalid key into the heap");
    BUG();
  }

  if (unlikely(node->key != ERTP_HEAP_NODE_KEY_INVALID)) {
    ertp_pr_debug("tried to insert node with already set key into the heap");
    BUG();
  }

  heap->size++;
  node->key = key;

  ertp_pr_trace(ERTP_LOG_HEAP, "insert new node: %p - key: %lu", node,
                node->key);

  if (heap->root == NULL) {
    ertp_heap_node_init(node);
    node->key = key;
    heap->root = node;
    heap->last = node;
    return;
  }

  parent_for_new_node = find_place_for_new_node(heap);

  ertp_pr_trace(ERTP_LOG_HEAP, "place found %p - P: %p, L: %p, R: %p, key: %lu",
                parent_for_new_node, parent_for_new_node->parent,
                parent_for_new_node->left, parent_for_new_node->right,
                parent_for_new_node->key);

  if (parent_for_new_node->left == NULL) {
    parent_for_new_node->left = node;
  } else if (parent_for_new_node->right == NULL) {
    parent_for_new_node->right = node;
  } else {
    BUG();
  }

  node->parent = parent_for_new_node;
  node->left = NULL;
  node->right = NULL;

  heap->last = node;

  ertp_pr_trace(ERTP_LOG_HEAP, "place used %p - P: %p, L: %p, R: %p, key: %lu",
                parent_for_new_node, parent_for_new_node->parent,
                parent_for_new_node->left, parent_for_new_node->right,
                parent_for_new_node->key);

  while (node->parent != NULL && !has_heap_property(node)) {
    ertp_pr_trace(ERTP_LOG_HEAP, "heapify N: %lu P: %lu", node->key,
                  node->parent->key);
    swap_node_with_parent(heap, node);
  }

  if (node->parent == NULL) {
    heap->root = node;
  }
}

static struct ertp_heap_node *find_previous_last_node(
    const struct ertp_heap *heap) {
  struct ertp_heap_node *current_node = heap->last;
  while (current_node->parent != NULL &&
         current_node == current_node->parent->left) {
    current_node = current_node->parent;
  }

  if (current_node->parent != NULL) {
    current_node = current_node->parent->left;
  }

  while (current_node->right != NULL) {
    current_node = current_node->right;
  }

  return current_node;
}

static void replace_node(struct ertp_heap_node *node_to_be_replaced,
                         struct ertp_heap_node *new_node) {
  new_node->left = node_to_be_replaced->left;
  new_node->right = node_to_be_replaced->right;
  new_node->parent = node_to_be_replaced->parent;

  if (node_to_be_replaced->left != NULL) {
    new_node->left->parent = new_node;
  }

  if (node_to_be_replaced->right != NULL) {
    new_node->right->parent = new_node;
  }

  if (node_to_be_replaced->parent != NULL) {
    if (node_to_be_replaced->parent->left == node_to_be_replaced) {
      node_to_be_replaced->parent->left = new_node;
    } else if (node_to_be_replaced->parent->right == node_to_be_replaced) {
      node_to_be_replaced->parent->right = new_node;
    } else {
      BUG();
    }
  }

  ertp_heap_node_init(node_to_be_replaced);
}

bool ertp_heap_remove(struct ertp_heap *heap, struct ertp_heap_node *node) {
  struct ertp_heap_node *new_last_node;
  struct ertp_heap_node *current_node;
  struct ertp_heap_node *current_child;

  ertp_pr_trace(ERTP_LOG_HEAP, "remove node: %p - key: %lu", node, node->key);

  if (node->key == ERTP_HEAP_NODE_KEY_INVALID) {
    return false;
  }

  if (node->parent == NULL && node->left == NULL && node->right == NULL) {
    heap->root = NULL;
    heap->last = NULL;
    goto node_removal;
  }

  new_last_node = find_previous_last_node(heap);

  ertp_pr_trace(ERTP_LOG_HEAP,
                "previous last node found: %p - P: %p, L: %p, R: %p, key: %lu",
                new_last_node, new_last_node->parent, new_last_node->left,
                new_last_node->right, new_last_node->key);

  if (heap->last->parent->left == heap->last) {
    heap->last->parent->left = NULL;
  } else if (heap->last->parent->right == heap->last) {
    heap->last->parent->right = NULL;
  } else {
    BUG();
  }

  if (heap->last == node) {
    ertp_pr_trace(ERTP_LOG_HEAP, "new last: %lu, previous last: %lu",
                  new_last_node->key, heap->last->key);
    heap->last = new_last_node;

    goto node_removal;
  }

  replace_node(node, heap->last);
  current_node = heap->last;

  if (heap->root == node) {
    heap->root = current_node;
  }

  if (new_last_node == node) {
    goto node_removal;
  }

  heap->last = new_last_node;

  while (current_node->parent != NULL && !has_heap_property(current_node)) {
    ertp_pr_trace(ERTP_LOG_HEAP, "heapify up N: %lu P: %lu", current_node->key,
                  current_node->parent->key);
    swap_node_with_parent(heap, current_node);
  }

  while (has_children(current_node)) {
    if (current_node->right == NULL) {
      current_child = current_node->left;
    } else if (current_node->left != NULL) {
      if (current_node->right->key < current_node->left->key) {
        current_child = current_node->right;
      } else {
        current_child = current_node->left;
      }
    } else {
      BUG();
    }

    if (!has_heap_property(current_child)) {
      ertp_pr_trace(ERTP_LOG_HEAP, "heapify down N: %lu P: %lu",
                    current_child->key, current_node->key);
      swap_node_with_parent(heap, current_child);
    } else {
      break;
    }
  }

node_removal:
  node->key = ERTP_HEAP_NODE_KEY_INVALID;
  heap->size--;
  return true;
}