#include <libk/assert.h>
#include <libk/rbtree.h>
#include <libk/std.h>
#include <libk/string.h>
#include <mm/liballoc.h>
#include <mm/pmm.h>
#include <proc/mutex.h>
#include <proc/proc.h>
#include <proc/resource.h>
#include <sync/spin_lock.h>
#include <sys/debug.h>

void proc_cleanup_resources (struct proc* proc) {
  struct proc_resource* resource = NULL;

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

  while (rnode) {
    struct rb_node_link* next;
    rbtree_next (rnode, next);
    resource = rbtree_entry (rnode, struct proc_resource, proc_resource_tree_link);
    proc_drop_resource (proc, resource);
    rnode = next;
  }

  assert (proc->resource_tree == NULL);
}

void proc_drop_resource (struct proc* proc, struct proc_resource* resource) {
  spin_lock_ctx_t ctxpr;

  DEBUG ("resource=%p, type=%d, rid=%d\n", resource, resource->type, resource->rid);

  if (atomic_fetch_sub (&resource->refs, 1) == 1) {
    spin_lock (&proc->lock, &ctxpr);
    rbtree_delete (&proc->resource_tree, &resource->proc_resource_tree_link);
    spin_unlock (&proc->lock, &ctxpr);

    resource->ops.cleanup (proc, resource);
    free (resource);
  }
}

static bool proc_create_resource_mem (struct proc_resource_mem* mem,
                                      struct proc_resource_mem_init* init) {
  if (init->pages == 0)
    return false;

  uintptr_t paddr = pmm_alloc (init->pages);
  if (paddr == PMM_ALLOC_ERR)
    return false;

  mem->paddr = paddr;
  mem->pages = mem->alive_pages = init->pages;

  return true;
}

static void proc_cleanup_resource_mem (struct proc* proc, struct proc_resource* resource) {
  (void)proc;
  pmm_free (resource->u.mem.paddr, resource->u.mem.pages);
}

static bool proc_create_resource_mutex (struct proc_mutex* mutex) {
  memset (mutex, 0, sizeof (*mutex));

  return true;
}

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

  proc_mutex_unlock (proc, mutex);
}

struct proc_resource* proc_create_resource (struct proc* proc, int rid, int type, int vis,
                                            void* data) {
  spin_lock_ctx_t ctxpr;
  /* Check if resource RID already exists */
  struct proc_resource* resource_check;
  rbtree_find (struct proc_resource, &proc->resource_tree, rid, resource_check,
               proc_resource_tree_link, rid);
  if (resource_check != NULL)
    return NULL;

  struct proc_resource* resource = malloc (sizeof (*resource));
  if (resource == NULL)
    return NULL;

  memset (resource, 0, sizeof (*resource));

  resource->lock = SPIN_LOCK_INIT;
  resource->type = type;
  resource->refs = 1;
  resource->rid = rid;
  resource->visibility = vis;

  switch (resource->type) {
  case PR_MEM: {
    struct proc_resource_mem_init* mem_init = data;
    proc_create_resource_mem (&resource->u.mem, mem_init);
    resource->ops.cleanup = &proc_cleanup_resource_mem;
    resource->u.mem.resource = resource;
    DEBUG ("PR_MEM resource=%p type=%d rid=%d paddr=%p, pages=%zu\n", resource, resource->type,
           resource->rid, resource->u.mem.paddr, resource->u.mem.pages);
  } break;
  case PR_MUTEX: {
    proc_create_resource_mutex (&resource->u.mutex);
    resource->ops.cleanup = &proc_cleanup_resource_mutex;
    resource->u.mutex.resource = resource;
    DEBUG ("PR_MUTEX resource=%p, type=%d rid=%d\n", resource, resource->type, resource->rid);
  } break;
  default: {
    free (resource);
    return NULL;
  } break;
  }

  spin_lock (&proc->lock, &ctxpr);
  rbtree_insert (struct proc_resource, &proc->resource_tree, &resource->proc_resource_tree_link,
                 proc_resource_tree_link, rid);
  spin_unlock (&proc->lock, &ctxpr);

  return resource;
}