#include <socks/types.h>
#include <socks/memblock.h>
#include <socks/vm.h>
#include <socks/util.h>
#include <socks/printk.h>
#include <socks/libc/string.h>

#define SIZE_N_CACHE(s) \
	{ .c_name = "size-" # s, .c_obj_size = s, .c_page_order = VM_PAGE_16K }

static int kmalloc_initialised = 0;

/* reserve space for the size-N caches: */
static vm_cache_t size_n_caches[] = {
	SIZE_N_CACHE(16),
	SIZE_N_CACHE(32),
	SIZE_N_CACHE(48),
	SIZE_N_CACHE(64),
	SIZE_N_CACHE(96),
	SIZE_N_CACHE(128),
	SIZE_N_CACHE(160),
	SIZE_N_CACHE(256),
	SIZE_N_CACHE(388),
	SIZE_N_CACHE(512),
	SIZE_N_CACHE(576),
	SIZE_N_CACHE(768),
	SIZE_N_CACHE(1024),
	SIZE_N_CACHE(1664),
	SIZE_N_CACHE(2048),
	SIZE_N_CACHE(3072),
	SIZE_N_CACHE(4096),
};
static const unsigned int nr_size_n_caches = sizeof size_n_caches / sizeof size_n_caches[0];

void kmalloc_init(void)
{
	for (unsigned int i = 0; i < nr_size_n_caches; i++) {
		vm_cache_init(&size_n_caches[i]);
	}

	kmalloc_initialised = 1;
	printk("vm: kmalloc size-N caches online");
}

void *kmalloc(size_t count, vm_flags_t flags)
{
	if (!count) {
		return NULL;
	}

	if (!kmalloc_initialised) {
		/* if alloc size is a power of 2, align pointer to that same power of 2,
		   otherwise, align to 8-byte boundary. this emulates the behaviour of
		   the cache allocator. */
		phys_addr_t align = 8;
		if (power_of_2(count)) {
			align = count;
		}

		return memblock_alloc(count, align);
	}

	vm_cache_t *best_fit = NULL;
	for (unsigned int i = 0; i < nr_size_n_caches; i++) {
		if (size_n_caches[i].c_obj_size >= count) {
			best_fit = &size_n_caches[i];
			break;
		}
	}

	if (!best_fit) {
		return NULL;
	}

	return vm_cache_alloc(best_fit, flags);
}

void *kzalloc(size_t count, vm_flags_t flags)
{
	void *p = kmalloc(count, flags);
	if (p) {
		memset(p, 0x0, count);
	}

	return p;
}

void kfree(void *p)
{
	if (!kmalloc_initialised) {
		/* p was probably allocated using memblock. however, memblock requires that
		   we specify the amount of memory to free, which isn't possible here. */
		return;
	}

	phys_addr_t phys = vm_virt_to_phys(p);
	vm_page_t *pg = vm_page_get(phys);
	if (!pg || !pg->p_slab) {
		return;
	}

	vm_cache_free(pg->p_slab->s_cache, p);
}