#include <libk/assert.h>
#include <libk/rbtree.h>
#include <libk/std.h>
#include <libk/string.h>
#include <mm/liballoc.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) {
  spin_lock (&proc->procgroup->lock);

  struct rb_node_link* rnode;
  rbtree_first (&proc->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;

    spin_lock (&resource->lock);

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

    if (resource->u.mutex.owner == proc && resource->u.mutex.locked) {
      spin_unlock (&resource->lock);

      proc_mutex_unlock (proc, &resource->u.mutex, rctx);
    }

    spin_unlock (&resource->lock);
  }

  spin_unlock (&proc->procgroup->lock);
}

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

  spin_lock (&mutex->resource->lock);
  spin_lock (&mutex->suspension_q.lock);

  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;

    /* we will relock during resume */
    spin_unlock (&mutex->suspension_q.lock);
    spin_unlock (&mutex->resource->lock);

    proc_sq_resume (suspended_proc, sq_entry, rctx);

    /* reacquire */
    spin_lock (&mutex->resource->lock);
    spin_lock (&mutex->suspension_q.lock);
  }

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

  spin_unlock (&mutex->suspension_q.lock);
  spin_unlock (&mutex->resource->lock);
}

void proc_mutex_lock (struct proc* proc, struct proc_mutex* mutex, struct reschedule_ctx* rctx) {
  spin_lock (&mutex->resource->lock);

  if (!mutex->locked || mutex->owner == proc) {
    mutex->locked = true;
    mutex->owner = proc;
    spin_unlock (&mutex->resource->lock);
    return;
  }

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

void proc_mutex_unlock (struct proc* proc, struct proc_mutex* mutex, struct reschedule_ctx* rctx) {
  spin_lock (&mutex->resource->lock);

  if (mutex->owner != proc) {
    spin_unlock (&mutex->resource->lock);
    return;
  }

  spin_lock (&mutex->suspension_q.lock);

  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;

    spin_unlock (&mutex->suspension_q.lock);
    spin_unlock (&mutex->resource->lock);

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

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

  spin_unlock (&mutex->suspension_q.lock);
  spin_unlock (&mutex->resource->lock);
}
