#include <stdbool.h>
#include <stddef.h>
#include <stdint.h>

#include "kprintf.h"
#include "memory.h"

#define MAX_FRAMES   1048319
#define CEIL(x, y)   (((x) + (y) - 1) / (y))
#define GET_FRAME(i) (frame_table[i / 8] & (1 << (i % 8)))
#define SET_FRAME(i, used)                                                     \
	do {                                                                   \
		if (used)                                                      \
			frame_table[i / 8] |= (1 << (i % 8));                  \
		else                                                           \
			frame_table[i / 8] &= ~(1 << (i % 8));                 \
	} while (0)

extern uint32_t end_kernel;

static uint8_t frame_table[CEIL(MAX_FRAMES, 8)];
static uint32_t remaining_frames = MAX_FRAMES;

void *kalloc_frame(uint32_t nb_frames)
{
	if (nb_frames > remaining_frames) {
		kprintf(KERN_CRIT, "Not enough frames (max: %d)\n", MAX_FRAMES);
		return NULL;
	}
	size_t i = 0;
	while (i < MAX_FRAMES) {
		size_t free_frames = 1;
		while (!GET_FRAME(i + free_frames) && free_frames < nb_frames &&
		       i + free_frames < remaining_frames)
			free_frames++;
		if (free_frames == nb_frames)
			goto end;
		i += free_frames;
	}
end:
	if (i != MAX_FRAMES) {
		for (size_t j = 0; j < nb_frames; j++)
			SET_FRAME(j + i, 1);
		remaining_frames -= nb_frames;
		return &end_kernel + i * PAGE_SIZE;
	}
	kprintf(KERN_WARNING, "Not enough frames available\n", MAX_FRAMES);
	return NULL;
}

void kfree_frame(void *frame, uint32_t nb_frames)
{
	const uint32_t start = (frame - (void *)&end_kernel) / PAGE_SIZE;

	if (start > MAX_FRAMES || frame < (void *)&end_kernel) {
		kprintf(KERN_WARNING, "Address out of range\n");
		return;
	} else if ((uint32_t)frame % PAGE_SIZE) {
		kprintf(KERN_WARNING, "Invalid address\n");
		return;
	} else if (start + nb_frames > MAX_FRAMES) {
		kprintf(KERN_WARNING, "Invalid number of frames\n");
		return;
	}
	for (size_t i = start; i < start + nb_frames; i++)
		SET_FRAME(i, 0);
	remaining_frames += nb_frames;
}