#include <libk/assert.h>
#include <libk/list.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>

static struct rb_node_link* resource_tree = NULL;
static rw_spin_lock_t resource_tree_lock = RW_SPIN_LOCK_INIT;

void proc_cleanup_resources (struct proc* proc) {
  spin_lock_ctx_t ctxrs;

  rw_spin_write_lock (&proc->resources->lock, &ctxrs);

  struct rb_node_link* rnode;
  rbtree_first (&proc->resources->tree, rnode);
  while (rnode) {
    struct rb_node_link* next;
    rbtree_next (rnode, next);

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

    rnode = next;

    proc_drop_resource (proc, resource, false);
  }

  rw_spin_write_unlock (&proc->resources->lock, &ctxrs);

  if (atomic_fetch_sub (&proc->resources->refs, 1) == 1) {
    free (proc->resources);
  }
}

void proc_drop_resource (struct proc* proc, struct proc_resource* resource, bool lock) {
  spin_lock_ctx_t ctxrs;

  if (atomic_fetch_sub (&resource->refs, 1) == 1) {
    DEBUG ("resource=%p created_by=%d vis=%d type=%d rid=%d refs=%d\n", resource,
           resource->created_by_pid, resource->visibility, resource->type, resource->rid,
           atomic_load (&resource->refs));

    switch (resource->visibility) {
    case RV_PRIVATE: {
      if (lock)
        rw_spin_write_lock (&proc->resources->lock, &ctxrs);

      rbtree_delete (&proc->resources->tree, &resource->local_resource_tree_link);

      if (lock)
        rw_spin_write_unlock (&proc->resources->lock, &ctxrs);
    } break;
    case RV_PUBLIC: {
      if (lock)
        rw_spin_write_lock (&resource_tree_lock, &ctxrs);

      rbtree_delete (&resource_tree, &resource->global_resource_tree_link);

      if (lock)
        rw_spin_write_unlock (&resource_tree_lock, &ctxrs);
    } break;
    default: {
      assert (0);
    } break;
    }

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

struct proc_resource* proc_find_resource (struct proc* proc, int rid, int vis) {
  struct proc_resource* resource = NULL;
  spin_lock_ctx_t ctxrs;

  switch (vis) {
  case RV_PRIVATE: {
    /* User wants to create a private resource, so search locally */
    rw_spin_read_lock (&proc->resources->lock, &ctxrs);
    rbtree_find (struct proc_resource, &proc->resources->tree, rid, resource,
                 local_resource_tree_link, rid);
    rw_spin_read_unlock (&proc->resources->lock, &ctxrs);
  } break;
  case RV_PUBLIC: {
    /* User wants to create a public resource, so search globally */
    rw_spin_read_lock (&resource_tree_lock, &ctxrs);
    rbtree_find (struct proc_resource, &resource_tree, rid, resource, global_resource_tree_link,
                 rid);
    rw_spin_read_unlock (&resource_tree_lock, &ctxrs);
  } break;
  default: {
    assert (0);
  } break;
  }

  return resource;
}

struct proc_resource* proc_create_resource (struct proc* proc, int rid, int type, int vis,
                                            void* data) {
  spin_lock_ctx_t ctxrs;

  /* Check if resource RID already exists */
  struct proc_resource* resource_check = proc_find_resource (proc, rid, vis);

  /* Resource was found either way, so it already exists */
  if (resource_check != NULL)
    return NULL;

  /* create the resource */

  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;
  resource->created_by_pid = proc->pid;

  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 created_by=%d, type=%d rid=%d paddr=%p, pages=%zu\n", resource,
           resource->created_by_pid, 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 created_by=%d type=%d rid=%d\n", resource,
           resource->created_by_pid, resource->type, resource->rid);
  } break;
  default: {
    free (resource);
    return NULL;
  } break;
  }

  switch (resource->visibility) {
  case RV_PRIVATE: {
    rw_spin_write_lock (&proc->resources->lock, &ctxrs);
    rbtree_insert (struct proc_resource, &proc->resources->tree,
                   &resource->local_resource_tree_link, local_resource_tree_link, rid);
    rw_spin_write_unlock (&proc->resources->lock, &ctxrs);
  } break;
  case RV_PUBLIC: {
    rw_spin_write_lock (&resource_tree_lock, &ctxrs);
    rbtree_insert (struct proc_resource, &resource_tree, &resource->global_resource_tree_link,
                   global_resource_tree_link, rid);
    rw_spin_write_unlock (&resource_tree_lock, &ctxrs);
  } break;
  default: {
    assert (0);
  } break;
  }

  return resource;
}