#include <libk/list.h>
#include <libk/minmax.h>
#include <libk/std.h>
#include <libk/string.h>
#include <mm/malloc.h>
#include <proc/mail.h>
#include <proc/proc.h>
#include <proc/resource.h>
#include <proc/suspension_q.h>
#include <sys/debug.h>
#include <sys/smp.h>

static void proc_mail_free_saved_buffer(void* udata) { free(udata); }

void proc_cleanup_resource_mail(struct proc_resource* resource, struct reschedule_ctx* rctx) {
  uint64_t fr, fssq;

  struct proc_mail* mail = &resource->u.mail;

  spin_lock(&mail->resource->lock, &fr);

  spin_lock(&mail->send_sq.lock, &fssq);

  while (mail->send_sq.proc_list != NULL) {
    struct list_node_link* node = mail->send_sq.proc_list;
    struct proc_sq_entry* sq_entry = list_entry(node, struct proc_sq_entry, sq_link);
    struct proc* suspended_proc = sq_entry->proc;

    spin_unlock(&mail->send_sq.lock, fssq);
    spin_unlock(&mail->resource->lock, fr);

    proc_sq_resume(suspended_proc, sq_entry, rctx);

    spin_lock(&mail->resource->lock, &fr);
    spin_lock(&mail->send_sq.lock, &fssq);
  }

  spin_unlock(&mail->send_sq.lock, fssq);
  spin_unlock(&mail->resource->lock, fr);
}

void proc_mail_send(struct proc* proc, struct proc_mail* mail, struct reschedule_ctx* rctx,
                    void* data, size_t data_size) {
  uint64_t fr, frsq;

  spin_lock(&mail->resource->lock, &fr);

  /* mail full */
  if (mail->packets_count == PROC_MAIL_MAX) {
    struct mail_saved_buffer* saved_buffer = malloc(sizeof(*saved_buffer));
    saved_buffer->buffer = data;
    saved_buffer->size = data_size;

    proc_sq_suspend(proc, &mail->send_sq, &mail->resource->lock, fr, rctx, saved_buffer,
                    &proc_mail_free_saved_buffer);
    return;
  }

  spin_lock(&mail->recv_sq.lock, &frsq);

  /* if receiver available, hand off directly */
  struct list_node_link* node = mail->recv_sq.proc_list;

  if (node != NULL) {
    struct proc_sq_entry* sq_entry = list_entry(node, struct proc_sq_entry, sq_link);
    struct proc* resumed_proc = sq_entry->proc;

    struct mail_saved_buffer* saved_buffer = sq_entry->udata;

    size_t copy_size = min(data_size, saved_buffer->size);
    memcpy(saved_buffer->buffer, data, copy_size);

    spin_unlock(&mail->recv_sq.lock, frsq);
    spin_unlock(&mail->resource->lock, fr);

    proc_sq_resume(resumed_proc, sq_entry, rctx);
    return;
  }

  spin_unlock(&mail->recv_sq.lock, frsq);

  /* mail is empty and nobody is waiting */
  void* mesg = malloc(data_size);
  if (mesg != NULL) {
    memcpy(mesg, data, data_size);

    struct mail_packet* packet = malloc(sizeof(*packet));

    if (packet == NULL) {
      free(mesg);
    } else {
      memset(packet, 0, sizeof(*packet));
      packet->packet_buffer = mesg;
      packet->packet_size = data_size;

      list_append(mail->packets, &packet->packets_link);
      mail->packets_count++;
    }
  }

  spin_unlock(&mail->resource->lock, fr);
}

void proc_mail_receive(struct proc* proc, struct proc_mail* mail, struct reschedule_ctx* rctx,
                       void* recv_buffer, size_t recv_size) {
  uint64_t fr, fssq;

  spin_lock(&mail->resource->lock, &fr);

  /* consume mesg if available */
  if (mail->packets_count > 0) {
    struct mail_packet* packet = list_entry(mail->packets, struct mail_packet, packets_link);
    list_remove(mail->packets, &packet->packets_link);
    mail->packets_count--;

    memcpy(recv_buffer, packet->packet_buffer, min(recv_size, packet->packet_size));
    free(packet->packet_buffer);
    free(packet);

    /* check for suspended sender */
    spin_lock(&mail->send_sq.lock, &fssq);

    struct list_node_link* node = mail->send_sq.proc_list;

    if (node != NULL) {
      struct proc_sq_entry* sq_entry = list_entry(node, struct proc_sq_entry, sq_link);
      struct proc* resumed_proc = sq_entry->proc;

      spin_unlock(&mail->send_sq.lock, fssq);
      spin_unlock(&mail->resource->lock, fr);

      proc_sq_resume(resumed_proc, sq_entry, rctx);
      return;
    }

    spin_unlock(&mail->send_sq.lock, fssq);
    spin_unlock(&mail->resource->lock, fr);

    return;
  }

  struct mail_saved_buffer* saved_buffer = malloc(sizeof(*saved_buffer));
  saved_buffer->buffer = recv_buffer;
  saved_buffer->size = recv_size;

  /* nothing to receive */
  proc_sq_suspend(proc, &mail->recv_sq, &mail->resource->lock, fr, rctx, saved_buffer,
                  &proc_mail_free_saved_buffer);
}