#include <fs/vfs.h>
#include <id/id_alloc.h>
#include <libk/rbtree.h>
#include <libk/std.h>
#include <libk/string.h>
#include <mm/liballoc.h>
#include <mm/pmm.h>
#include <proc/proc.h>
#include <proc/procgroup.h>
#include <proc/reschedule.h>
#include <sync/spin_lock.h>
#include <sys/debug.h>
#include <sys/mm.h>

#define PGIDS_MAX 1024

static struct rb_node_link* procgroup_tree = NULL;
static spin_lock_t procgroup_tree_lock = SPIN_LOCK_INIT;

static struct id_alloc pgid_alloc;

void procgroup_pgid_alloc_init (void) { id_alloc_init (&pgid_alloc, PGIDS_MAX); }

struct procgroup* procgroup_find (int pgid) {
  struct procgroup* procgroup = NULL;

  spin_lock (&procgroup_tree_lock);
  rbtree_find (struct procgroup, &procgroup_tree, pgid, procgroup, procgroup_tree_link, pgid);
  spin_unlock (&procgroup_tree_lock);

  return procgroup;
}

uintptr_t procgroup_map (struct procgroup* procgroup, uintptr_t vaddr, size_t pages, uint32_t flags,
                         uintptr_t* out_paddr) {
  spin_lock (&procgroup->lock);

  vaddr = (vaddr == 0) ? procgroup->map_base : vaddr;

  struct proc_mapping* mapping = malloc (sizeof (*mapping));

  if (mapping == NULL) {
    spin_unlock (&procgroup->lock);
    return 0;
  }

  uintptr_t paddr = pmm_alloc (pages);

  if (paddr == PMM_ALLOC_ERR) {
    free (mapping);
    spin_unlock (&procgroup->lock);
    return 0;
  }

  if (out_paddr != NULL)
    *out_paddr = paddr;

  mapping->paddr = paddr;
  mapping->vaddr = vaddr;
  mapping->size = pages * PAGE_SIZE;

  procgroup->map_base += pages * PAGE_SIZE;

  list_append (procgroup->mappings, &mapping->proc_mappings_link);

  for (uintptr_t vpage = vaddr, ppage = paddr; vpage < vaddr + pages * PAGE_SIZE;
       vpage += PAGE_SIZE, ppage += PAGE_SIZE) {
    mm_map_page (&procgroup->pd, ppage, vpage, flags);
  }

  spin_unlock (&procgroup->lock);

  return vaddr;
}

