#include <libk/rbtree.h>
#include <libk/std.h>
#include <mm/liballoc.h>
#include <proc/proc.h>
#include <proc/procgroup.h>
#include <sync/rw_spin_lock.h>
#include <sync/spin_lock.h>
#include <sys/debug.h>
#include <sys/mm.h>

static struct rb_node_link* procgroup_tree = NULL;
static rw_spin_lock_t procgroup_tree_lock = RW_SPIN_LOCK_INIT;
static atomic_int pgids = 0;

struct procgroup* procgroup_create (void) {
  spin_lock_ctx_t ctxpgtr;

  struct procgroup* procgroup = malloc (sizeof (*procgroup));
  if (procgroup == NULL) {
    return NULL;
  }

  procgroup->refs = 0;
  procgroup->memb_proc_tree = NULL;
  procgroup->lock = SPIN_LOCK_INIT;
  procgroup->pgid = atomic_fetch_add (&pgids, 1);
  procgroup->pd.lock = SPIN_LOCK_INIT;
  procgroup->pd.cr3_paddr = mm_alloc_user_pd_phys ();

  rw_spin_write_lock (&procgroup_tree_lock, &ctxpgtr);
  rbtree_insert (struct procgroup, &procgroup_tree, &procgroup->procgroup_tree_link,
                 procgroup_tree_link, pgid);
  rw_spin_write_unlock (&procgroup_tree_lock, &ctxpgtr);

  return procgroup;
}

void procgroup_attach (struct procgroup* procgroup, struct proc* proc) {
  spin_lock_ctx_t ctxpg, ctxpr;

  spin_lock (&procgroup->lock, &ctxpg);
  spin_lock (&proc->lock, &ctxpr);

  rbtree_insert (struct proc, &procgroup->memb_proc_tree, &proc->procgroup_memb_tree_link,
                 procgroup_memb_tree_link, pid);
  atomic_fetch_add (&procgroup->refs, 1);
  DEBUG ("procgrpup attach PID %d to PGID %d\n", proc->pid, procgroup->pgid);

  spin_unlock (&proc->lock, &ctxpr);
  spin_unlock (&procgroup->lock, &ctxpg);
}

void procgroup_detach (struct procgroup* procgroup, struct proc* proc) {
  spin_lock_ctx_t ctxpg, ctxpr, ctxpgtr;

  spin_lock (&procgroup->lock, &ctxpg);
  spin_lock (&proc->lock, &ctxpr);

  rbtree_delete (&procgroup->memb_proc_tree, &proc->procgroup_memb_tree_link);
  int refs = atomic_fetch_sub (&procgroup->refs, 1);
  DEBUG ("procgrpup detach PID %d to PGID %d\n", proc->pid, procgroup->pgid);

  spin_unlock (&proc->lock, &ctxpr);
  spin_unlock (&procgroup->lock, &ctxpg);

  if (refs == 1) {
    rw_spin_write_lock (&procgroup_tree_lock, &ctxpgtr);
    spin_lock (&procgroup->lock, &ctxpg);

    rbtree_delete (&procgroup_tree, &procgroup->procgroup_tree_link);

    spin_unlock (&procgroup->lock, &ctxpg);
    rw_spin_write_unlock (&procgroup_tree_lock, &ctxpgtr);

    /* unlink resources */
    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;

      proc_resource_unlink (resource);
    }

    free (procgroup);
  }
}

int procgroup_get_sys_rid (struct procgroup* procgroup) {
  return atomic_fetch_add (&procgroup->sys_rids, 1);
}