#include <amd64/fx.h>
#include <amd64/gdt.h>
#include <amd64/proc.h>
#include <aux/elf.h>
#include <fs/vfs.h>
#include <libk/align.h>
#include <libk/list.h>
#include <libk/rbtree.h>
#include <libk/std.h>
#include <libk/string.h>
#include <limine/requests.h>
#include <mm/liballoc.h>
#include <mm/pmm.h>
#include <proc/mutex.h>
#include <proc/proc.h>
#include <proc/procgroup.h>
#include <proc/resource.h>
#include <sync/spin_lock.h>
#include <sys/debug.h>
#include <sys/proc.h>

struct proc* proc_from_elf (uint8_t* elf_contents) {
  struct limine_hhdm_response* hhdm = limine_hhdm_request.response;

  struct proc* proc = malloc (sizeof (*proc));
  if (proc == NULL)
    return NULL;

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

  proc->lock = SPIN_LOCK_INIT;
  proc->state = PROC_READY;
  proc->pid = proc_alloc_pid ();

  if (proc->pid < 0) {
    free (proc);
    return NULL;
  }

  proc->procgroup = procgroup_create ();
  if (proc->procgroup == NULL) {
    free (proc);
    return NULL;
  }
  procgroup_attach (proc->procgroup, proc);

  uintptr_t kstack_paddr = pmm_alloc (KSTACK_SIZE / PAGE_SIZE);
  proc->pdata.kernel_stack = kstack_paddr + (uintptr_t)hhdm->offset + KSTACK_SIZE;

  procgroup_map (proc->procgroup, PROC_USTACK_TOP - USTACK_SIZE, USTACK_SIZE / PAGE_SIZE,
                 MM_PG_USER | MM_PG_PRESENT | MM_PG_RW, NULL);

  proc->flags |= PROC_USTK_PREALLOC;

  struct elf_aux aux = proc_load_segments (proc, elf_contents);

  proc->pdata.regs.ss = GDT_UDATA | 0x03;
  proc->pdata.regs.rsp = (uint64_t)PROC_USTACK_TOP;
  proc->pdata.regs.rflags = 0x202;
  proc->pdata.regs.cs = GDT_UCODE | 0x03;
  proc->pdata.regs.rip = aux.entry;
  fx_init (proc->pdata.fx_env);

  proc->exec_pid = -1;

  return proc;
}

struct proc* proc_clone (struct proc* proto, uintptr_t vstack_top, uintptr_t entry,
                         uintptr_t argument_ptr) {
  struct limine_hhdm_response* hhdm = limine_hhdm_request.response;

  struct proc* proc = malloc (sizeof (*proc));
  if (proc == NULL)
    return NULL;

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

  proc->lock = SPIN_LOCK_INIT;
  proc->state = PROC_READY;
  proc->pid = proc_alloc_pid ();

  if (proc->pid < 0) {
    free (proc);
    return NULL;
  }

  spin_lock (&proto->lock);

  proc->procgroup = proto->procgroup;
  procgroup_attach (proc->procgroup, proc);

  spin_unlock (&proto->lock);

  uintptr_t kstack_paddr = pmm_alloc (KSTACK_SIZE / PAGE_SIZE);
  proc->pdata.kernel_stack = kstack_paddr + (uintptr_t)hhdm->offset + KSTACK_SIZE;

  proc->pdata.regs.ss = GDT_UDATA | 0x03;
  proc->pdata.regs.rsp = (uint64_t)vstack_top;
  proc->pdata.regs.rflags = 0x202;
  proc->pdata.regs.cs = GDT_UCODE | 0x03;
  proc->pdata.regs.rip = (uint64_t)entry;
  fx_init (proc->pdata.fx_env);

  proc->uvaddr_argument = argument_ptr;

  proc_init_tls (proc);

  proc->exec_pid = -1;

  return proc;
}

void proc_cleanup (struct proc* proc, struct reschedule_ctx* rctx) {
  spin_lock (&proc->lock);
  spin_lock (&proc->done_sq.lock);

  while (proc->done_sq.proc_list != NULL) {
    struct list_node_link* node = proc->done_sq.proc_list;
    struct proc_sq_entry* sq_entry = list_entry (node, struct proc_sq_entry, sq_link);
    struct proc* suspended_proc = sq_entry->proc;

    spin_unlock (&proc->done_sq.lock);
    spin_unlock (&proc->lock);

    proc_sq_resume (suspended_proc, sq_entry, rctx);

    spin_lock (&proc->lock);
    spin_lock (&proc->done_sq.lock);
  }

  spin_unlock (&proc->done_sq.lock);
  spin_unlock (&proc->lock);

  proc_sqs_cleanup (proc);
  proc_mutexes_cleanup (proc, rctx);

  pmm_free (proc->pdata.kernel_stack, KSTACK_SIZE / PAGE_SIZE);
  procgroup_unmap (proc->procgroup, proc->pdata.tls_vaddr, proc->procgroup->tls.tls_tmpl_pages);

  procgroup_detach (proc->procgroup, proc, rctx);

  vfs_volume_close (proc, proc->cwv, rctx);
  proc_free_pid (proc->pid);
  /* clean the process */
  free (proc);
}

void proc_init_tls (struct proc* proc) {
  struct limine_hhdm_response* hhdm = limine_hhdm_request.response;

  if (proc->procgroup->tls.tls_tmpl == NULL)
    return;

  size_t tls_size = proc->procgroup->tls.tls_tmpl_size;
  size_t pages = proc->procgroup->tls.tls_tmpl_pages;

  uintptr_t tls_paddr;
  uint32_t flags = MM_PG_USER | MM_PG_PRESENT | MM_PG_RW;

  uintptr_t tls_vaddr = procgroup_map (proc->procgroup, 0, pages, flags, &tls_paddr);

  uintptr_t k_tls_addr = (uintptr_t)hhdm->offset + tls_paddr;

  memset ((void*)k_tls_addr, 0, pages * PAGE_SIZE);
  memcpy ((void*)k_tls_addr, (void*)proc->procgroup->tls.tls_tmpl, tls_size);

  uintptr_t ktcb = k_tls_addr + tls_size;
  uintptr_t utcb = tls_vaddr + tls_size;

  *(uintptr_t*)ktcb = utcb;

  proc->pdata.fs_base = utcb;
  proc->pdata.tls_vaddr = tls_vaddr;
}
