#include <kernel/bsp.h>
#include <kernel/handle.h>
#include <kernel/printk.h>
#include <kernel/sched.h>
#include <kernel/util.h>
#include <kernel/vm-object.h>
#include <kernel/vm-region.h>

#define BOOTSTRAP_STACK_SIZE 0x10000

static struct boot_module bsp_location = {0};

void bsp_set_location(const struct boot_module *mod)
{
	memcpy(&bsp_location, mod, sizeof bsp_location);
}

void bsp_get_location(struct boot_module *out)
{
	memcpy(out, &bsp_location, sizeof bsp_location);
}

kern_status_t bsp_load(struct bsp *bsp, const struct boot_module *mod)
{
	size_t trailer_offset = mod->mod_size - sizeof(struct bsp_trailer);
	const char *p = vm_phys_to_virt(mod->mod_base);
	const struct bsp_trailer *trailer_bigendian
		= (const struct bsp_trailer *)(p + trailer_offset);

	bsp->bsp_trailer.bsp_magic
		= big_to_host_u32(trailer_bigendian->bsp_magic);
	bsp->bsp_trailer.bsp_fs_offset
		= big_to_host_u64(trailer_bigendian->bsp_fs_offset);
	bsp->bsp_trailer.bsp_fs_len
		= big_to_host_u32(trailer_bigendian->bsp_fs_len);
	bsp->bsp_trailer.bsp_exec_offset
		= big_to_host_u64(trailer_bigendian->bsp_exec_offset);
	bsp->bsp_trailer.bsp_exec_len
		= big_to_host_u32(trailer_bigendian->bsp_exec_len);
	bsp->bsp_trailer.bsp_text_faddr
		= big_to_host_u64(trailer_bigendian->bsp_text_faddr);
	bsp->bsp_trailer.bsp_text_vaddr
		= big_to_host_u64(trailer_bigendian->bsp_text_vaddr);
	bsp->bsp_trailer.bsp_text_size
		= big_to_host_u64(trailer_bigendian->bsp_text_size);
	bsp->bsp_trailer.bsp_data_faddr
		= big_to_host_u64(trailer_bigendian->bsp_data_faddr);
	bsp->bsp_trailer.bsp_data_vaddr
		= big_to_host_u64(trailer_bigendian->bsp_data_vaddr);
	bsp->bsp_trailer.bsp_data_size
		= big_to_host_u64(trailer_bigendian->bsp_data_size);
	bsp->bsp_trailer.bsp_exec_entry
		= big_to_host_u64(trailer_bigendian->bsp_exec_entry);

	if (bsp->bsp_trailer.bsp_magic != BSP_MAGIC) {
		return KERN_INVALID_ARGUMENT;
	}

	bsp->bsp_vmo = vm_object_create_in_place(
		"bsp",
		3,
		mod->mod_base,
		mod->mod_size,
		VM_PROT_READ | VM_PROT_EXEC | VM_PROT_USER);
	if (!bsp->bsp_vmo) {
		return KERN_INVALID_ARGUMENT;
	}

	return KERN_OK;
}

static kern_status_t map_executable_dyn(
	struct bsp *bsp,
	struct task *task,
	virt_addr_t *entry)
{
	kern_status_t status = KERN_OK;
	size_t exec_size = 0;
	if (bsp->bsp_trailer.bsp_text_vaddr > bsp->bsp_trailer.bsp_data_vaddr) {
		exec_size = bsp->bsp_trailer.bsp_text_vaddr
		          + bsp->bsp_trailer.bsp_text_size;
	} else {
		exec_size = bsp->bsp_trailer.bsp_data_vaddr
		          + bsp->bsp_trailer.bsp_data_size;
	}

	struct vm_region *region;
	status = vm_region_create(
		task->t_address_space,
		"exec",
		4,
		VM_REGION_ANY_OFFSET,
		exec_size,
		VM_PROT_READ | VM_PROT_WRITE | VM_PROT_EXEC | VM_PROT_USER,
		&region);
	if (status != KERN_OK) {
		return status;
	}

	struct vm_object *data = vm_object_create(
		".data",
		5,
		bsp->bsp_trailer.bsp_data_size,
		VM_PROT_READ | VM_PROT_WRITE | VM_PROT_USER);
	if (!data) {
		return KERN_NO_MEMORY;
	}

	virt_addr_t text_base = 0, data_base = 0;

	off_t text_foffset = bsp->bsp_trailer.bsp_exec_offset
	                   + bsp->bsp_trailer.bsp_text_faddr;
	off_t data_foffset = 0;
	off_t text_voffset = bsp->bsp_trailer.bsp_text_vaddr;
	off_t data_voffset = bsp->bsp_trailer.bsp_data_vaddr;

#if 0
	size_t tmp = 0;
	status = vm_object_copy(
		data,
		0,
		bsp->bsp_vmo,
		bsp->bsp_trailer.bsp_data_faddr,
		bsp->bsp_trailer.bsp_data_size,
		&tmp);

	tracek("read %zuB of data from executable", tmp);
#endif

	tracek("text_foffset=%06llx, data_foffset=%06llx",
	       text_foffset,
	       data_foffset);
	tracek("text_voffset=%08llx, data_voffset=%08llx",
	       text_voffset,
	       data_voffset);

	status = vm_region_map_object(
		region,
		text_voffset,
		bsp->bsp_vmo,
		text_foffset,
		bsp->bsp_trailer.bsp_text_size,
		VM_PROT_READ | VM_PROT_EXEC | VM_PROT_USER,
		&text_base);
	if (status != KERN_OK) {
		return status;
	}

	status = vm_region_map_object(
		region,
		data_voffset,
		data,
		data_foffset,
		bsp->bsp_trailer.bsp_data_size,
		VM_PROT_READ | VM_PROT_WRITE | VM_PROT_USER,
		&data_base);
	if (status != KERN_OK) {
		return status;
	}

	tracek("text_base=%08llx, data_base=%08llx", text_base, data_base);

	*entry = text_base + bsp->bsp_trailer.bsp_exec_entry;
	return KERN_OK;
}

