#include <memory.h>
#include <stdint.h>
#include <lib.h>
#include <stddef.h>

#define MEMORY_INTERNAL
#include <memory/physalloc.h>
#include <memory/virtalloc.h>

extern uintptr_t kernel_start;
extern uintptr_t kernel_end;

// between memory_start and kernel_start will be the bitmap
static uintptr_t memory_start = 0;

typedef unsigned char page[4096];

struct memory_area {
	uint64_t len;
	uintptr_t addr;
};

static uint64_t *bitmap;
static uint64_t total_memory;
static uint64_t free_memory;
static uint64_t page_count;
static uint64_t page_free_start;
static struct memory_area *page_start;

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

static void *page_at(int i) {
	int cur_page = 0;
	for (struct memory_area *m = page_start; m != NULL; m++) {
		int pages = n_pages(m);
		if (i - cur_page < pages) {
			return (void *) (m->addr + (PAGE_SIZE * (i - cur_page)));
		}
		cur_page += pages;
	}
	return NULL;
}

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

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

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

void *alloc_phys_page(void) {
	return alloc_phys_pages(1);
}

void *alloc_phys_pages(int pages) {
	if (pages < 1) return NULL;

	int n_contiguous = 0;
	int free_region_start = 0;
	bool first = true;
	for (uint64_t i = page_free_start; i < page_count; i++) {
		bool free = !bitmap_get(i);

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

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

	return NULL;
}

void free_phys_page(void *ptr) {
	free_phys_pages(ptr, 1);
}

void free_phys_pages(void *ptr, int pages) {
	long idx = page_idx(ptr);
	if (idx == -1) return;

	if ((unsigned) idx < page_free_start) page_free_start = idx;

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

static bool segment_invalid(const struct memory_segment *segment) {
	if (segment->type != 1) return false;
	if (segment->addr < kernel_start) return false;
	if (segment->addr + segment->len < memory_start) return false;
	if (segment->addr + segment->len < kernel_start) return false;
	return true;
}

static struct memory_area segment_to_area(const struct memory_segment *segment) {
	uint64_t length = segment->len;
	uintptr_t addr = segment->addr;

	uintptr_t start;
	if (memory_start)
		start = memory_start;
	else
		start = kernel_end;

	if (segment->addr < start) {
		addr = start;
		length -= addr - segment->addr;
	} else {
		addr = segment->addr;
	}

	struct memory_area temp;
	temp.len = length;
	temp.addr = addr;

	return temp;
}

static uintptr_t page_align(uintptr_t ptr) {
	return ptr + PAGE_SIZE - 1 / PAGE_SIZE;
}

void memory_init(struct memory_map *map) {

	memory_lock();

	virtaddr_init();

	bitmap = NULL;
	total_memory = 0;
	free_memory = 0;
	page_count = 0;
	page_free_start = 0;
	page_start = NULL;

	uintptr_t end = (uintptr_t) map;
	end += map->size;

	struct memory_segment *segment = &map->entries[0];
	int segment_count = 0;

	for(; (uintptr_t) segment < end; segment++) {
		if (segment_invalid(segment))
			continue;

		struct memory_area temp = segment_to_area(segment);
		page_count += n_pages(&temp);
		segment_count++;
	}

	long bitmap_pages = page_count / 64 / PAGE_SIZE + 1;
	free_memory = page_count * PAGE_SIZE;

	//HACK: terrible hack bad bad bad bad
	long bitmap_size = bitmap_pages * PAGE_SIZE;
	bitmap = (uint64_t *) page_at(page_count - bitmap_pages);

	long page_area_size = segment_count * sizeof(struct memory_area);
	char *page_area_addr = (char *)bitmap + bitmap_size;
	
	bitmap = mmap(bitmap, bitmap_size);
	memset(bitmap, 0, bitmap_size);
	
	memory_start = page_align(kernel_end + bitmap_size + page_area_size);

	page_area_addr = mmap(page_area_addr, page_area_size);
	page_start = (struct memory_area *) page_area_addr;
	
	struct memory_area *area = page_start;
	segment = &map->entries[0];
	
	for(; (uintptr_t) segment < end; segment++) {
		if (segment_invalid(segment))
			continue;

		struct memory_area temp = segment_to_area(segment);
		*area = temp;
		area++;
	}
	
	page_count -= bitmap_pages;

	memory_unlock();

}

void *alloc_page(void) {
	return alloc_pages(1);
}

uint64_t memory_total(void) {
	return total_memory;
}

uint64_t memory_free(void) {
	return free_memory;
}

uint64_t memory_used(void) {
	return total_memory - free_memory;
}

// stubs
// simon do these ik you want to
// :3

void *malloc(size_t size) {
	//TODO: implement
	(void)size;
	return NULL;
}

void *realloc(void *ptr, size_t size) {
	//TODO: implement
	(void)ptr;
	(void)size;
	return NULL;
}

void free(void *ptr) {
	//TODO: implement
	(void)ptr;
}