#include <stdint.h>
#include <stdlib.h>
#include <stdbool.h>
#include "hal/hal.h"
#include "spinlock/spinlock.h"
#include "proc.h"
#include "dlmalloc/malloc.h"
#include "pmm/pmm.h"
#include "util/util.h"
#include "kprintf.h"
#include "elf.h"
#include "errors.h"
#include "vfs/vfs.h"
#include "bootinfo/bootinfo.h"
#include "ipc/pipe/pipe.h"
#include "kproc/kproc.h"
#include "ps2kbproc/ps2kbproc.h"
#include "termproc/termproc.h"
#include "serialproc/serialproc.h"
#include "sysdefs/processctl.h"

#define PROC_REAPER_FREQ 30

uint64_t pids = 0;
uint64_t sched_ticks = 0;

Procs PROCS;

bool proc_checkelf(uint8_t *elf) {
  if (elf[0] != 0x7f || elf[1] != 'E' || elf[2] != 'L' || elf[3] != 'F') {
    return false;
  }
  return true;
}

ElfAuxval proc_load_elf_segs(Proc *proc, uint8_t *data) {
  ElfAuxval aux = {0};

  Elf64_Ehdr *elfhdr = (Elf64_Ehdr *)data;
  aux.entry = elfhdr->e_entry;
  aux.phnum = elfhdr->e_phnum;
  aux.phent = elfhdr->e_phentsize;

  for (uint64_t seg = 0; seg < elfhdr->e_phnum; seg++) {
    Elf64_Phdr *phdr = (Elf64_Phdr *)(data + elfhdr->e_phoff + (elfhdr->e_phentsize * seg));

    switch (phdr->p_type) {
      case PT_PHDR: {
        aux.phdr = (uint64_t)phdr->p_vaddr;
      } break;
      case PT_LOAD: {
        uint64_t off = phdr->p_vaddr & (HAL_PAGE_SIZE - 1);
        uint64_t blocks = (off + phdr->p_memsz + HAL_PAGE_SIZE - 1) / HAL_PAGE_SIZE;

        uint8_t *physaddr = pmm_alloc(blocks);
        uint8_t *virtaddr = (uint8_t *)(phdr->p_vaddr & ~(HAL_PAGE_SIZE - 1));

        hal_memset(VIRT(physaddr), 0, blocks * HAL_PAGE_SIZE);
        hal_memcpy(VIRT(physaddr) + off, (data + phdr->p_offset), phdr->p_filesz);

        uint32_t pgflags = HAL_PG_USER | HAL_PG_PRESENT;
        if (phdr->p_flags & PF_W) {
          pgflags |= HAL_PG_RW;
        }
        hal_vmm_map_range(proc->platformdata.cr3, virtaddr, physaddr, blocks * HAL_PAGE_SIZE, pgflags);

        VasRange *range = dlmalloc(sizeof(*range));
        range->virtstart = virtaddr;
        range->physstart = physaddr;
        range->size = blocks * HAL_PAGE_SIZE;
        range->pgflags = pgflags;

        LL_APPEND(proc->vas, range);
      } break;
    }
  }
  return aux;
}

Proc *proc_spawnkern(void (*ent)(void), char *name) {
  if (pids >= PROC_MAX) {
    return NULL;
  }

  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 *pstackp = (uint8_t *)pmm_alloc(PROC_STACKBLOCKS) + PROC_STACKSIZE;
  uint8_t *kstackp = (uint8_t *)pmm_alloc(PROC_STACKBLOCKS) + PROC_STACKSIZE;
  proc->platformdata.kstack = kstackp;
  proc->platformdata.pstack = pstackp;

  hal_memset(&proc->platformdata.trapframe, 0, sizeof(proc->platformdata.trapframe));

  proc->platformdata.trapframe.ss = 0x10;
  proc->platformdata.trapframe.rsp = (uint64_t)VIRT(pstackp);
  proc->platformdata.trapframe.rflags = 0x202;
  proc->platformdata.trapframe.cs = 0x08;
  proc->platformdata.trapframe.rip = (uint64_t)ent;
  proc->platformdata.cr3 = KERNEL_CR3;
  proc->state = PROC_EMBRYO;
  proc->pid = pids++;
  spinlock_init(&proc->bcast_pipes.spinlock);
  spinlock_init(&proc->pipes_spinlock);

  return proc;
}

