#include "sched/sched.h"
#include "arch/x86_64/idt.h"
#include "arch/x86_64/msr.h"
#include "mm/memop.h"
#include "mm/pmm.h"
#include "mm/vmm.h"
#include "sys/log.h"
#include <lib/string.h>
#include <mm/liballoc/liballoc.h>
#include <stddef.h>

sched_process *proc_list;
sched_process *curr_proc;
int current_pid = 0;
int standby = 0;

void map_range_to_pagemap(pagemap_t *dest_pagemap, pagemap_t *src_pagemap,
                          uint64_t start, uint64_t size, int user,
                          uint64_t flags) {
  for (uint64_t offset = 0; offset < size; offset += PMM_PAGE_SIZE) {
    uint64_t phys = virt_to_phys(src_pagemap, start + offset);
    if (phys) {
      if (user)
        vmm_map_user(dest_pagemap, start + offset, phys, flags);
      else
        vmm_map(dest_pagemap, start + offset, phys, flags);
    }
  }
}

void sched_init() {
  // TODO: It may be good to implement heap memory to save space.

  // We must initialize the process list.
  // By default, sched_create will append to this list.
  proc_list = malloc(sizeof(sched_process));
  memcpy(proc_list->name, "System\0", 7);
  proc_list->pid = -1;
  proc_list->type = SCHED_EMPTY;
  proc_list->flags = SCHED_KERNEL_PROCESS;
  proc_list->pm = vmm_kernel_pm;

  curr_proc = proc_list;

  standby = 1;
  log("sched - As there's nothing "
      "to schedule, the scheduler entered standby"
      "mode.\n");
}

sched_process *sched_create(char *name, uint64_t entry_point, pagemap_t *pm,
                            uint32_t flags) {

  sched_process *proc = malloc(sizeof(sched_process));
  memset(proc, 0, sizeof(sched_process));

  memcpy(proc->name, name, strlen(name) > 128 ? 128 : strlen(name));
  proc->pid = current_pid;
  proc->type = SCHED_RUNNING;
  proc->flags = flags;

  // We are about to setup the registers ourself.
  // If it's broken, it's a boom in the ass of your computer
  // (and a CPU exception)

  proc->pm = pm;

  uint64_t *stack_phys = pmm_request_page();
  uint64_t *kstack_phys = pmm_request_page();
  uint64_t *stack_virt = (uint64_t *)0x40000000;

  proc->kernel_stack = HIGHER_HALF(kstack_phys) + PMM_PAGE_SIZE;
  if (flags == SCHED_KERNEL_PROCESS) {
    proc->stack_base = HIGHER_HALF(stack_phys);
    proc->stack_base_physical = stack_phys;
    proc->stack_end = proc->stack_base + PMM_PAGE_SIZE;
  } else if (flags == SCHED_USER_PROCESS) {
    vmm_map_user(proc->pm, (uint64_t)stack_virt, (uint64_t)stack_phys,
                 VMM_PRESENT | VMM_WRITABLE | VMM_USER);
    proc->stack_base = stack_virt;
    proc->stack_base_physical = stack_phys;
    proc->stack_end = proc->stack_base + PMM_PAGE_SIZE;
  }
  proc->regs.rip = (uint64_t)entry_point;

  if (flags == SCHED_KERNEL_PROCESS) {
    proc->regs.cs = 0x28; // Run in kernel mode
    proc->regs.ss = 0x30;
  } else if (flags == SCHED_USER_PROCESS) {
    proc->regs.cs = 0x40 | 3; // Run in user mode
    proc->regs.ss = 0x38 | 3;
  }
  proc->regs.rflags = 0x202; // Enable interrupts
  proc->regs.rsp = (uint64_t)proc->stack_end;
  proc->regs.rbp = 0;

  proc->next = curr_proc->next;
  curr_proc->next = proc;

  current_pid++;

  if (flags == SCHED_USER_PROCESS)
    map_range_to_pagemap(pm, vmm_kernel_pm, 0x1000, 0x10000, 1,
                         VMM_PRESENT | VMM_WRITABLE | VMM_USER);
  else
    map_range_to_pagemap(pm, vmm_kernel_pm, 0x1000, 0x10000, 0,
                         VMM_PRESENT | VMM_WRITABLE);

  if (standby) {
    // Disable standby mode as there's actually something to
    // run, now.
    standby = 0;
    log("sched - Standby mode has been"
        "disabled.\n");
  }

  log("sched - created process '%s' (pid: %d, rip: %p)\n", proc->name,
      proc->pid, proc->regs.rip);
  return proc;
}

