#include <stdint.h>
#include <stdlib.h>
#include "hal/hal.h"
#include "vmm/vmm.h"
#include "spinlock/spinlock.h"
#include "proc.h"
#include "dlmalloc/malloc.h"
#include "pmm/pmm.h"
#include "util/util.h"
#include "kprintf.h"

#define PROC_REAPER_FREQ 30

uint64_t pids = 0;
uint64_t sched_ticks = 0;

Procs PROCS;

Proc *proc_spawnkern(void (*ent)(void), char *name) {
  Proc *proc = dlmalloc(sizeof(*proc));
  if (proc == NULL) {
    return NULL;
  }
  hal_memset(proc, 0, sizeof(*proc));
  hal_memcpy(proc->name, name, PROC_NAME_MAX);

  proc->kern = true;

  uint8_t *sp = (uint8_t *)pmm_alloc(PROC_STACKBLOCKS) + PROC_STACKSIZE;

  #if defined(__x86_64__)
    proc->platformdata.kstack = sp;
    hal_memset(&proc->platformdata.trapframe, 0, sizeof(proc->platformdata.trapframe));

    proc->platformdata.trapframe.ss = 0x10;
    proc->platformdata.trapframe.rsp = (uint64_t)VIRT(sp);
    proc->platformdata.trapframe.rflags = 0x202;
    proc->platformdata.trapframe.cs = 0x08;
    proc->platformdata.trapframe.rip = (uint64_t)ent;
    proc->platformdata.cr3 = hal_vmm_current_cr3();
    proc->state = PROC_READY;
    proc->pid = pids++;
  #endif

  return proc;
}

void proc_register(Proc *proc) {
  spinlock_acquire(&PROCS.spinlock);
  LL_APPEND(PROCS.procs, proc);
  spinlock_release(&PROCS.spinlock);
}

Proc *proc_nextready(void) {
  Proc *next = PROCS.current->next;
  Proc *proc = next == NULL ? PROCS.procs : next;
  return proc;
}

void proc_sched(void *cpustate) {
  hal_intr_disable();
  sched_ticks++;

  if (sched_ticks % PROC_REAPER_FREQ == 0) {
    Proc *head = PROCS.procs;
    while (head) {
      if (head->state == PROC_ZOMBIE) {
        Proc *zombie = head;
        head = head->next;

        LL_REMOVE(PROCS.procs, zombie);
        pmm_free((uintptr_t)(zombie->platformdata.kstack - PROC_STACKSIZE), PROC_STACKBLOCKS);
        dlfree(zombie);
      } else {
        head = head->next;
      }
    }
  }

  #if defined(__x86_64__)
    IntrStackFrame *frame = cpustate;

    PROCS.current->platformdata.trapframe = *frame;
    PROCS.current->platformdata.fsbase = hal_cpu_rdmsr(HAL_CPU_FSBASE);
  #endif
    
  PROCS.current = proc_nextready();
  PROCS.current->state = PROC_RUNNING;

  #if defined(__x86_64__)
    hal_cpu_wrmsr(HAL_CPU_FSBASE, PROCS.current->platformdata.fsbase);
    hal_switchproc(&PROCS.current->platformdata.trapframe, (void *)PROCS.current->platformdata.cr3);
  #endif
}

void proc_kill(Proc *proc) {
  proc->state = PROC_ZOMBIE;
}

void proc_killself(void) {
  spinlock_acquire(&PROCS.spinlock);
  Proc *proc = PROCS.current;
  proc_kill(proc);
  spinlock_release(&PROCS.spinlock);
}

void proc_idle(void) {
  for(;;) {
    // .
  }
}

int a = 0;

void proc_print(void) {
  for (;;) {
    for (volatile size_t i = 0; i < 0x1000000; i++) {
    }
    
    if (a == 10) {
      proc_killself();
      for(;;);
    }
    kprintf("A");
    a++;
  }
}

int b = 0;

void proc_print2(void) {
  for (;;) {
    for (volatile size_t i = 0; i < 0x1000000; i++) {
    }
    
    if (b == 20) {
      proc_killself();
      for(;;);
    }
    kprintf("B");
    b++;
  }
}

void proc_init(void) {
  spinlock_init(&PROCS.spinlock);
  PROCS.procs = NULL;

  Proc *idle = proc_spawnkern(&proc_idle, "kIDLE");
  proc_register(idle);
  PROCS.current = idle;

  proc_register(proc_spawnkern(&proc_print, "print"));
  proc_register(proc_spawnkern(&proc_print2, "print"));

  #if defined(__x86_64__)
    hal_switchproc(&PROCS.current->platformdata.trapframe, (void *)PROCS.current->platformdata.cr3);
  #endif
}