Proc *proc_spawnuser(char *mountpoint, char *path) {
  VfsObj *vobj = vfs_open(mountpoint, path, VFS_FLAG_READ);
  if (vobj == NULL) {
    return NULL;
  }

  VfsStat stat;
  if (vobj->stat(vobj, &stat) != E_OK) {
    vfs_close(vobj);
    return NULL;
  }

  if (stat.type != VFS_TYPE_FILE) {
    vfs_close(vobj);
    return NULL;
  }

  uint8_t *data = dlmalloc(stat.size);
  if (data == NULL) {
    vfs_close(vobj);
    return NULL;
  }

  if (vobj->read(vobj, data, stat.size, 0) != E_OK) {
    dlfree(data);
    vfs_close(vobj);
    return NULL;
  }

  vfs_close(vobj);

  Proc *proc = dlmalloc(sizeof(*proc));
  hal_memset(proc, 0, sizeof(*proc));
  ksprintf(proc->name, "%s:%s", mountpoint, path);
  
  hal_memset(&proc->platformdata.trapframe, 0, sizeof(proc->platformdata.trapframe));
  proc->platformdata.cr3 = hal_vmm_userproc_pml4_phys(proc);

  uint8_t *kstackp = (uint8_t *)pmm_alloc(PROC_STACKBLOCKS) + PROC_STACKSIZE;
  proc->platformdata.kstack = kstackp;
  uint8_t *pstackp = (uint8_t *)pmm_alloc(PROC_STACKBLOCKS);
  proc->platformdata.pstack = pstackp;

  uint64_t virttop = PROC_USER_STACK_TOP;
  uint64_t virtbase = virttop - PROC_STACKSIZE;

  uint32_t flags = HAL_PG_RW | HAL_PG_PRESENT | HAL_PG_USER;
  hal_vmm_map_range(proc->platformdata.cr3, (void *)virtbase, (void *)pstackp, PROC_STACKSIZE, flags);

  VasRange *range = dlmalloc(sizeof(*range));
  range->virtstart = (uint8_t *)virtbase;
  range->physstart = (uint8_t *)pstackp;
  range->size = PROC_STACKSIZE;
  range->pgflags = flags;
  LL_APPEND(proc->vas, range);

  ElfAuxval aux = proc_load_elf_segs(proc, data);
  dlfree(data);

  proc->platformdata.trapframe.ss = 0x20 | 0x3;
  proc->platformdata.trapframe.rsp = (uint64_t)virttop;
  proc->platformdata.trapframe.rflags = 0x202;
  proc->platformdata.trapframe.cs = 0x18 | 0x3;
  proc->platformdata.trapframe.rip = aux.entry;
  proc->state = PROC_EMBRYO;
  proc->pid = pids++;
  proc->mman_map_base = PROC_MMAN_MAP_BASE;
  spinlock_init(&proc->bcast_pipes.spinlock);
  spinlock_init(&proc->pipes_spinlock);

  proc->pipes[0] = dlmalloc(sizeof(IpcPipe));
  ipc_pipeinit(proc->pipes[0], proc->pid);
  proc->pipes[1] = dlmalloc(sizeof(IpcPipe));
  ipc_pipeinit(proc->pipes[1], proc->pid);

  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 *proc = PROCS.current->next;
  for (;;) {
    if (proc == NULL) {
      proc = PROCS.procs;
    }
    if (proc->state == PROC_READY) {
      return proc;
    }
    proc = proc->next;
  }
}