bool procgroup_unmap (struct procgroup* procgroup, uintptr_t start_vaddr, size_t pages) {
  size_t unmap_size = pages * PAGE_SIZE;
  uintptr_t end_vaddr = start_vaddr + unmap_size;

  struct list_node_link *mapping_link, *mapping_link_tmp;

  bool used_tail_mapping = false;

  struct proc_mapping* tail_mapping = malloc (sizeof (*tail_mapping));
  if (tail_mapping == NULL)
    return false;

  spin_lock (&procgroup->lock);

  list_foreach (procgroup->mappings, mapping_link, mapping_link_tmp) {
    struct proc_mapping* mapping =
        list_entry (mapping_link, struct proc_mapping, proc_mappings_link);

    uintptr_t m_start = mapping->vaddr;
    uintptr_t m_end = mapping->vaddr + mapping->size;

    /* check overlap */
    if ((start_vaddr < m_end) && (end_vaddr > mapping->vaddr)) {
      uintptr_t free_vstart = (start_vaddr > m_start) ? start_vaddr : m_start;
      uintptr_t free_vend = (end_vaddr < m_end) ? end_vaddr : m_end;
      size_t free_size = free_vend - free_vstart;

      uintptr_t ppage_to_free = mapping->paddr + (free_vstart - m_start);
      pmm_free (ppage_to_free, free_size / PAGE_SIZE);

      /* split in the middle */
      if ((start_vaddr > m_start) && (end_vaddr < m_end)) {
        tail_mapping->vaddr = end_vaddr;
        tail_mapping->paddr = mapping->paddr + (end_vaddr - m_start);
        tail_mapping->size = m_end - end_vaddr;

        mapping->size = start_vaddr - m_start;

        list_insert_after (procgroup->mappings, &mapping->proc_mappings_link,
                           &tail_mapping->proc_mappings_link);

        used_tail_mapping = true;

        break;
      } else if ((start_vaddr <= m_start) && (end_vaddr < m_end)) { /* shrink left */
        size_t diff = end_vaddr - m_start;
        mapping->vaddr += diff;
        mapping->paddr += diff;
        mapping->size -= diff;
      } else if ((start_vaddr > m_start) && (end_vaddr >= m_end)) { /* shrink right */
        mapping->size = start_vaddr - m_start;
      } else { /* full overlap */
        list_remove (procgroup->mappings, &mapping->proc_mappings_link);
        free (mapping);
      }
    }
  }

  if (!used_tail_mapping)
    free (tail_mapping);

  for (uintptr_t vpage = start_vaddr; vpage < end_vaddr; vpage += PAGE_SIZE) {
    mm_unmap_page (&procgroup->pd, vpage);
  }

  spin_unlock (&procgroup->lock);

  return true;
}

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

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

  if (!id_alloc_init (&procgroup->rid_alloc, PROCGROUP_RESOURCES_MAX)) {
    free (procgroup);
    return NULL;
  }

  procgroup->pgid = id_alloc (&pgid_alloc);

  if (procgroup->pgid < 0) {
    id_alloc_fini (&procgroup->rid_alloc);
    free (procgroup);
    return NULL;
  }

  procgroup->memb_proc_tree = NULL;
  procgroup->lock = SPIN_LOCK_INIT;
  procgroup->pd.cr3_paddr = mm_alloc_user_pd_phys ();
  procgroup->map_base = PROC_MAP_BASE;

  if (proc_create_resource_mail (procgroup) == NULL) {
    id_alloc_fini (&procgroup->rid_alloc);
    free (procgroup);
    return NULL;
  }

  spin_lock (&procgroup_tree_lock);
  rbtree_insert (struct procgroup, &procgroup_tree, &procgroup->procgroup_tree_link,
                 procgroup_tree_link, pgid);
  spin_unlock (&procgroup_tree_lock);

  return procgroup;
}

void procgroup_attach (struct procgroup* procgroup, struct proc* proc) {
  spin_lock (&procgroup->lock);
  spin_lock (&proc->lock);

  rbtree_insert (struct proc, &procgroup->memb_proc_tree, &proc->procgroup_memb_tree_link,
                 procgroup_memb_tree_link, pid);

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

static void procgroup_delete (struct procgroup* procgroup, struct reschedule_ctx* rctx) {
  spin_lock (&procgroup_tree_lock);
  spin_lock (&procgroup->lock);

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

  spin_unlock (&procgroup->lock);
  spin_unlock (&procgroup_tree_lock);

  /* delete 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_delete_resource (procgroup, resource, rctx);
  }

  /* delete mappings */
  struct list_node_link *mapping_link, *mapping_link_tmp;
  list_foreach (procgroup->mappings, mapping_link, mapping_link_tmp) {
    struct proc_mapping* mapping =
        list_entry (mapping_link, struct proc_mapping, proc_mappings_link);

    pmm_free (mapping->paddr, mapping->size / PAGE_SIZE);
    free (mapping);
  }

  pmm_free (procgroup->pd.cr3_paddr, 1);

  free (procgroup->tls.tls_tmpl);

  id_alloc_fini (&procgroup->rid_alloc);
  id_free (&pgid_alloc, procgroup->pgid);

  free (procgroup);
}

void procgroup_detach (struct procgroup* procgroup, struct proc* proc,
                       struct reschedule_ctx* rctx) {
  spin_lock (&procgroup->lock);
  spin_lock (&proc->lock);

  rbtree_delete (&procgroup->memb_proc_tree, &proc->procgroup_memb_tree_link);
  struct rb_node_link* memb_tree = procgroup->memb_proc_tree;

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

  if (memb_tree == NULL) {
    procgroup_delete (procgroup, rctx);
  }
}
