#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/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) {
  switch (fd) {
  case 0:
    return current_process->stdin->read(current_process->stdin, max, to);
  case 1:
    return current_process->stdout->read(current_process->stdout, max, to);
  default:
    if (!current_process->fds[fd].node) {
      return (uint64_t)-1;
    }
    if (current_process->fds[fd].offset >= current_process->fds[fd].node->size) {
      return 0;
    }
    uint64_t remaining = current_process->fds[fd].node->size - current_process->fds[fd].offset;
    uint64_t to_read = remaining > max ? max : remaining;
    fs_read(current_process->fds[fd].node, current_process->fds[fd].offset, to_read, to);
    current_process->fds[fd].offset += to_read;
    return to_read;
  }
  return (uint64_t)-1;
}

static uint64_t write(uint64_t fd, const char *from, uint64_t bytes) {
  switch (fd) {
  case 0:
    return current_process->stdin->write(current_process->stdin, from, bytes);
  case 1:
    return current_process->stdout->write(current_process->stdout, from, bytes);
  default:
    return (uint64_t)-1;
  }
  return (uint64_t)-1;
}

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 (uint64_t)-1;
  }
  uint8_t *code = memory_allocate(node->size);
  fs_read(node, 0, node->size, code);

  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, code, node->size);

  memory_free(code);
  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));

  process_t *parent = current_process;
  current_process = child;
  tss.rsp0 = (uint64_t)child->kernel_stack;

  process_switch_to(&parent->kernel_rsp, child->kernel_rsp, VIRT_TO_PHYS(child->pml4));

  tss.rsp0 = (uint64_t)parent->kernel_stack;
  current_process = parent;
  process_destroy(child);

  return 0;
}

static uint64_t readdir(const char *path, uint64_t index, uint64_t max, char *to) {
  fs_node_t *dir = path_open_directory(current_process->root, current_process->cwd, path);
  if (!dir) {
    return (uint64_t)-1;
  }
  fs_node_t *item = fs_open_at(dir, index);
  if (!item) {
    if (dir != current_process->root) {
      fs_close(dir);
    }
    return 0;
  }
  uint64_t length = string_length(item->name);
  uint64_t size = length > max - 1 ? max - 1 : length;
  memory_copy(item->name, size, to);
  to[size] = '\0';

  fs_close(item);
  if (dir != current_process->root) {
    fs_close(dir);
  }

  return size;
}

static uint64_t open(const char *path) {
  if (current_process->free_fd >= MAX_FDS) {
    return (uint64_t)-1;
  }
  fs_node_t *node = path_open_node(current_process->root, current_process->cwd, path);
  if (!node) {
    return (uint64_t)-1;
  }
  current_process->fds[current_process->free_fd].node = node;
  current_process->fds[current_process->free_fd].offset++;
  current_process->free_fd++;
  return current_process->free_fd - 1;
}

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

static void exit() {
  process_switch_to(&current_process->kernel_rsp, current_process->parent->kernel_rsp, VIRT_TO_PHYS(current_process->parent->pml4));
}

uint64_t syscall_dispatch(uint64_t func, uint64_t arg1, uint64_t arg2, uint64_t arg3, 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_READDIR:
    return readdir((const char *)arg1, arg2, arg3, (char *)arg4);
  case SYSCALL_OPEN:
    return open((const char *)arg1);
  case SYSCALL_CLOSE:
    return close(arg1);
  case SYSCALL_EXIT:
    exit();
    return 0;
  default:
    return (uint64_t)-1;
  }
}