#include <stdint.h>
#include <stdbool.h>
#include <string.h>
#include <stdlib.h>

#ifdef __TESTING__
#include <stdio.h>
#endif

#include <kernel/_kernel.h>
#include <kernel/pmm.h>

/**
 * What i want to do is create a linked list of all the memory structures
 *
 * Theres one at the very start of the memory
 *
 * one at 1MB
 *
 * and then one provided by ram.
 *
 *
 * So the idea is to create a way to access memory through this such that,
 * when you give a bit block number, it'll go through the first item in the linked list, if the block is out of that range, it 
 * traverses to the next node, tries to find it there, and then continues until it either runs out of memory, or finds a location
 */

#define PMM_PAGE_SIZE 4096
#define PMM_BLOCKS_PER_BYTE 8

struct pmm_mem_info {
	uint32_t startaddr;
	uint32_t len; // in kb
	uint32_t* bitmap;

	uint32_t max_blocks;
	uint32_t free_blocks;
	uint32_t used_blocks;
};

#define PMM_GET_MEM_BLOCKS(x) x.len / PMM_PAGE_SIZE

struct pmm_mem_info main_mem;

void pmm_set(uint32_t bit);
void pmm_unset(uint32_t bit);
bool pmm_test(uint32_t bit);
int pmm_first_free(void);
void pmm_init(void);
void* pmm_alloc_block(void);
void pmm_free_block(void* p);

void __pmm_set(uint32_t bit, struct pmm_mem_info* mem_block)
{
	(mem_block->bitmap)[bit / 32] |= (1 << (bit % 32));
	mem_block->used_blocks++;
	mem_block->free_blocks--;
}

void __pmm_unset(uint32_t bit, struct pmm_mem_info* mem_block)
{
	(mem_block->bitmap)[bit / 32] &= ~(1 << (bit % 32));
	mem_block->used_blocks--;
	mem_block->free_blocks++;
}

bool __pmm_test(uint32_t bit, struct pmm_mem_info* mem_block)
{
	return (mem_block->bitmap)[bit / 32] & (1 << (bit % 32));
}

int __pmm_first_free(struct pmm_mem_info* mem_block)
{
	for (uint32_t i = 0; i < PMM_GET_MEM_BLOCKS((*mem_block)) / 32; i++) {
		if (mem_block->bitmap[i] == 0xFFFFFFFF) // this segment is full
			continue;
		for (int j = 0; j < 32; j++) {
			if (mem_block->bitmap[i] & (1 << j))
				continue; // this page is used 
			return (i * 32) + j; // i * 32 is the chunk of 32, plus j to get to the page in the chunk
		}
	}
	return -1;
}

void __pmm_add_mem_block(uint32_t addr, int32_t len, struct pmm_mem_info* mem_block)
{
	mem_block->startaddr = addr; 
	mem_block->len = len;
	mem_block->bitmap = 0;
}

void __pmm_init(struct pmm_mem_info* mem_block)
{
	// TODO same as above
	mem_block->used_blocks = 0;
	mem_block->max_blocks = PMM_GET_MEM_BLOCKS(main_mem);
	mem_block->free_blocks = mem_block->max_blocks;

	memset(mem_block->bitmap, 0x0, PMM_GET_MEM_BLOCKS((*mem_block)) / PMM_BLOCKS_PER_BYTE); // declare all memory available

#ifdef __TESTING__
	printf("Initialized %1 blocks of memory (%1KiB available)\n", mem_block->max_blocks, mem_block->free_blocks * 4096);
#endif
	__pmm_set(0, mem_block); // first block must always be set
}

void* __pmm_alloc_block(struct pmm_mem_info* mem_block)
{
	if (mem_block->free_blocks == 0) {
		kerror("OUT OF MEMORY");
		return 0;
	}

	int block_in_map = pmm_first_free();
	if (block_in_map == -1) {
		kerror("OUT OF MEMORY");
		return 0;
	}

	__pmm_set(block_in_map, mem_block);

	return (void*) (mem_block->startaddr + (block_in_map * PMM_PAGE_SIZE));
}

void __pmm_free_block(void* p, struct pmm_mem_info* mem_block)
{
	uint64_t* addr = (uint64_t*) &p;
	uint32_t idx = ((*addr) - mem_block->startaddr) / PMM_PAGE_SIZE;

	// TODO this might still be a little flaky
	// should we be able to free any pointer? or just ones that we've given out?
	if (idx == 0)
		panic("Trying to free reserved memory!");

	if (pmm_test(idx) == 0)
		panic("Trying to free a block that was already free!");

	__pmm_unset(idx, mem_block);
}

void pmm_set(uint32_t bit)
{
	/**
	 * Here we want to calculate if the bit is over the length
	 * subtract the length and bit amount so that we compensate for the bit map
	 *
	 * i.e. (length / 4096) == amount of blocks in that specific mem region
	 * if (bit > amt of blocks), 
	 * go to next node, subtract amt of blocks from bit, and pass that
	 *
	 * below is merely a temporary solution
	 */
	__pmm_set(bit, &main_mem); 
}

void pmm_unset(uint32_t bit)
{
	 // TODO: same as above
	 __pmm_unset(bit, &main_mem);
}

bool pmm_test(uint32_t bit)
{
	// TODO: same as above
	return __pmm_test(bit, &main_mem);
}

int pmm_first_free(void) //TODO implement a free_s where it finds a series of free pages
{
	// TODO: same as above
	int ret = __pmm_first_free(&main_mem);
	if (ret == -1)
		kerror("OUT OF MEMORY");
	return ret;
}


void pmm_init(void)
{
	// TODO same as above
	__pmm_init(&main_mem);
}

void* pmm_alloc_block(void)
{
	return __pmm_alloc_block(&main_mem);
}

void pmm_free_block(void* p)
{
	__pmm_free_block(p, &main_mem);
}

void pmm_add_mem_block(uint32_t addr, uint32_t len)
{
	// TODO: make this add to a linked list
	__pmm_add_mem_block(addr, len, &main_mem);
}

#ifdef __TESTING__
void print_main_mem()
{
	printf("Available blocks: %1\nUsed blocks: %1\nMax blocks: %1\n",
			main_mem.free_blocks, main_mem.used_blocks, main_mem.max_blocks);
}
#endif