#include <stdint.h>
#include <limits.h>
#include <stddef.h>
#include <stdbool.h>
#include "memory.h"
#include "multiboot.h"
#include "bees.h"
#include "misc.h"
#include "string.h"

#include "liballoc.h"

extern mem_t __kernel_start;
extern mem_t __kernel_end;

struct memory_area {
	uint32_t len;
	struct memory_area *next;
	struct memory_area *prev;
};

uint32_t memory_available = 0;
uint32_t total_memory = 0;
static uint32_t total_pages = 0;
static struct memory_area *first_area = NULL;

static uint32_t *bitmap;
typedef unsigned char page[4096];
static int first_free = 0;


static int n_pages(const struct memory_area *m) {
	return (m->len - sizeof(*m)) / sizeof(page);
}

static void *page_at(int i) {
	int cur_page = 0;
	for (struct memory_area *m = first_area; m != NULL; m = m->next) {
		int pages = n_pages(m);
		if (i - cur_page < pages) {
			page *page_array = (page *) (m + 1);
			return page_array[i - cur_page];
		}
		cur_page += pages;
	}
	return NULL;
}

static int page_idx(page p) {
	uintptr_t addr = (uintptr_t) p;
	int cur_page = 0;
	for (struct memory_area *m = first_area; m != NULL; m = m->next) {
		if ((uintptr_t) m + m->len > addr) {
			return cur_page + (addr - (uintptr_t) m) / sizeof(page);
		}
		cur_page += n_pages(m);
	}
	return -1;
}

static inline bool bitmap_get(int i) {
	return (bitmap[i / 32] >> i % 32) & 1;
}

static inline void bitmap_set(int i, bool v) {
	int idx = i / 32;
	bitmap[idx] &= ~(1 << i % 32);
	bitmap[idx] |=  (v << i % 32);
}

// pay attention to these
int liballoc_lock() {
	return 0;
}
int liballoc_unlock() {
	return 0;
}

void *liballoc_alloc(int n) {
	if (n < 1) return NULL;

	int n_contiguous = 0;
	int free_region_start = 0;
	bool first = true;
	for (int i = first_free; i < total_pages; i++) {
		bool free = !bitmap_get(i);

		if (first) {
			first = false;
			first_free = i;
		}

		if (free) {
			if (n_contiguous == 0) free_region_start = i;
			n_contiguous++;
			if (n_contiguous == n) {
				for (int j = 0; j < n; j++)
					bitmap_set(free_region_start + j, true);
				return page_at(free_region_start);
			}
		} else n_contiguous = 0;
	}

	return NULL;
}

int liballoc_free(void *p, int n) {
	int idx = page_idx(p);
	if (idx == -1) return 1;

	if (idx < first_free) first_free = idx;

	for (int i = 0; i < n; i++)
		bitmap_set(idx + n, false);
	return 0;
}

void init_memory(multiboot_info_t *mb) {
	struct memory_area *prev = NULL;
	for (
		multiboot_memory_map_t *e = (multiboot_memory_map_t *) mb->mmap_addr;
		(unsigned int) e < mb->mmap_addr + mb->mmap_length; 
		e += 1
	) {
		if (e->type != MULTIBOOT_MEMORY_AVAILABLE) continue;
		// memory below 1MB is not safe to use.
		if (e->addr < (uint32_t) __kernel_start) continue;
		// memory above UINT_MAX is unusable.
		if (e->addr > UINT_MAX) continue;

		uint32_t length;
		// if the memory reaches outside of UINT_MAX, truncate length value
		if (e->addr + e->len > UINT_MAX)
			length = UINT_MAX - e->addr;
		else length = e->len;

		// free memory comes after the kernel
		uintptr_t addr = e->addr;
		if (addr < (uintptr_t) __kernel_end) {
			addr = (uintptr_t) __kernel_end;
			length -= addr - e->addr;
		}

		// join memory area to linked list
		struct memory_area *cur = (struct memory_area *) addr;

		cur->len = length;
		cur->next = NULL;
		cur->prev = prev;

		if (prev != NULL)
			prev->next = cur;

		total_pages += n_pages(cur);
		total_memory += length;

		prev = cur;

		if (first_area == NULL) first_area = cur;
	}
	int bitmap_pages = total_pages / 32 / sizeof(page) + 1;

	bitmap = (uint32_t *) page_at(total_pages - bitmap_pages);
	total_pages -= bitmap_pages;
	memset(bitmap, 0, bitmap_pages * sizeof(page));

	memory_available = total_pages * sizeof(page);
}