#include <socks/kext.h>
#include <socks/printk.h>
#include <socks/vm.h>
#include <socks/util.h>
#include <socks/libc/stdio.h>

static struct kext *self = NULL;

extern struct btree kext_tree;
extern char __kexts_start[];
extern char __kexts_end[];

static kern_status_t collect_dependencies(struct kext *kext, struct kext_info *info)
{
	if (!info->k_dependencies) {
		return KERN_OK;
	}

	unsigned int i;
	for (i = 0; info->k_dependencies[i]; i++);
	kext->k_nr_dependencies = i;

	if (!kext->k_nr_dependencies) {
		return KERN_OK;
	}

	kext->k_dependencies = kmalloc(kext->k_nr_dependencies * sizeof (struct kext *), VM_NORMAL);
	if (!kext->k_dependencies) {
		kext_release(kext);
		return KERN_NO_MEMORY;
	}

	for (i = 0; info->k_dependencies[i]; i++) {
		struct kext *dep = kext_get_by_id(info->k_dependencies[i]);
		if (!dep) {
			kfree(kext->k_dependencies);

			printk("kxld: internal kext has unresolved dependency:");
			printk("kxld:   * kext '%s'", kext->k_ident);
			printk("kxld:     depends on unknown kext '%s'", info->k_dependencies[i]);
			return KERN_NO_ENTRY;
		}

		kext->k_dependencies[i] = dep;
	}

	return KERN_OK;
}

static kern_status_t create_kext_from_info(struct kext_info *info, struct kext **out)
{
	struct kext *kext = kext_alloc();
	if (!kext) {
		return KERN_NO_MEMORY;
	}

	strncpy(kext->k_ident, info->k_ident, sizeof kext->k_ident);
	kext->k_ident[sizeof kext->k_ident - 1] = 0;
	kext->k_ident_hash = hash_string(kext->k_ident);
	kext->k_flags = KEXT_INTERNAL;
	kext->k_online = info->k_online;
	kext->k_offline = info->k_offline;

	kext->k_dependencies = NULL;
	kext->k_nr_dependencies = 0;

	if (info->k_dependencies) {
		unsigned int i;
		for (i = 0; info->k_dependencies[i]; i++);
		kext->k_nr_dependencies = i;
	}

	*out = kext;
	return KERN_OK;
}

struct kext *kernel_kext(void)
{
	return self;
}

kern_status_t register_internal_kexts(void)
{
	struct kext_info *cur = (struct kext_info *)__kexts_start;
	struct kext_info *end = (struct kext_info *)__kexts_end;
	while (cur < end) {
		struct kext *kext;
		kern_status_t status = create_kext_from_info(cur, &kext);

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

		status = kext_register(kext);
		if (status != KERN_OK) {
			kext_release(kext);
			return status;
		}

		cur = (struct kext_info *)((char *)cur + __KEXT_INFO_ALIGNMENT);
	}

	return KERN_OK;
}

kern_status_t resolve_internal_dependencies(void)
{
	struct kext_info *cur = (struct kext_info *)__kexts_start;
	struct kext_info *end = (struct kext_info *)__kexts_end;
	while (cur < end) {
		struct kext *kext = kext_get_by_id(cur->k_ident);

		if (!kext) {
			return KERN_NO_ENTRY;
		}

		kern_status_t status = collect_dependencies(kext, cur);
		kext_release(kext);

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

		cur = (struct kext_info *)((char *)cur + __KEXT_INFO_ALIGNMENT);
	}

	return KERN_OK;
}

kern_status_t init_kernel_kext(void)
{
	kext_cache_init();

	self = kext_alloc();
	if (!self) {
		return KERN_NO_MEMORY;
	}

	snprintf(self->k_ident, sizeof self->k_ident, "%s", KERNEL_KEXT_ID);
	self->k_flags = KEXT_INTERNAL | KEXT_ONLINE;

	self->k_nr_dependencies = 0;
	self->k_dependencies = NULL;

	kext_register(self);
	return KERN_OK;
}

kern_status_t scan_internal_kexts(void)
{
	kern_status_t status = register_internal_kexts();
	if (status != KERN_OK) {
		return status;
	}

	status = resolve_internal_dependencies();

	return status;
}

kern_status_t bring_internal_kexts_online(void)
{
	struct btree_node *cur = btree_first(&kext_tree);
	while (cur) {
		struct kext *kext = BTREE_CONTAINER(struct kext, k_node, cur);

		kern_status_t status = kext_bring_online(kext);
		if (status != KERN_OK) {
			return status;
		}

		cur = btree_next(cur);
	}

	return KERN_OK;
}