#include "src/kernel/process.h"
#include "src/kernel/memory.h"
#include "src/kernel/panic.h"
#include "src/kernel/path.h"
#include "src/lib/layout.h"
#include "src/lib/memory.h"

process_t *current_process;

process_t *process_create(const process_t *parent, const uint8_t *code, uint64_t size) {
  ASSERT(size <= USER_VIRTUAL_HEAP - USER_VIRTUAL_CODE, "process_create: binary too big");

  process_t *proc = memory_allocate(sizeof(process_t));

  proc->parent = parent;

  proc->kernel_stack = (uint8_t *)memory_allocate(PAGE_SIZE) + PAGE_SIZE;

  uint64_t *kstackv = (uint64_t *)((uint8_t *)proc->kernel_stack);
  *(--kstackv) = (uint64_t)process_trampoline; // "return address" for switch_to's ret
  *(--kstackv) = 0;                            // r15
  *(--kstackv) = 0;                            // r14
  *(--kstackv) = 0;                            // r13
  *(--kstackv) = 0;                            // r12
  *(--kstackv) = 0;                            // rbx
  *(--kstackv) = 0;                            // rbp
  proc->kernel_rsp = kstackv;

  proc->pml4 = PHYS_TO_VIRT(memory_page_allocate());
  memory_set(0, PAGE_SIZE / 2, proc->pml4);
  memory_copy((uint64_t *)KERNEL_VIRTUAL_PML4 + 256, PAGE_SIZE / 2, proc->pml4 + 256);

  void *ustackp = memory_page_allocate();
  proc->user_stack = PHYS_TO_VIRT((uint64_t)ustackp + PAGE_SIZE);
  memory_page_map(proc->pml4, (void *)USER_VIRTUAL_STACK, ustackp, PAGE_PRESENT | PAGE_WRITABLE | PAGE_USER);

  for (uint64_t page = 0; size > 0; page++) {
    void *p = memory_page_allocate();
    uint64_t s = size < PAGE_SIZE ? size : PAGE_SIZE;
    memory_page_map(proc->pml4, (void *)(USER_VIRTUAL_CODE + page * PAGE_SIZE), p,
                    PAGE_PRESENT | PAGE_WRITABLE | PAGE_USER); // TODO Restrict writeability of code.
    memory_copy(code, s, PHYS_TO_VIRT(p));
    code += s;
    size -= s;
  }

  for (uint64_t page = 0; page < HEAP_PAGE_COUNT; page++) {
    memory_page_map(proc->pml4, (void *)(USER_VIRTUAL_HEAP + page * PAGE_SIZE), memory_page_allocate(),
                    PAGE_PRESENT | PAGE_WRITABLE | PAGE_USER);
  }

  path_t *cwd = memory_allocate(sizeof(path_t));
  cwd->depth = parent->cwd->depth;
  for (uint8_t i = 0; i < parent->cwd->depth; i++) {
    cwd->stack[i] = fs_open_again(parent->cwd->stack[i]);
  }

  proc->cwd = cwd;
  proc->root = parent->root;
  proc->stdin = parent->stdin;
  proc->stdout = parent->stdout;
  proc->free_fd = 2;

  return proc;
}

void process_destroy(process_t *proc) {
  path_close((path_t *)proc->cwd);

  for (uint16_t pml4i = 0; pml4i < 256; pml4i++) {
    if (!(proc->pml4[pml4i] & PAGE_PRESENT)) {
      continue;
    }

    uint64_t *pdpt = PHYS_TO_VIRT(proc->pml4[pml4i] & ~(uint64_t)0xFFF);
    for (uint16_t pdpti = 0; pdpti < 512; pdpti++) {
      if (!(pdpt[pdpti] & PAGE_PRESENT)) {
        continue;
      }

      uint64_t *pd = PHYS_TO_VIRT(pdpt[pdpti] & ~(uint64_t)0xFFF);
      for (uint16_t pdi = 0; pdi < 512; pdi++) {
        if (!(pd[pdi] & PAGE_PRESENT)) {
          continue;
        }

        if (pd[pdi] & PAGE_HUGE) {
          memory_page_free((void *)(pd[pdi] & ~(uint64_t)0xFFF));
          continue;
        }

        uint64_t *pt = PHYS_TO_VIRT(pd[pdi] & ~(uint64_t)0xFFF);
        for (uint16_t pti = 0; pti < 512; pti++) {
          if (!(pt[pti] & PAGE_PRESENT)) {
            continue;
          }

          memory_page_free((void *)(pt[pti] & ~(uint64_t)0xFFF));
        }
        memory_page_free((void *)(pd[pdi] & ~(uint64_t)0xFFF));
      }
      memory_page_free((void *)(pdpt[pdpti] & ~(uint64_t)0xFFF));
    }
    memory_page_free((void *)(proc->pml4[pml4i] & ~(uint64_t)0xFFF));
  }

  memory_page_free(VIRT_TO_PHYS(proc->pml4));

  memory_free(proc->kernel_stack - PAGE_SIZE);

  memory_free(proc);
}