#include <arch/msr.h>
#include <kernel/machine/cpu.h>
#include <kernel/machine/thread.h>
#include <kernel/thread.h>

#define MAX_REG_ARGS 6
#define REG_ARG_0    rdi
#define REG_ARG_1    rsi
#define REG_ARG_2    rdx
#define REG_ARG_3    rcx
#define REG_ARG_4    r8
#define REG_ARG_5    r9

/* this is the context information restored by ml_thread_switch.
 * since ml_thread_switch only jumps to kernel-mode, IRETQ isn't used,
 * and the extra register values needed by IRETQ aren't present. */
struct thread_ctx {
	uint64_t r15, r14, r13, r12, r11, r10, r9, r8;
	uint64_t rdi, rsi, rbp, unused_rsp, rbx, rdx, rcx, rax;
	uint64_t rfl;
} __packed;

void ml_thread_prepare_kernel_context(uintptr_t ip, uintptr_t *sp)
{
	(*sp) -= sizeof(uintptr_t);
	uintptr_t *dest_ip = (uintptr_t *)(*sp);
	*dest_ip = ip;

	(*sp) -= sizeof(struct thread_ctx);
	struct thread_ctx *ctx = (struct thread_ctx *)(*sp);
	memset(ctx, 0x0, sizeof *ctx);

	ctx->rfl = 0x202;
}

extern kern_status_t ml_thread_prepare_user_context(
	virt_addr_t ip,
	virt_addr_t user_sp,
	virt_addr_t *kernel_sp,
	const uintptr_t *args,
	size_t nr_args)
{
	(*kernel_sp) -= sizeof(struct ml_cpu_context);

	struct ml_cpu_context *ctx = (struct ml_cpu_context *)(*kernel_sp);
	memset(ctx, 0x0, sizeof *ctx);
	ctx->rip = ip;
	ctx->rsp = user_sp;
	ctx->ss = 0x1b;
	ctx->cs = 0x23;
	ctx->rflags = 0x202;
	ctx->rdi = 0;   // arg 0
	ctx->rsi = 0;   // arg 1

	for (size_t i = 0; i < nr_args; i++) {
		switch (i) {
		case 0:
			ctx->REG_ARG_0 = args[i];
			break;
		case 1:
			ctx->REG_ARG_1 = args[i];
			break;
		case 2:
			ctx->REG_ARG_2 = args[i];
			break;
		case 3:
			ctx->REG_ARG_3 = args[i];
			break;
		case 4:
			ctx->REG_ARG_4 = args[i];
			break;
		case 5:
			ctx->REG_ARG_5 = args[i];
			break;
		default:
			return KERN_INVALID_ARGUMENT;
		}
	}

	return KERN_OK;
}

kern_status_t ml_thread_config_get(
	struct thread *thread,
	kern_config_key_t key,
	void *out,
	size_t max)
{
	return KERN_OK;
}

kern_status_t ml_thread_config_set(
	struct thread *thread,
	kern_config_key_t key,
	const void *ptr,
	size_t len)
{
	switch (key) {
	case THREAD_CFG_FSBASE:
		if (len != sizeof(thread->tr_ml.tr_fsbase)) {
			return KERN_INVALID_ARGUMENT;
		}

		thread->tr_ml.tr_fsbase = *(virt_addr_t *)ptr;
		if (thread == current_thread()) {
			wrmsr(MSR_FS_BASE, thread->tr_ml.tr_fsbase);
		}

		break;
	case THREAD_CFG_GSBASE:
		if (len != sizeof(thread->tr_ml.tr_gsbase)) {
			return KERN_INVALID_ARGUMENT;
		}

		thread->tr_ml.tr_gsbase = *(virt_addr_t *)ptr;
		if (thread == current_thread()) {
			/* we're in the kernel right now, so the user and kernel
			 * gs-base registers are swapped. when we return to
			 * usermode, this value will be swapped back into
			 * the user gs-base register */
			wrmsr(MSR_KERNEL_GS_BASE, thread->tr_ml.tr_gsbase);
		}

		break;
	default:
		return KERN_INVALID_ARGUMENT;
	}

	return KERN_OK;
}