#include <amd64/apic.h>
#include <amd64/fx.h>
#include <amd64/gdt.h>
#include <amd64/intr.h>
#include <amd64/mm.h>
#include <amd64/msr-index.h>
#include <amd64/msr.h>
#include <libk/lengthof.h>
#include <libk/list.h>
#include <libk/string.h>
#include <mm/malloc.h>
#include <proc/proc.h>
#include <proc/reschedule.h>
#include <status.h>
#include <sys/debug.h>
#include <sys/intr.h>
#include <sys/smp.h>
#include <syscall/syscall.h>
#include <syscall_defs.h>

extern void syscall_entry(void);

uintptr_t syscall_dispatch(void* stack_ptr) {
  uint64_t fp, fc;

  load_kernel_cr3();
  struct saved_regs* regs = stack_ptr;

  spin_lock(&thiscpu->lock, &fc);
  struct proc* caller = thiscpu->proc_current;

  spin_lock(&caller->lock, &fp);

  int caller_pid = caller->pid;

  memcpy(&caller->pdata.regs, regs, sizeof(struct saved_regs));

  fx_save(caller->pdata.fx_env);

  spin_unlock(&caller->lock, fp);
  spin_unlock(&thiscpu->lock, fc);

  int syscall_num = regs->rax;
  syscall_handler_func_t func = syscall_find_handler(syscall_num);

  if (func == NULL) {
    return -ST_SYSCALL_NOT_FOUND;
  }

  struct reschedule_ctx rctx;
  memset(&rctx, 0, sizeof(rctx));

  lapic_timer_mask();

  intr_enable();

  uintptr_t r =
      func(caller, regs, &rctx, regs->rdi, regs->rsi, regs->rdx, regs->r10, regs->r8, regs->r9);

  intr_disable();

  lapic_timer_unmask();

  caller = proc_find_pid(caller_pid);

  if (caller != NULL) {
    spin_lock(&caller->lock, &fp);
    caller->pdata.regs.rax = r;
    spin_unlock(&caller->lock, fp);
  }

  bool do_thiscpu = false;
  for (size_t i = 0; i < lengthof(rctx.cpus); i++) {
    if (rctx.cpus[i] != NULL && rctx.cpus[i] != thiscpu)
      cpu_request_sched(rctx.cpus[i], true);
    else
      do_thiscpu = true;
  }

  if (do_thiscpu)
    cpu_request_sched(thiscpu, true);

  return r;
}

void syscall_init(void) {
  wrmsr(MSR_STAR, ((uint64_t)GDT_KCODE << 32) | ((uint64_t)(GDT_KDATA | 0x03) << 48));
  wrmsr(MSR_LSTAR, (uint64_t)&syscall_entry);
  wrmsr(MSR_SYSCALL_MASK, (1ULL << 9));
  wrmsr(MSR_EFER, rdmsr(MSR_EFER) | EFER_SCE);
}