#define MSG_IMPLEMENTATION
#define MSG_NO_MALLOC

#include "elf.h"

#include <errno.h>
#include <fcntl.h>
#include <heap/heap.h>
#include <mango/log.h>
#include <mango/msg.h>
#include <mango/task.h>
#include <mango/types.h>
#include <mango/vm.h>
#include <rosetta/bootstrap.h>
#include <rosetta/fs.h>
#include <stdarg.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <sys/mman.h>
#include <sys/remote.h>
#include <unistd.h>

static const char *search_paths[] = {
	"/usr/lib",
};
static const size_t nr_search_paths
	= sizeof search_paths / sizeof search_paths[0];

static void report_error(const char *name, int err, const char *msg, ...)
{
	char buf[1024];
	va_list arg;
	va_start(arg, msg);
	vsnprintf(buf, sizeof buf, msg, arg);
	va_end(arg);

	kern_tracef("%s: %s: %s", name, buf, strerror(err));
}

static const char *get_image_name(const char *path)
{
	const char *last_slash = NULL;
	for (size_t i = 0; path[i]; i++) {
		if (path[i] == '/') {
			last_slash = path + i;
		}
	}

	return last_slash ? last_slash + 1 : path;
}

static int find_image(struct elf_image *img)
{
	const char *name = img->e_leaf.l_name;
	char path[4096];

	for (size_t i = 0; i < nr_search_paths; i++) {
		snprintf(path, sizeof path, "%s/%s", search_paths[i], name);
		int fd = open(path, O_RDONLY);
		if (fd < 0) {
			continue;
		}

		kern_tracef("found %s -> %s", name, path);
		img->e_fd = fd;
		img->e_status = ELF_IMAGE_OPEN;
		return SUCCESS;
	}

	return ENOENT;
}

static int load_images(struct image_list *list)
{
	int status = SUCCESS;
	struct image_list_iterator it;
	image_list_iterator_begin(&it, list);

	while (it.it_leaf) {
		struct elf_image *image
			= QUEUE_CONTAINER(struct elf_image, e_leaf, it.it_leaf);
		kern_tracef(
			"image: %s [%s]",
			it.it_leaf->l_name,
			elf_image_status_to_string(image->e_status));

		int new_dependencies = 0;
		switch (image->e_status) {
		case ELF_IMAGE_NONE:
			/* Find the image using its name */
			status = find_image(image);
			if (status != SUCCESS) {
				return status;
			}
		case ELF_IMAGE_OPEN:
			/* parse the image */
			status = elf_image_parse(image);
			if (status != SUCCESS) {
				return status;
			}
		case ELF_IMAGE_PARSED:
			/* load the image */
			status = elf_image_load(image);
			if (status != SUCCESS) {
				return status;
			}
		case ELF_IMAGE_LOADED:
			/* collect dependencies */
			new_dependencies
				= elf_image_collect_dependencies(image, list);
		default:
			break;
		}

		if (new_dependencies < 0) {
			return -new_dependencies;
		}

		if (new_dependencies > 0) {
			image_list_iterator_begin(&it, list);
		} else {
			image_list_iterator_move_next(&it);
		}
	}

	return SUCCESS;
}

static int link_images(struct image_list *list)
{
	int status = SUCCESS;
	struct image_list_iterator it;
	image_list_iterator_begin(&it, list);

	kern_trace("linking all images");
	while (it.it_leaf) {
		struct elf_image *image
			= QUEUE_CONTAINER(struct elf_image, e_leaf, it.it_leaf);
		kern_tracef(
			"image: %s [%s]",
			it.it_leaf->l_name,
			elf_image_status_to_string(image->e_status));

		status = elf_image_link(image);

		image_list_iterator_move_next(&it);
	}
	return SUCCESS;
}

int main(const struct rosetta_bootstrap *bs)
{
	kern_tracef("ld");

	for (size_t i = 0; i < bs->bs_argc; i++) {
		kern_tracef("argv[%zu]: %s", i, bs->bs_argv[i]);
	}

	sys_remote_set(SYS_REMOTE_NSD, 0, 0);
	const char *exec_path = bs->bs_argv[1];
	const char *task_name = bs->bs_argv[2];
	const char *image_name = get_image_name(exec_path);

	struct image_list images;
	image_list_init(&images);

	struct elf_image *exec = NULL;
	int err = elf_image_open(exec_path, &exec);
	if (err != SUCCESS) {
		report_error(
			task_name,
			err,
			"error while loading %s",
			exec_path);
		return -1;
	}

	snprintf(
		exec->e_leaf.l_name,
		sizeof exec->e_leaf.l_name,
		"%s",
		image_name);

	image_list_put(&images, &exec->e_leaf);
	err = load_images(&images);
	if (err != SUCCESS) {
		report_error(
			task_name,
			err,
			"error while loading %s",
			exec_path);
		return -1;
	}

	err = link_images(&images);
	if (err != SUCCESS) {
		report_error(
			task_name,
			err,
			"error while loading %s",
			exec_path);
		return -1;
	}

	struct image_list_iterator it;
	image_list_iterator_begin(&it, &images);

	while (it.it_leaf) {
		struct elf_image *image
			= QUEUE_CONTAINER(struct elf_image, e_leaf, it.it_leaf);
		kern_tracef(
			"image: %s [%s]",
			it.it_leaf->l_name,
			elf_image_status_to_string(image->e_status));

		image_list_iterator_move_next(&it);
	}

	kern_tracef("ld finished");
	int (*entry)(int, const char **)
		= (int (*)(int, const char **))exec->e_entry;
	return entry(bs->bs_argc - 2, bs->bs_argv + 2);
}