#include <toxic/errno.h>
#include <toxic/string.h>
#include <toxic/math.h>
#include <toxic/kernel.h>
#include <mm/heap.h>

static int
heap_validate_table(
	struct heap_table *tbl,
	void *start,
	void *end)
{
	int	ret = 0;

	size_t	tbl_size = (size_t)(end - start);
	size_t	tbl_pages = tbl_size / PAGE_SIZE;

	if (tbl->count != tbl_pages)
		ret = -EINVAL;

	return ret;
}

static int
heap_validate_alignment(void *ptr)
{
	return ((uint32_t)ptr % PAGE_SIZE) == 0;
}

int
heap_create(
	struct heap	*heap,
	void		*start,
	void		*end,
	struct		heap_table *tbl)
{
	int		ret = -EINVAL;
	size_t		tbl_size = HEAP_ENTRY_SIZE * tbl->count;

	if (!heap_validate_alignment(start) ||
	    !heap_validate_alignment(end))
		goto out;

	memset(heap, 0, sizeof(struct heap));
	heap->start_addr = start;
	heap->table = tbl;

	ret = heap_validate_table(tbl, start, end);

	if (ret)
		goto out;

	memset(tbl->entries, 0, tbl_size);

out:
	return ret;
}

static int
heap_entry_is_used(unsigned char entry)
{
	return (entry & HEAP_ENTRY_USED_MASK) == HEAP_ENTRY_INUSE;
}

int
heap_get_start_page(
	struct heap		*heap,
	uint32_t		pages)
{
	struct heap_table	*tbl = heap->table;
	int			start_page = -1;
	int			cur_page = 0;
	size_t			i;

	for (i = 0; i < tbl->count; i++) {
		if (heap_entry_is_used(tbl->entries[i])) {
			start_page = -1;
			cur_page = 0;
			continue;
		}

		if (start_page < 0)
			start_page = i;

		/* Did we find enough blocks? */
		if (++cur_page == pages)
			break;
	}

	if (start_page < 0)
		return -ENOMEM;

	return start_page;
}

void *
heap_page_to_addr(struct heap *heap, uint32_t page)
{
	return heap->start_addr + (page * PAGE_SIZE);
}

uint32_t
heap_addr_to_page(struct heap *heap, void *ptr)
{
	return (uint32_t)(ptr - heap->start_addr) / PAGE_SIZE;
}

void
heap_mark_pages_used(struct heap *heap, uint32_t start_page, uint32_t page_count)
{
	int i;
	uint32_t end = start_page + page_count - 1;
	unsigned char cur = HEAP_ENTRY_IS_FIRST | HEAP_ENTRY_INUSE;

	if (page_count > 1)
		cur |= HEAP_ENTRY_HAS_NEXT;

	for (i = start_page; i <= end; i++) {
		heap->table->entries[i] = cur;

		cur = HEAP_ENTRY_INUSE;
		if (i != (end - 1))
			cur |= HEAP_ENTRY_HAS_NEXT;
	}
}

void
heap_mark_pages_free(struct heap *heap, uint32_t start_page)
{
	int i;
	struct heap_table *tbl = heap->table;

	for (i = start_page; i < tbl->count; i++) {
		unsigned char e = tbl->entries[i];

		tbl->entries[i] = HEAP_ENTRY_FREE;

		if (!(e & HEAP_ENTRY_HAS_NEXT))
			break;
	}


}

void *
heap_malloc_pages(struct heap *heap, uint32_t pages)
{
	void *addr = 0;

	uint32_t start_page = heap_get_start_page(heap, pages);

	if (start_page < 0)
		goto out;

	addr = heap_page_to_addr(heap, start_page);

	heap_mark_pages_used(heap, start_page, pages);

out:
	return addr;
}

void *
heap_malloc(struct heap *heap, size_t size)
{
	size_t aligned = roundup32((uint32_t) size, PAGE_SIZE);
	uint32_t pages = aligned / PAGE_SIZE;

	return heap_malloc_pages(heap, pages);
}

void
heap_free(struct heap *heap, void *ptr)
{
	heap_mark_pages_free(heap, heap_addr_to_page(heap, ptr));
}