#include <stddef.h>
#include <stdio.h>
#include <unistd.h>

#define ALIGN4(s) (((((s)-1)>>2)<<2)+4)
#define BLOCK_SIZE sizeof(struct block)
#define MINIMUM_BLOCK_SIZE  16

/// The memory block's header.
struct block
{
	size_t size;
	struct block* prev;
	struct block* next;
	int free;
};

struct block* first = NULL;
struct block* last = NULL;

/// Extends heap memory upwards, towards zero.
/// @param [in] s The size of the memory needed aligned by 4 bytes.
/// @returns The new memory block.
struct block* extend_heap(size_t s)
{
	// Ensure the allocated size is at least the minimum block size
	if (s < MINIMUM_BLOCK_SIZE)
		s = MINIMUM_BLOCK_SIZE;

	struct block* b = sbrk(0);

	if (sbrk(BLOCK_SIZE + s) == (void*)-1)
		return NULL;

	b->size = s;
	b->prev = last;
	b->next = NULL;
	b->free = 0;

	if (last)
		last->next = b;

	last = b;

	return b;
}

/// Finds the first block that will fit the given size.
/// @param [in] s The 4 byte aligned size to look for.
/// @returns The matching available memory block.
struct block* find_first(size_t s)
{
	struct block* current = first;
	while (current && (!current->free || current->size < s))
		current = current->next;

	return current;
}

/// Fragments an existing free memory block into the given size.
/// @param [in] in The memory block to fragment.
/// @param [in] s The size of the new memory block.
/// @returns The new memory block.
struct block* fragment_block(struct block* in, size_t s)
{
	size_t totalSize = BLOCK_SIZE + s;
	struct block* b = (struct block*)((char*)(in + 1) + (in->size - totalSize));
	b->size = s;
	b->prev = in;
	b->next = in->next;
	b->free = 0;

	in->size -= totalSize;
	in->next = b;

	return b;
}

/// Will find or allocate a memory block.
/// @param [in] size The size of the memory block to request.
/// @returns The requested memory on the heap.
/// @todo Fragmenting functionality.
void* malloc(size_t size)
{
	struct block* b;

	size = ALIGN4(size);

	if (first)
	{
		b = find_first(size);
		if (!b)
			b = extend_heap(size);
		else if (b->size > BLOCK_SIZE + size)
			b = fragment_block(b, size);
	}
	else
	{
		b = extend_heap(size);
		if (!b)
			return NULL;

		first = b;
	}

	return b + 1;
}

void* realloc(void* ptr, size_t new_size)
{
	if (!ptr)
		return malloc(new_size);

	if (!new_size)
		return NULL;

	struct block* b = (struct block*)ptr - 1;
	if (b->free)
		return NULL;

	if (b->size == new_size)
		return ptr;
	else if (b->size > BLOCK_SIZE + new_size)
		return fragment_block(b, new_size) + 1;

	return NULL;
}

/// Will flag the provided memory as free and will defragment other blocks adjacent to it.
/// @param [in] ptr The memory to flag as free.
/// @note If all data after the provided memory is free, it will reduce the heap size.
void free(void* ptr)
{
	if (!ptr)
		return;

	struct block* b = (struct block*)ptr - 1;
	if (b->free)
		return;

	b->free = 1;

	while (b->prev && b->prev->free)
		b = b->prev;

	while (b->next && b->next->free)
	{
		b->size += BLOCK_SIZE + b->next->size;
		b->next = b->next->next;
	}

	if (!b->next)
		brk(b);
}

/// Retrieves the total heap usage in bytes.
/// @returns Usage in bytes.
size_t get_heap_size()
{
	if (!first || !last)
		return 0;

	return ((char*)last + last->size) - (char*)first;
}

int main()
{
	int* a = (int*)malloc(sizeof(int));
	int* b = (int*)malloc(sizeof(int));

	*a = 5;
	*b = 12;

	printf("Test 1: %i\n", *a);
	printf("Test 2: %i\n", *b);
	printf("Heap Size: %zu Bytes\n", get_heap_size());

	free(a);
	free(b);

	int* c = (int*)malloc(sizeof(int));

	return 0;
}