#include <kernel/arch/generic.h>
#include <kernel/malloc.h>
#include <kernel/panic.h>
#include <kernel/util.h>
#include <shared/mem.h>
#include <stdbool.h>
#include <stdint.h>

#define SCMIN  6 /* 1<<6  == 64 */
#define SCMAX 12 /* 1<<11 == 2048 */

static const char *tagnames[] = {
	"Invalid", /* TagInvalid */
	"KernelFs", /* TagKernelFs */
	"ProcFs", /* TagProcFs */
	"UserFs", /* TagUserFs */
	"OpenPath", /* TagOpenPath */
	"MountPath", /* TagMountPath */
	"MountUser", /* TagMountUser */
	"MountRoot", /* TagMountRoot */
	"Execbuf", /* TagExecbuf */
	"VfsReq", /* TagVfsReq */
	"Handleset", /* TagHandleset */
	"Handle", /* TagHandle */
	"Process", /* TagProcess */
	"ProcessHandle", /* TagProcessHandle */
	"DevTime", /* TagDevTime */
	"RootCache", /* TagRootCache */
	"PageRefcount", /* TagPageRefcount */
};
static_assert(sizeof(tagnames) == sizeof(const char *) * TagLast);

typedef struct Slab Slab;
struct Slab {
	/* The slab is divided up into 1<<sizeclass sized regions.
	 * SCMIN <= sizeclass < SCMAX */
	int sizeclass;
	/* Conveniently, 4096/64 = 64, so I can fit the entire region bitmap
	 * in a word. */
	uint64_t bitmap;
	Slab **slot, *next;
};

/* The code assumes the header fits within the first region. */
static_assert(sizeof(Slab) <= 1<<SCMIN);

/* stored at the end of a region */
typedef struct AllocHdr AllocHdr;
struct AllocHdr {
	uint16_t len : 10;
	uint16_t tag : 6;
};
static_assert(sizeof(AllocHdr) == 2); /* just making sure */

/* Yes, I'm wasting a bit of memory here. */
static Slab *freeslabs[SCMAX] = {NULL};
static Slab *fullslabs[SCMAX] = {NULL};
static uint64_t curslabs = 0;

static uint64_t
findfreebit(uint64_t bitmap)
{
	/* I'm hoping gcc is smart enough to optimize this into a BSF */
	bitmap = ~bitmap;
	for (uint64_t i = 0; i < 64; i++) {
		if (bitmap & (1ull<<i)) return i;
	}
	return 64;
}

static bool
isfull(uint64_t bitmap, int sizeclass)
{
	return findfreebit(bitmap)<<sizeclass == PAGE_SIZE;
}

static uint64_t
crawllist(Slab *al, int sizeclass, bool full)
{
	(void) full;
	uint64_t tally = 0;
	int regsize = 1ull << sizeclass;
	for (; al; al = al->next) {
		kprintf(
			"%08p %016llx %s %db \n",
			al, al->bitmap, full ? "full" : "part", 1ull<<al->sizeclass
		);

		assert(al->sizeclass == sizeclass);
		assert(al->bitmap & 1);
		assert(al->slot && *al->slot == al);
		assert(isfull(al->bitmap, sizeclass) == full);
		tally++;

		for (int i = 1; i * regsize < PAGE_SIZE; i++) {
			char *p = ((char*)al) + i * regsize;
			AllocHdr *hdr = (void*)p + regsize - sizeof(AllocHdr);
			if (al->bitmap & (1ull<<i)) {
				assert(hdr->tag != TagInvalid && hdr->tag < TagLast);
				kprintf("%08p %4db %s\n", p, hdr->len, tagnames[hdr->tag]);
			}
		}
	}
	return tally;
}

void
kmalloc_debugprint(void)
{
	uint64_t tally = 0;
	for (int i = 0; i < SCMAX; i++) {
		tally += crawllist(freeslabs[i], i, false);
		tally += crawllist(fullslabs[i], i, true);
	}
	assert(tally == curslabs);
	// TODO count waste
}

static int
getsizeclass(size_t len)
{
	int i;
	for (i = SCMIN; i < SCMAX-1; i++) { /* last iteration can be skipped */
		if (len <= (1ull<<i)) break;
	}
	assert(len <= (1ull<<i));
	return i;
}

static Slab *
getheader(void *addr)
{
	Slab *hdr = (void*)((uintptr_t)addr & ~PAGE_MASK);
	assert(SCMIN <= hdr->sizeclass && hdr->sizeclass < SCMAX);
	assert(hdr->bitmap & 1);
	return hdr;
}

void *
kmalloc(size_t len, enum MallocTag tag)
{
	Slab *al;
	void *ret;
	int sizeclass, regsize;
	uint64_t idx;

	assert(tag != TagInvalid && tag < TagLast);
	assert(len <= KMALLOC_MAX);
	len += sizeof(AllocHdr);
	sizeclass = getsizeclass(len);
	regsize = 1ull << sizeclass;

	if (freeslabs[sizeclass] == NULL) {
		al = page_alloc(1);
		al->sizeclass = sizeclass;
		al->bitmap = 1; /* mark header as allocated */
		al->slot = &freeslabs[sizeclass];
		al->next = NULL;
		freeslabs[sizeclass] = al;
		curslabs++;
	}
	al = freeslabs[sizeclass];

	idx = findfreebit(al->bitmap);
	assert(1 <= idx && idx < 64);
	assert((al->bitmap & (1ull<<idx)) == 0);
	assert(regsize * idx < PAGE_SIZE);
	ret = ((char*)al) + regsize * idx;
	al->bitmap |= 1ull<<idx;

	AllocHdr *hdr = ret + regsize - sizeof(AllocHdr);
	hdr->len = len - sizeof(AllocHdr);
	hdr->tag = tag;

	/* if this was the last free region, move to fullslabs */
	if (isfull(al->bitmap, sizeclass)) {
		*al->slot = al->next;
		al->next->slot = al->slot;

		al->next = fullslabs[sizeclass];
		al->next->slot = &al->next;
		fullslabs[sizeclass] = al;
		al->slot = &fullslabs[sizeclass];
	}

	return ret;
}

void
kfree(void *ptr)
{
	assert(ptr);
	Slab *al = getheader(ptr);
	int sizeclass = al->sizeclass;
	int regsize = 1ull << sizeclass;

	uint64_t off = ((uint64_t)ptr) & PAGE_MASK;
	uint64_t idx = off >> sizeclass;
	bool full = isfull(al->bitmap, sizeclass);

	assert(((char*)al) + idx * regsize == ptr); /* aligned? */
	assert(1 <= idx); /* not the header? */
	assert(al->bitmap & (1ull<<idx)); /* actually allocated? */
	al->bitmap &= ~(1ull<<idx);

	memset(ptr, 0xCC, regsize);

	if (al->bitmap == 1) { /* we're empty */
		*al->slot = al->next;
		al->next->slot = al->slot;
		page_free(al, 1);
		curslabs--;
	} else if (full) { /* we're no longer full */
		*al->slot = al->next;
		al->next->slot = al->slot;

		al->next = freeslabs[sizeclass];
		al->next->slot = &al->next;
		freeslabs[sizeclass] = al;
		al->slot = &freeslabs[sizeclass];
		assert(fullslabs[sizeclass] != al);
	}
}