#include "kjarna.h"
#include <stdnoreturn.h>
#include <asm/x86_64/msr.h>
#include <bits/x86_64/msr.h>
#include <bits/x86_64/descriptor.h>
#include <libc/stdio.h>
#include <posix/sys/mman.h>
#include <posix/unistd.h>

/*
 * Fake syscall mechanism -
 *
 * We have a GDT with entries for supervisory mode. These serve as the
 * entries to satisfy the requirements of the SYSCALL instruction, as
 * for some reason (possibly intentionally) the OVMF GDT is not laid out
 * in a way to make use of the SYSCALL instruction possible.
 *
 * This causes us to have to work inside of constraints during loading time
 * - All "user" mode execution entirely blocks interrupt processing. That
 *   means that "user" mode code must not execute "hlt", or the system will
 *   be locked.
 * - When (if?) user input is required, it is always buffered. It is possible
 *   to simulate unbuffered input, at the cost of one syscall per transfer
 *   from "kernel" side to "user" side. This will cause high input latency.
 *   We will not be running Quake in this environment.
 *
 * These constraints are probably fine, as the loading process only needs to 
 * open files, map memory, etc.
 */


void *stack_alloc(void **stack_pointer, size_t alloc_size)
{
	void *block = *(char **)stack_pointer -= alloc_size;

	// maintain alignment
	*(char **)stack_pointer -= alloc_size % sizeof(size_t);

	return block;
}

static struct segment_descriptor const fake_syscall_gdt[] =
{
	{ 0 },
	{ 0xffff, 0, 0, 0x9a, 0xaf, 0},
	{ 0xffff, 0, 0, 0x92, 0xcf, 0}
};

typedef uint64_t syscall_delegate_parameters[6];

static struct
{
	uint16_t cs;
	uint16_t ds;
	uint64_t rflags;
	struct descriptor_table_register_long gdtr;
}
uefi_context;

static void save_uefi_context()
{
	__asm__ volatile (
			"movw %%cs, %0\n"
			"movw %%ds, %1\n"
			"pushfq\n"
			"popq %3\n"
			"sgdt %2\n"
			:
				"=m"(uefi_context.cs),
				"=m"(uefi_context.ds),
				"=m"(uefi_context.gdtr),
				"=g"(uefi_context.rflags)
			:
		);
}

static void restore_uefi_context()
{
	__asm__ volatile (
			"lgdt %0\n"
			"mov %w1, %%ds\n"
			"mov %w1, %%es\n"
			"mov %w1, %%fs\n"
			"mov %w1, %%gs\n"
			"mov %w1, %%gs\n"
			"mov %w1, %%ss\n"
			"push %q2\n"
			"lea .Lflush(%%rip), %%rax\n"
			"push %%rax\n"
			"lretq\n"
			".Lflush:\n"
			"push %q3\n"
			"popfq\n"
			:
			:
				"m"(uefi_context.gdtr),
				"m"(uefi_context.ds),
				"m"(uefi_context.cs),
				"g"(uefi_context.rflags)
			:
				"rax"
			);

}

struct syscall_context_return_state
{
	uint64_t rflags;
	struct descriptor_table_register_long gdtr;
	uint64_t return_ds;
	uintptr_t return_rip;
	uint64_t return_cs;
};

struct syscall_context
{
	uint64_t pad;
	uint64_t syscall_index;
	syscall_delegate_parameters parameters;
	struct syscall_context_return_state state;
};


typedef int64_t (syscall_delegate)(syscall_delegate_parameters params);

int64_t syscall_delegate_write(syscall_delegate_parameters params)
{
	return write((int)params[0], (void *)params[1], (size_t)params[2]);
}

constexpr int NR_SYSCALLS = 10;
static syscall_delegate *syscall_handler_delegates[NR_SYSCALLS] =
{
	nullptr,
	syscall_delegate_write
};

SYSV_ABI int64_t fake_syscall_handler(struct syscall_context *context)
{
	if (context->syscall_index > NR_SYSCALLS)
	{
		// TODO: ENOSYS
		return -1;
	}

	syscall_delegate *delegate = syscall_handler_delegates[context->syscall_index];

	if (delegate == nullptr)
	{
		// TODO: ENOSYS
		return -1;
	}

	restore_uefi_context();
	int64_t result = delegate(context->parameters);
	save_uefi_context();
	
	return result;
}

SYSV_ABI void fake_syscall_entry(void);
SYSV_ABI void fake_syscall_start(void *stack);

static void install_syscall_handler(void)
{
	union msr_lstar lstar = { (uintptr_t)fake_syscall_entry };
	union msr_star star = { { 0, 1 << 3, 1 << 3 | 3 } };

	msr_write(MSR_INDEX_LSTAR, lstar.value);
	msr_write(MSR_INDEX_STAR, star.value);

	uint64_t efer = msr_read(MSR_INDEX_EFER);
	efer |= 1;
	msr_write(MSR_INDEX_EFER, efer);
}

static void *create_fake_syscall_stack(void *entry_addr)
{
	constexpr size_t fake_syscall_stack_size = 0x20000; // 128KiB stack
	char *stack_pointer_base = mmap(nullptr, fake_syscall_stack_size, 0, 0, -1, 0);
	void *stack_pointer_head = stack_pointer_base + fake_syscall_stack_size;

	struct syscall_context *context = stack_alloc(&stack_pointer_head, sizeof(*context));

	context->state.rflags = 0;
	context->state.gdtr.base = (uintptr_t)&fake_syscall_gdt;
	context->state.gdtr.limit = sizeof(fake_syscall_gdt) - 1;
	context->state.return_ds = 16;
	context->state.return_cs = 8;
	context->state.return_rip = (uintptr_t)entry_addr;

	return stack_pointer_head;
}

int main(int argc, char **argv)
{
	(void)argc;
	(void)argv;

	struct kjarna_boot_image boot_image = get_boot_image();
	install_syscall_handler();
	void *fake_syscall_stack = create_fake_syscall_stack((void *)boot_image.entry);
	save_uefi_context();
	fake_syscall_start(fake_syscall_stack);
}