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

#define MAX_PROCESSES 256

static process_t *processes[MAX_PROCESSES];
static uint64_t free_pid = 1;

process_t *current_process;

process_t *process_create(const process_t *parent, const uint8_t *code, uint64_t size) {
  if (size >= USER_VIRTUAL_HEAP - USER_VIRTUAL_CODE) {
    return NUL;
  }

  uint64_t free_pi;
  for (free_pi = 0; free_pi <= MAX_PROCESSES; free_pi++) {
    if (!processes[free_pi]) {
      break;
    }
  }
  if (free_pi == MAX_PROCESSES) {
    return NUL;
  }

  process_t *proc = processes[free_pi] = 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);
  }

  proc->pid = free_pid++;
  proc->state = PROCESS_RUNNING;
  proc->cwd = path_open_again(parent->cwd);
  proc->fds[0] = parent->fds[0];
  proc->fds[1] = parent->fds[1];
  proc->free_fd = 2;
  proc->code = EXIT_CODE_OK;

  return proc;
}

process_t *process_get(uint64_t pid) {
  for (uint64_t i = 0; i < MAX_PROCESSES; i++) {
    if (processes[i] && processes[i]->pid == pid) {
      return processes[i];
    }
  }
  return NUL;
}

void process_next() {
  process_t *next = NUL;
  uint64_t cpi = 0;

  for (uint64_t i = 0; i < MAX_PROCESSES; i++) {
    if (processes[i] == current_process) {
      cpi = i;
      break;
    }
  }

  for (uint64_t i = 1; i <= MAX_PROCESSES; i++) {
    process_t *candidate = processes[(cpi + i) % MAX_PROCESSES];
    if (candidate && candidate->state == PROCESS_RUNNING) {
      next = candidate;
      break;
    }
  }

  if (!next) {
    outw(0x604, 0x2000);
    __asm__ volatile("cli; hlt");
  }

  if (next == current_process) {
    return;
  }

  process_t *prev = current_process;
  tss.rsp0 = (uint64_t)next->kernel_stack;
  current_process = next;
  __asm__ volatile("cli");
  process_switch_to(&prev->kernel_rsp, next->kernel_rsp, VIRT_TO_PHYS(next->pml4));
  __asm__ volatile("sti");
}

void process_destroy(process_t *proc) {
  for (uint64_t i = 2; i < MAX_FDS; i++) {
    if (proc->fds[i]) {
      proc->fds[i]->close(proc->fds[i]);
      proc->fds[i] = NUL;
    }
  }

  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);

  for (uint64_t i = 0; i < MAX_PROCESSES; i++) {
    if (processes[i] == proc) {
      processes[i] = NUL;
      break;
    }
  }

  memory_free(proc);
}