void proc_reaper(void) {
  spinlock_acquire(&PROCS.spinlock);
  Proc *head, *tmp;
  LL_FOREACH_SAFE(PROCS.procs, head, tmp) {
    if (head->state != PROC_ZOMBIE) {
      continue;
    }

    Proc *zombie = head;
    LL_REMOVE(PROCS.procs, zombie);

    for (size_t i = 0; i < zombie->vobjcnt; i++) {
      if (zombie->vobjs[i] != NULL) {
        vfs_close(zombie->vobjs[i]);
        zombie->vobjs[i] = NULL;
      }
    }

    for (size_t i = 0; i < PROC_PIPEHANDLES_MAX; i++) {
      if (zombie->pipes[i] != NULL && zombie->pipes[i]->ownerpid == zombie->pid) {
        dlfree(zombie->pipes[i]);
        ipc_pipefree(zombie->pipes[i]);
        zombie->pipes[i] = NULL;
      }
    }

    pmm_free((uintptr_t)(zombie->platformdata.kstack - PROC_STACKSIZE), PROC_STACKBLOCKS);
    pmm_free((uintptr_t)(zombie->platformdata.pstack - PROC_STACKSIZE), PROC_STACKBLOCKS);
    
    if (!zombie->kern) {
      VasRange *vas, *vastmp;
      LL_FOREACH_SAFE(zombie->vas, vas, vastmp) {
        hal_vmm_unmap_range(zombie->platformdata.cr3, vas->virtstart, vas->physstart, vas->size);
        pmm_free((uintptr_t)vas->physstart, vas->size / HAL_PAGE_SIZE);
        dlfree(vas);
      }

      pmm_free((uintptr_t)zombie->platformdata.cr3, 1);

      ProcArg *arg, *argtmp;
      LL_FOREACH_SAFE(zombie->procargs.list, arg, argtmp) {
        dlfree(arg->string);
        dlfree(arg);
      }
    }
    dlfree(zombie);
  }
  spinlock_release(&PROCS.spinlock);
}

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

  IntrStackFrame *frame = cpustate;

  PROCS.current->platformdata.trapframe = *frame;
  
  PROCS.current = proc_nextready();
  
  if (sched_ticks % PROC_REAPER_FREQ == 0) {
    proc_reaper();
  }

  tss.rsp0 = (uint64_t)VIRT(PROCS.current->platformdata.kstack);
  hal_switchproc(&PROCS.current->platformdata.trapframe, (void *)PROCS.current->platformdata.cr3);
}

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_status(void) {
  static const char *statuses[] = {"ready", "running", "zombie", "waiting"};
  for (;;) {
    spinlock_acquire(&PROCS.spinlock);
    Proc *head = PROCS.procs;
    while (head) {
      kprintf("%s %s %s\n", head->kern ? "kern" : "user", statuses[head->state], head->name);
      head = head->next;
    }
    kprintf("\n\n");
    spinlock_release(&PROCS.spinlock);
  
    hal_wait(3 * 1000);
  }
}

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

  kproc_init(proc_spawnkern(&kproc_fn, "kproc"));
  proc_register(KPROC);
  PROCS.current = KPROC;
  KPROC->state = PROC_READY;

  ps2kbproc_init(proc_spawnkern(&ps2kbproc_fn, "ps2kbproc"));
  proc_register(PS2KBPROC);
  PS2KBPROC->state = PROC_READY;

  termproc_init(proc_spawnkern(&termproc_fn, "termproc"));
  proc_register(TERMPROC);
  TERMPROC->state = PROC_READY;

  /* serialproc_init(proc_spawnkern(&serialproc_fn, "serialproc")); */
  /* proc_register(SERIALPROC); */

  Proc *init = proc_spawnuser("base", "/bin/init");
  ipc_pipefree(init->pipes[0]);
  dlfree(init->pipes[0]);
  init->pipes[0] = TERMPROC->pipes[1];
  proc_register(init);
  init->state = PROC_READY;

  tss.rsp0 = (uint64_t)VIRT(PROCS.current->platformdata.kstack);
  hal_switchproc(&PROCS.current->platformdata.trapframe, (void *)PROCS.current->platformdata.cr3);
}