static kern_status_t map_executable_exec(
	struct bsp *bsp,
	struct task *task,
	virt_addr_t *entry)
{
	kern_status_t status = KERN_OK;
	struct vm_object *data = vm_object_create(
		".data",
		5,
		bsp->bsp_trailer.bsp_data_size,
		VM_PROT_READ | VM_PROT_WRITE | VM_PROT_USER);
	if (!data) {
		return KERN_NO_MEMORY;
	}

	virt_addr_t text_base = 0, data_base = 0;

	off_t text_foffset = bsp->bsp_trailer.bsp_exec_offset
	                   + bsp->bsp_trailer.bsp_text_faddr;
	off_t data_foffset = 0;
	off_t text_voffset = bsp->bsp_trailer.bsp_text_vaddr;
	off_t data_voffset = bsp->bsp_trailer.bsp_data_vaddr;

	text_voffset -= vm_region_get_base_address(task->t_address_space);
	data_voffset -= vm_region_get_base_address(task->t_address_space);

#if 0
	size_t tmp = 0;
	status = vm_object_copy(
		data,
		0,
		bsp->bsp_vmo,
		bsp->bsp_trailer.bsp_data_faddr,
		bsp->bsp_trailer.bsp_data_size,
		&tmp);

	tracek("read %zuB of data from executable", tmp);
#endif

	tracek("text_foffset=%06llx, data_foffset=%06llx",
	       text_foffset,
	       data_foffset);
	tracek("text_voffset=%08llx, data_voffset=%08llx",
	       text_voffset,
	       data_voffset);

	status = vm_region_map_object(
		task->t_address_space,
		text_voffset,
		bsp->bsp_vmo,
		text_foffset,
		bsp->bsp_trailer.bsp_text_size,
		VM_PROT_READ | VM_PROT_EXEC | VM_PROT_USER,
		&text_base);
	if (status != KERN_OK) {
		return status;
	}

	status = vm_region_map_object(
		task->t_address_space,
		data_voffset,
		data,
		data_foffset,
		bsp->bsp_trailer.bsp_data_size,
		VM_PROT_READ | VM_PROT_WRITE | VM_PROT_USER,
		&data_base);
	if (status != KERN_OK) {
		return status;
	}

	tracek("text_base=%08llx, data_base=%08llx", text_base, data_base);

	*entry = bsp->bsp_trailer.bsp_exec_entry;
	return KERN_OK;
}

kern_status_t bsp_launch_async(struct bsp *bsp, struct task *task)
{
	virt_addr_t stack_buffer, bsp_data_base;
	virt_addr_t entry, sp;
	kern_status_t status;

	struct vm_object *user_stack = vm_object_create(
		"stack",
		5,
		BOOTSTRAP_STACK_SIZE,
		VM_PROT_READ | VM_PROT_WRITE | VM_PROT_USER);
	if (!user_stack) {
		return KERN_NO_ENTRY;
	}

	status = vm_region_map_object(
		task->t_address_space,
		VM_REGION_ANY_OFFSET,
		user_stack,
		0,
		BOOTSTRAP_STACK_SIZE,
		VM_PROT_READ | VM_PROT_WRITE | VM_PROT_USER,
		&stack_buffer);

	if (status != KERN_OK) {
		return status;
	}

	status = vm_region_map_object(
		task->t_address_space,
		VM_REGION_ANY_OFFSET,
		bsp->bsp_vmo,
		0,
		bsp->bsp_trailer.bsp_exec_offset,
		VM_PROT_READ | VM_PROT_USER,
		&bsp_data_base);

	if (status != KERN_OK) {
		return status;
	}

	status = map_executable_exec(bsp, task, &entry);
	if (status != KERN_OK) {
		return status;
	}
#ifdef TRACE
	vm_region_dump(task->t_address_space);
#endif

	sp = stack_buffer + BOOTSTRAP_STACK_SIZE;
	tracek("bootstrap: entry=%llx, sp=%llx", entry, sp);

	kern_handle_t self, self_address_space;
	task_open_handle(task, &task->t_base, 0, &self);
	task_open_handle(
		task,
		&task->t_address_space->vr_base,
		0,
		&self_address_space);

	const uintptr_t args[] = {
		0,                    // int argc
		0,                    // const char ** argv
		self,                 // kern_handle_t task
		self_address_space,   // kern_handle_t address_space

		/* this parameter is specific to the bsp bootstrap program, so
	         * that it can access the rest of the bsp image. */
		bsp_data_base,
	};
	const size_t nr_args = sizeof args / sizeof args[0];

	struct thread *init_thread = task_create_thread(task);
	thread_init_user(init_thread, entry, sp, args, nr_args);
	schedule_thread_on_cpu(init_thread);

	return KERN_OK;
}