#include <libk/assert.h>
#include <libk/rbtree.h>
#include <libk/std.h>
#include <libk/string.h>
#include <mm/malloc.h>
#include <proc/mutex.h>
#include <proc/proc.h>
#include <proc/suspension_q.h>
#include <sync/spin_lock.h>
#include <sys/debug.h>
#include <sys/smp.h>
#include <sys/spin_lock.h>

void proc_mutexes_cleanup(struct proc* proc, struct reschedule_ctx* rctx) {
  struct procgroup* procgroup = proc->procgroup;

  struct rb_node_link* rnode;
  rbtree_first(&procgroup->resource_tree, rnode);

  while (rnode) {
    struct rb_node_link* next;
    rbtree_next(rnode, next);

    struct proc_resource* resource = rbtree_entry(rnode, struct proc_resource, resource_tree_link);

    rnode = next;

    if (resource->type != PR_MUTEX) {
      continue;
    }

    if (resource->u.mutex.owner == proc && resource->u.mutex.locked) {
      proc_mutex_unlock(proc, &resource->u.mutex, rctx);
    }
  }
}

void proc_cleanup_resource_mutex(struct proc_resource* resource, struct reschedule_ctx* rctx) {
  struct proc_mutex* mutex = &resource->u.mutex;

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

    proc_sq_resume(suspended_proc, sq_entry, rctx, 0);
  }

  mutex->locked = false;
  mutex->owner = NULL;
}

void proc_mutex_lock(struct proc* proc, struct proc_mutex* mutex, struct reschedule_ctx* rctx) {
  if (!mutex->locked || mutex->owner == proc) {
    mutex->locked = true;
    mutex->owner = proc;
    return;
  }

  proc_sq_suspend(proc, &mutex->suspension_q, rctx, NULL, NULL);
}

void proc_mutex_unlock(struct proc* proc, struct proc_mutex* mutex, struct reschedule_ctx* rctx) {
  if (mutex->owner != proc) {
    return;
  }

  struct list_node_link* node = mutex->suspension_q.proc_list;

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

    mutex->owner = resumed_proc;
    mutex->locked = true;

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

  mutex->locked = false;
  mutex->owner = NULL;
}