sched_process *sched_from_program(program_t *prog) {
  // D̶̡͖̼̪̩̽̃̿͐͐̀̑͐̑ǫ̴͚̘̩̗̠̣̗͇̳̀̅̈͌͆͆ ̸̢̢̯͎̳̟̑͋́N̷̹͛̋͗ö̴͇̼̖̝̼́̀̾̿̾̓̔̌t̵̯̘͇̙̬̤̫͚͓̠͔̭̚ͅ ̵͔̲͈̭͛̆̈Ú̷̱̈́͑́̇̀̈́̔͝͠s̴̢͇̳̯̬͔̦̗̯̜̀̓̎̆̅̈́̈́̑̾͆͌̆͝ê̴̛͎̪̩̙͕̖͖͍͓͐̊̈͋̆̂̌̍̂͘͘̚ ̵̢̢͖̦̬͔̻͇̠̺̂͒͌͌̈́̇̏͑͋̓̔̅̓͂͝S̴̡͇͎̠̯͙̖̲̮͓͋͊̐̋̑́͂͆̿C̵̤̲͉̫̜͓͇̙͇̪̥̈̂́̈́͋̃̿̕H̶̡͓̠̝͉͋͋̎̃E̶͉̬̞̖̮͚̙͗̔̋͆Ḑ̴̛͎̙̥̱̱͎͊͊͋̈́̆̎̓͑̚̚_̷͚̯̻̖͈̳̺̖̈͗́̽͠Ḱ̴̮͓̲̠̫͈̆̿͂͂̀̂͆̀́̔͝ͅḘ̷̡͔̘̒̒̐͆̍̀R̵̭̰̐́̾̈́̈́̋͗̎̉͐͝N̴̰͆̂͛̂̈́̄̊͌̓͆̕͘͠͠Ě̴̥͙̼̳̩̠͓͕̫͔̘̠̻̀͌̓͑́̓̂̋̌͘͜L̴̙͓̣̭̙͙͈̺̩̄̈͆͆͒͗͛̈́̎̉̈́͘͘͜͠_̵̭̱͇͍̩̬͌̿͂͊̐͘̕ͅP̵̡̪̹̀͂̀̋̇̓͘͜Ŕ̸̭͍̮̼̙͓͕̳̫̇̏͌͌̋͜͝Ǫ̵̛͎̭̱͈̫̮̘͉͔̪̟̺̥̉̈̈́̊́̕̕͠ͅC̵̡̡͕̱̠̝̭͓̗̜͇̼̳̲͓̑̽̔̑͠Ę̷̛̦̑̈̒̃̊̈́̓̚̚͜͠͝Ş̴̥͇̬̱̠̟͙̳̖̿͗̀͌͊̽̐͒̃͜S̴͉͖͉͎͎̱̝͔̙̜̭̍͑͊̒͆̃͗̂͌̽͑̓ͅ ̶̼̥̜͙͈̦̲̰̞̈́̑͆̂͝H̸̩̭̹̘̩̊͗̔̉̓̈́ȩ̷̩͍͚̮̰̭̘͍͇͈̺͇̻̓̀̅̊́̓̈́͗̇̔̓͛͜͝ṙ̶̡̩̜̜̭̝̪͔͙̟̣͉̆͗̉͜͠͝ͅę̷̡̦̘̰͖̲͉̤̼͈̺̺̦̯̿͌́͐̃̃͂̾́̊́.̶̡̧͈͍͈̟̠͙̘̤̝̪̽̊̈͐̕͜
  return sched_create("unknown program", prog->entry, prog->pm, SCHED_USER_PROCESS);
}
void sched_exit(int exit_code) {
  log("sched - Process %d exited with code %d!", curr_proc->pid, exit_code);
  curr_proc->type = SCHED_DIED;
  schedule(&curr_proc->regs);
}

void schedule(registers_t *regs) {
  if (standby) {
    // log("sched - Sched is in standby.\n");
    return;
  }

  memcpy(&curr_proc->regs, regs, sizeof(registers_t));

  if (curr_proc->type == SCHED_DIED) {
    sched_process *prev_proc = proc_list;
    while (prev_proc->next != curr_proc) {
      prev_proc = prev_proc->next;
    }

    prev_proc->next = curr_proc->next;

    vmm_release_pm(curr_proc->pm);
    pmm_free_page(curr_proc->stack_base_physical);

    // R.I.P. process
    pmm_free_page(curr_proc);

    return;
  }

  curr_proc = curr_proc->next;
  if (curr_proc == NULL)
    curr_proc = proc_list;

  memcpy(regs, &curr_proc->regs, sizeof(registers_t));

  wrmsr(IA32_GS_KERNEL_MSR, (uint64_t)curr_proc);
  // log("sched - proc %d\n", curr_proc->pid);
  vmm_load_pagemap(curr_proc->pm);
}