#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>

typedef struct FreePage FreePage;
struct FreePage {
	FreePage *next;
};
static FreePage *firstfreepage;

extern uint8_t pbitmap[]; /* linker.ld */
static size_t pbitmap_len; /* in bytes */
static void *watermark; /* first nonallocated page */
static void *memtop; /* end of physical memory */
static int curallocs = 0;

static size_t
toindex(void *p)
{
	assert((void*)pbitmap <= p);
	return ((uintptr_t)p - (uintptr_t)pbitmap) / PAGE_SIZE;
}

static bool
pbitmap_get(void *p)
{
	size_t i = toindex(p);
	size_t b = i / 8;
	uint8_t m = 1 << (i&7);
	assert(b < pbitmap_len);
	return (pbitmap[b]&m) != 0;
}

static bool
pbitmap_set(void *p, bool v)
{
	size_t i = toindex(p);
	size_t b = i / 8;
	uint8_t m = 1 << (i&7);
	assert(b < pbitmap_len);
	bool prev = (pbitmap[b]&m) != 0;
	if (v) {
		pbitmap[b] |= m;
	} else {
		pbitmap[b] &= ~m;
	}
	return prev;
}

/* watermark allocator. doesn't mark the page as allocated in the bitmap. */
static void *
wmalloc(void)
{
	/* skip over reserved pages */
	while (watermark < memtop && pbitmap_get(watermark)) {
		watermark += PAGE_SIZE;
	}
	if (watermark < memtop) {
		void *p = watermark;
		watermark += PAGE_SIZE;
		return p;
	} else {
		return NULL;
	}
}

void
mem_init(void *p)
{
	memtop = p;
	kprintf("memory   %8x -> %8x\n", &_bss_end, memtop);
	pbitmap_len = (toindex(memtop) + 7) / 8; /* rounds up */
	memset(pbitmap, 0, pbitmap_len);
	mem_reserve(pbitmap, pbitmap_len);
	mem_reserve(memtop, 4096 * 8);
	watermark = pbitmap;
}

void
mem_reserve(void *addr, size_t len)
{
	kprintf("reserved %8x -> %8x\n", addr, addr + len);

	void *top = min(addr + len, memtop);
	addr = (void*)((uintptr_t)addr & ~PAGE_MASK); /* round down to page */
	for (void *p = max(addr, (void*)pbitmap); p < top; p += PAGE_SIZE) {
		/* This doesn't allow overlapping reserved regions, but, more
		 * importantly, it prevents reserving an already allocated page.
		 * Note that a reserved page can be freed. */
		if (pbitmap_get(p)) {
			panic_invalid_state();
		}
		pbitmap_set(p, true);
		curallocs++;
	}
}

void *
page_zalloc(size_t pages)
{
	void *p = page_alloc(pages);
	memset(p, 0, pages * PAGE_SIZE);
	return p;
}

void *
page_alloc(size_t pages)
{
	void *ret;
	assert(pages == 1);
	if (firstfreepage != NULL) {
		ret = firstfreepage;
		firstfreepage = firstfreepage->next;
	} else {
		ret = wmalloc();
	}

	if (ret == NULL) {
		mem_debugprint();
		kprintf("we ran out of memory :(\ngoodbye.\n");
		panic_unimplemented();
	} else {
		assert(pbitmap_get(ret) == false);
		pbitmap_set(ret, true);
		curallocs += pages;
		return ret;
	}
}

void
page_free(void *addr, size_t pages)
{
	assert((void*)pbitmap <= addr);
	for (size_t i = 0; i < pages; i++) {
		FreePage *fp = addr + i*PAGE_SIZE;
		assert(pbitmap_get(fp) == true);
		pbitmap_set(fp, false);
		fp->next = firstfreepage;
		firstfreepage = fp;
	}
	curallocs -= pages;
}

void
mem_debugprint(void)
{
	kprintf(
		"[kern] %d pages allocated, watermark at %08p / %08p\n",
		curallocs, watermark, memtop
	);
	kmalloc_debugprint();
}