#include <mango/bsp.h>
#include <mango/printk.h>
#include <mango/sched.h>
#include <mango/util.h>
#include <mango/vm-object.h>
#include <mango/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",
		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(
	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",
		VM_REGION_ANY_MAP_ADDRESS,
		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",
		bsp->bsp_trailer.bsp_data_size,
		VM_PROT_READ | VM_PROT_WRITE | VM_PROT_USER);
	/* TODO copy .data from executable to memory */
	if (!data) {
		return KERN_NO_MEMORY;
	}

	off_t text_offset = bsp->bsp_trailer.bsp_exec_offset
	                  + bsp->bsp_trailer.bsp_text_faddr;
	off_t data_offset = 0;
	virt_addr_t text_base = region->vr_entry.e_base_address
	                      + bsp->bsp_trailer.bsp_text_vaddr;
	virt_addr_t data_base = region->vr_entry.e_base_address
	                      + bsp->bsp_trailer.bsp_data_vaddr;

	tracek("exec_offset=%llx, text_faddr=%llx",
	       bsp->bsp_trailer.bsp_exec_offset,
	       bsp->bsp_trailer.bsp_text_faddr);
	tracek("text_offset=%llx, data_offset=%llx", text_offset, data_offset);
	tracek("text_base=%llx, data_base=%llx", text_base, data_base);

	status = vm_region_map_object(
		region,
		text_base,
		bsp->bsp_vmo,
		text_offset,
		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_base,
		data,
		data_offset,
		bsp->bsp_trailer.bsp_data_size,
		VM_PROT_READ | VM_PROT_WRITE | VM_PROT_USER,
		&data_base);
	if (status != KERN_OK) {
		return status;
	}

	*entry = text_base + 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;
	virt_addr_t entry, sp;
	kern_status_t status = map_executable(bsp, task, &entry);
	if (status != KERN_OK) {
		return status;
	}

	struct vm_object *user_stack = vm_object_create(
		"stack",
		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_MAP_ADDRESS,
		user_stack,
		0,
		BOOTSTRAP_STACK_SIZE,
		VM_PROT_READ | VM_PROT_WRITE | VM_PROT_USER,
		&stack_buffer);

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

#ifdef TRACE
	vm_region_dump(task->t_address_space, 0);
#endif

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

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

	return KERN_OK;
}