#include "src/kernel/syscall.h"
#include "src/kernel/fs.h"
#include "src/kernel/gdt.h"
#include "src/kernel/path.h"
#include "src/kernel/process.h"
#include "src/kernel/stream.h"
#include "src/kernel/tss.h"
#include "src/lib/layout.h"
#include "src/lib/memory.h"
#include "src/lib/string.h"
#include "src/lib/util.h"

#define MSR_EFER 0xC0000080
#define MSR_STAR 0xC0000081
#define MSR_LSTAR 0xC0000082
#define MSR_FMASK 0xC0000084

static inline uint64_t rdmsr(uint32_t msr) {
  uint32_t lo, hi;
  __asm__ volatile("rdmsr" : "=a"(lo), "=d"(hi) : "c"(msr));
  return ((uint64_t)hi << 32) | lo;
}

static inline void wrmsr(uint32_t msr, uint64_t value) {
  __asm__ volatile("wrmsr" : : "c"(msr), "a"((uint32_t)value), "d"((uint32_t)(value >> 32)));
}

extern void syscall_entry();

void syscall_init() {
  wrmsr(MSR_EFER, rdmsr(MSR_EFER) | 0x01);
  wrmsr(MSR_STAR, ((uint64_t)KERNEL_DS << 48) | ((uint64_t)KERNEL_CS << 32));
  wrmsr(MSR_LSTAR, (uint64_t)syscall_entry);
  wrmsr(MSR_FMASK, 0x200);
}

static uint64_t read(uint64_t fd, uint64_t max, char *to) {
  if (fd >= MAX_FDS || !current_process->fds[fd]) {
    return (uint64_t)-1;
  }
  return current_process->fds[fd]->read(current_process->fds[fd], max, to);
}

static uint64_t write(uint64_t fd, const char *from, uint64_t bytes) {
  if (fd >= MAX_FDS || !current_process->fds[fd]) {
    return (uint64_t)-1;
  }
  return current_process->fds[fd]->write(current_process->fds[fd], from, bytes);
}

static uint64_t getcwd(uint64_t max, char *to) {
  if (max == 0) {
    return 0;
  }

  if (max == 1) {
    to[0] = '\0';
    return 0;
  }

  if (max == 2 || !current_process->cwd->depth) {
    to[0] = '/';
    to[1] = '\0';
    return 1;
  }

  uint64_t i = 0;
  for (uint8_t j = 0; j < current_process->cwd->depth; j++) {
    uint64_t l = string_length(current_process->cwd->stack[j]->name);
    if (i + l + 1 > max) {
      break;
    }
    to[i++] = '/';
    memory_copy(current_process->cwd->stack[j]->name, l, to + i);
    i += l;
  }

  to[i] = '\0';
  return i;
}

static uint64_t chdir(const char *path) {
  path_t *cwd = path_open(current_process->root, current_process->cwd, path);
  if (!cwd) {
    return (uint64_t)-1;
  }
  if (!cwd->depth || !cwd->stack[cwd->depth - 1]->is_dir) {
    path_close(cwd);
    return (uint64_t)-1;
  }

  path_close((path_t *)current_process->cwd);
  current_process->cwd = cwd;
  return current_process->cwd->depth;
}

static uint64_t spawn(const char *path, uint64_t argc, const char **argv) {
  fs_node_t *node = path_open_file(current_process->root, current_process->cwd, path);
  if (!node) {
    return EXIT_CODE_NOT_FOUND;
  }
  uint8_t *bin = memory_allocate(node->size);
  fs_read(node, 0, node->size, bin);

  uint64_t size = 0;
  uint64_t sizes[16];
  uint64_t offsets[16];
  for (uint64_t i = 0; i < argc; i++) {
    offsets[i] = size;
    size += (sizes[i] = string_length(argv[i]) + 1);
  }
  char *blob = memory_allocate(size);
  for (uint64_t i = 0; i < argc; i++) {
    memory_copy(argv[i], sizes[i], blob + offsets[i]);
  }

  process_t *child = process_create(current_process, bin, node->size);

  memory_free(bin);
  fs_close(node);

  __asm__ volatile("mov %0, %%cr3" : : "r"((uint64_t)VIRT_TO_PHYS((void *)KERNEL_VIRTUAL_PML4)) : "memory");

  uint8_t *stack = (uint8_t *)child->user_stack;

  stack -= size;
  memory_copy(blob, size, stack);
  memory_free(blob);

  stack = (uint8_t *)((uint64_t)stack & ~7ULL);

  stack -= argc * sizeof(char *);
  for (uint64_t i = 0; i < argc; i++) {
    ((char **)stack)[i] = (char *)(USER_VIRTUAL_STACK_TOP - size + offsets[i]);
  }
  stack -= sizeof(uint64_t);
  *(uint64_t *)stack = argc;

  child->user_rsp = (void *)(USER_VIRTUAL_STACK_TOP - ((uint64_t)child->user_stack - (uint64_t)stack));

  __asm__ volatile("mov %0, %%cr3" : : "r"((uint64_t)VIRT_TO_PHYS((void *)current_process->pml4)) : "memory");

  return child->pid;
}

static exit_code_t wait(uint64_t pid) {
  process_t *child = process_get(pid);
  if (!child) {
    return EXIT_CODE_GENERAL_FAILURE;
  }

  current_process->state = PROCESS_WAITING;
  current_process->waiting_for = pid;
  while (child->state != PROCESS_ZOMBIE) {
    process_next();
  }
  current_process->state = PROCESS_RUNNING;

  exit_code_t code = child->code;
  process_destroy(child);

  return code;
}

static uint64_t open(const char *path) {
  if (current_process->free_fd >= MAX_FDS) {
    return (uint64_t)-1;
  }
  stream_t *stream = path_open_stream(current_process->root, current_process->cwd, path);
  if (!stream) {
    return (uint64_t)-1;
  }
  current_process->fds[current_process->free_fd++] = stream;
  return current_process->free_fd - 1;
}

static uint64_t close(uint64_t fd) {
  if (fd >= MAX_FDS || !current_process->fds[fd]) {
    return (uint64_t)-1;
  }
  current_process->fds[fd]->close(current_process->fds[fd]);
  current_process->fds[fd] = NUL;
  return 0;
}

static void exit(exit_code_t code) {
  current_process->state = PROCESS_ZOMBIE;
  current_process->code = code;

  if (current_process->parent) {
    process_t *parent = (process_t *)current_process->parent;
    if (parent->state == PROCESS_WAITING && parent->waiting_for == current_process->pid) {
      parent->state = PROCESS_RUNNING;
    }
  }

  process_next();
}

uint64_t syscall_dispatch(uint64_t func, uint64_t arg1, uint64_t arg2, uint64_t arg3, __attribute__((unused)) uint64_t arg4) {
  switch (func) {
  case SYSCALL_READ:
    return read(arg1, arg2, (char *)arg3);
  case SYSCALL_WRITE:
    return write(arg1, (const char *)arg2, arg3);
  case SYSCALL_GETCWD:
    return getcwd(arg1, (char *)arg2);
  case SYSCALL_CHDIR:
    return chdir((const char *)arg1);
  case SYSCALL_SPAWN:
    return spawn((const char *)arg1, arg2, (const char **)arg3);
  case SYSCALL_WAIT:
    return wait(arg1);
  case SYSCALL_OPEN:
    return open((const char *)arg1);
  case SYSCALL_CLOSE:
    return close(arg1);
  case SYSCALL_EXIT:
    exit(arg1);
    return 0;
  default:
    return (uint64_t)-1;
  }
}