/*
 * 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_qos_queue.h"

#include <linux/ioprio.h>

#include "ertp_logs.h"
#include "ertp_stats.h"

static enum ertp_qos_queue_ids ertp_qos_queue_prio_to_id(
    unsigned short io_priority_class) {
  switch (io_priority_class) {
    case IOPRIO_CLASS_RT:
      return ERTP_QOS_QUEUE_RT;
    case IOPRIO_CLASS_BE:
      return ERTP_QOS_QUEUE_BE;
    case IOPRIO_CLASS_IDLE:
      return ERTP_QOS_QUEUE_IDLE;
  }
  return ERTP_QOS_QUEUE_BE;
}

static void unblock_and_unref_event(struct ertp_event* event) {
  complete_all(&event->wait_qos);
  ertp_event_unblock(event, 0);
  ertp_event_unref(event);
}

static void deinit_node(struct ertp_heap_node* node) {
  struct ertp_event* event = (struct ertp_event*)(node);
  unblock_and_unref_event(event);
}

void ertp_qos_queue_init(struct ertp_qos_queue_t* queue) {
  size_t i;
  for (i = 0; i < N_ERTP_QOS_QUEUE_IDS; i++) {
    ertp_heap_init(&queue->events[i], deinit_node);
  }
}

void ertp_qos_queue_deinit(struct ertp_qos_queue_t* queue) {
  size_t i;
  for (i = 0; i < N_ERTP_QOS_QUEUE_IDS; i++) {
    ertp_heap_deinit(&queue->events[i]);
  }
}

void ertp_qos_queue_swap_and_reinit(struct ertp_qos_queue_t* queue,
                                    struct ertp_qos_queue_t* swap) {
  size_t i;
  for (i = 0; i < N_ERTP_QOS_QUEUE_IDS; i++) {
    ertp_heap_swap(&queue->events[i], &swap->events[i]);
    ertp_heap_init(&queue->events[i], deinit_node);
  }
}

bool ertp_qos_queue_event_ready(struct ertp_qos_queue_t* queue) {
  size_t i;
  for (i = 0; i < N_ERTP_QOS_QUEUE_IDS; i++) {
    if (!ertp_heap_empty(&queue->events[i])) {
      return true;
    }
  }
  return false;
}

enum ertp_qos_queue_ids ertp_qos_queue_insert(struct ertp_qos_queue_t* queue,
                                              struct ertp_event* event) {
  enum ertp_qos_queue_ids queue_id =
      ertp_qos_queue_prio_to_id(event->qos_data.io_priority_class);

  ertp_event_ref(event);
  ertp_heap_insert(&queue->events[queue_id], &event->node,
                   event->qos_data.current_deadline);
  ertp_stats_report_qos_heap_size(ertp_heap_size(&queue->events[queue_id]),
                                  queue_id);

  return queue_id;
}

bool ertp_qos_queue_remove(struct ertp_qos_queue_t* queue,
                           enum ertp_qos_queue_ids id,
                           struct ertp_event* event) {
  BUG_ON(id >= N_ERTP_QOS_QUEUE_IDS);

  if (ertp_heap_remove(&queue->events[id], &event->node)) {
    ertp_stats_report_qos_heap_size(ertp_heap_size(&queue->events[id]), id);
    ertp_event_unref(event);
    return true;
  }
  return false;
}

struct ertp_event* ertp_qos_queue_get_next(struct ertp_qos_queue_t* queue,
                                           size_t max_size) {
  struct ertp_event* found_event = NULL;
  size_t found_queue;

  struct ertp_event* event = NULL;
  size_t i;

  for (i = 0; i < N_ERTP_QOS_QUEUE_IDS; i++) {
    event = (struct ertp_event*)ertp_heap_find_min(&queue->events[i]);
    if (event == NULL) {
      continue;
    }

    if (ertp_qos_event_data_is_after_deadline(&event->qos_data)) {
      found_event = event;
      found_queue = i;
      break;
    }

    if (found_event == NULL) {
      found_event = event;
      found_queue = i;
    }
  }

  if (found_event) {
    if (found_event->flex_members_size(found_event) > max_size) {
      return ERR_PTR(-ENOBUFS);
    }

    if (unlikely(!ertp_heap_remove(&queue->events[found_queue],
                                   &found_event->node))) {
      BUG();
    }
    return found_event;
  }

  return NULL;
}