#include "elf.h"
#include "stack.h"

#include <launch.h>
#include <mango/handle.h>
#include <mango/log.h>
#include <mango/task.h>
#include <mango/vm.h>
#include <rosetta/bootstrap.h>
#include <stdio.h>
#include <string.h>

#define STACK_SIZE 0x10000

enum launch_status launch_ctx_init(struct launch_ctx *ctx)
{
	memset(ctx, 0x0, sizeof *ctx);
	return LAUNCH_OK;
}

void launch_ctx_cleanup(struct launch_ctx *ctx)
{
}

static kern_handle_t get_library(
	struct launch_ctx *ctx,
	const char *name,
	void *arg)
{
	enum launch_status status = LAUNCH_ERR_CANNOT_RESOLVE_DEPENDENCY;
	kern_handle_t result = KERN_HANDLE_INVALID;
	if (ctx->ctx_resolve_library) {
		status = ctx->ctx_resolve_library(ctx, name, &result, arg);
	}

	return result;
}

static virt_addr_t write_bootstrap_data(
	struct stack_writer *stack,
	const struct launch_parameters *params)
{
	virt_addr_t bs_remote;
	struct rosetta_bootstrap *bs
		= stack_writer_put(stack, NULL, sizeof *bs, &bs_remote);

	memset(bs, 0x0, sizeof *bs);

	bs->bs_argc = params->p_argc;
	bs->bs_envc = params->p_envc;
	bs->bs_handles_count = params->p_handle_count;
	bs->bs_channels_count = params->p_channel_count;

	const char **argv, **envp;

	if (bs->bs_argc > 0) {
		virt_addr_t remote_argv;
		argv = stack_writer_put(
			stack,
			NULL,
			bs->bs_argc * sizeof(char *),
			&remote_argv);
		bs->bs_argv = (const char **)remote_argv;
	}

	if (bs->bs_envc > 0) {
		virt_addr_t remote_envp;
		envp = stack_writer_put(
			stack,
			NULL,
			bs->bs_envc * sizeof(char *),
			&remote_envp);
		bs->bs_envp = (const char **)remote_envp;
	}

	for (size_t i = 0; i < params->p_argc; i++) {
		virt_addr_t arg_ptr;
		stack_writer_put_string(stack, params->p_argv[i], &arg_ptr);
		argv[i] = (const char *)arg_ptr;
	}

	for (size_t i = 0; i < params->p_envc; i++) {
		virt_addr_t env_ptr;
		stack_writer_put_string(stack, params->p_envp[i], &env_ptr);
		envp[i] = (const char *)env_ptr;
	}

	return bs_remote;
}

enum launch_status launch_ctx_execute(
	struct launch_ctx *ctx,
	const struct launch_parameters *params,
	enum launch_flags flags,
	struct launch_result *result)
{
	kern_status_t kstatus;
	kern_handle_t stack_vmo;

	kstatus = vm_object_create(
		"stack",
		5,
		STACK_SIZE,
		VM_PROT_READ | VM_PROT_WRITE | VM_PROT_USER,
		&stack_vmo);

	size_t name_len = params->p_task_name ? strlen(params->p_task_name) : 0;
	kern_handle_t remote_task = KERN_HANDLE_INVALID,
		      remote_address_space = KERN_HANDLE_INVALID;
	kstatus = task_create(
		params->p_parent_task,
		params->p_task_name,
		name_len,
		&remote_task,
		&remote_address_space);
	if (kstatus != KERN_OK) {
		kern_handle_close(stack_vmo);
		return LAUNCH_ERR_TASK_CREATION_FAILED;
	}

	struct elf_image image;
	elf_image_init(&image);

	enum launch_status status = elf_image_load(
		&image,
		params->p_executable,
		params->p_local_address_space,
		remote_address_space);

	if (status == LAUNCH_ERR_INTERPRETER_REQUIRED) {
		kern_handle_t interp = get_library(
			ctx,
			image.e_interp,
			params->p_resolver_arg);
		if (interp == KERN_HANDLE_INVALID) {
			elf_image_cleanup(&image);
			kern_handle_close(stack_vmo);
			kern_handle_close(remote_address_space);
			kern_handle_close(remote_task);
			return status;
		}

		elf_image_init(&image);
		status = elf_image_load(
			&image,
			interp,
			params->p_local_address_space,
			remote_address_space);
	}

	if (status != LAUNCH_OK) {
		elf_image_cleanup(&image);
		kern_handle_close(stack_vmo);
		kern_handle_close(remote_address_space);
		kern_handle_close(remote_task);
		return status;
	}

	virt_addr_t remote_stack_buf, local_stack_buf;
	kstatus = vm_region_map_relative(
		remote_address_space,
		VM_REGION_ANY_OFFSET,
		stack_vmo,
		0,
		STACK_SIZE,
		VM_PROT_READ | VM_PROT_WRITE | VM_PROT_USER,
		&remote_stack_buf);
	kstatus = vm_region_map_relative(
		params->p_local_address_space,
		VM_REGION_ANY_OFFSET,
		stack_vmo,
		0,
		STACK_SIZE,
		VM_PROT_READ | VM_PROT_WRITE | VM_PROT_USER,
		&local_stack_buf);
	kern_handle_close(stack_vmo);

	if (kstatus != KERN_OK) {
		elf_image_cleanup(&image);
		kern_handle_close(remote_address_space);
		kern_handle_close(remote_task);
		return LAUNCH_ERR_MEMORY_MAP_FAILED;
	}

	struct stack_writer stack;
	stack_writer_init(
		&stack,
		local_stack_buf + STACK_SIZE,
		remote_stack_buf + STACK_SIZE);
	virt_addr_t bsdata = write_bootstrap_data(&stack, params);

	virt_addr_t ip = image.e_hdr.e_entry + image.e_remote_base;

	kern_handle_t thread;
	kstatus = task_create_thread(
		remote_task,
		ip,
		stack.w_remote_sp,
		&bsdata,
		1,
		&thread);
	if (kstatus != KERN_OK) {
		elf_image_cleanup(&image);
		kern_handle_close(remote_address_space);
		kern_handle_close(remote_task);
		return LAUNCH_ERR_THREAD_CREATION_FAILED;
	}

	thread_start(thread);

	kern_handle_close(thread);
	elf_image_cleanup(&image);

	return LAUNCH_OK;
}