#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 "sysdefs/proc.h"
#include "sysdefs/fs.h"
#include "time/time.h"
#include "pico_stack.h"

#define PROC_REAPER_FREQ 30

uint64_t pids = 0;
uint64_t sched_ticks = 0;

uint64_t uptime_ms = 0;
uint64_t last_secs = 0;

uint64_t get_wallclock_secs(void) {
  Time t;
  time_get(&t);
  return time_totalhours(&t) * 3600 + t.minute * 60 + t.second;
}

void update_time(void) {
  uint64_t now_s = get_wallclock_secs();
  uint64_t delta_s;

  if (now_s >= last_secs)
    delta_s = now_s - last_secs;
  else
    delta_s = 0;

  uptime_ms += (uint64_t)delta_s * 1000;
  last_secs = now_s;
}

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_spawnuser(char *mountpoint, char *path) {
  FsStat stat;
  if (vfs_stat(mountpoint, path, &stat) != E_OK) {
    return NULL;
  }

  if (stat.type != FSSTAT_FILE) {
    return NULL;
  }

  VfsObj *vobj = vfs_open(mountpoint, path, VFS_FLAG_READ);
  if (vobj == NULL) {
    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();

  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->pipes_spinlock);
  spinlock_init(&proc->devs_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);

  time_get(&proc->time);

  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 < PROC_VFSHANDLES_MAX; 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);
    
    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++;
  update_time();
  pico_stack_tick();

  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_init(void) {
  spinlock_init(&PROCS.spinlock);
  PROCS.procs = NULL;

  Proc *init = proc_spawnuser("base", "/bin/init");
  PROCS.current = init;
  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);
}
