#include "CHM.h"

#include <stddef.h>
#include <stdio.h>
#include <unistd.h>
#include <stdlib.h>
#include <sys/syscall.h>

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

enum MemMode mode = MEM_MODE_BALANCED;
struct Block *first = NULL;
struct Block *last = NULL;

/// Sets the memory mode for the custom heap manager.
/// @param [in] new_mode The new mode to use.
void set_memory_mode(const enum MemMode new_mode)
{
	if (new_mode > MEM_MODE_BALANCED)
	{
		printf("The given memory mode, \"%i\", is invalid.", new_mode);
		return;
	}

	mode = new_mode;
}

/// Gets the current memory mode for the custom heap manager.
/// @returns The current mode.
enum MemMode get_memory_mode()
{
	return mode;
}

/// 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(const size_t s)
{
	struct Block* b = (struct Block *)syscall(SYS_brk, NULL);

	if ((void *)syscall(SYS_brk, (char *)(b + 1) + 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;
}

struct Block *extend_block(struct Block *in, const size_t newSize)
{
	if (!in->next)
	{
		void *addr = (struct Block *)syscall(SYS_brk, NULL);

		if ((void *)syscall(SYS_brk, (char *)addr + (newSize - in->size)) == (void *)-1)
			return NULL;

		in->size = newSize;

		return in;
	}

	size_t totalSize = in->size;
	struct Block *start = in->prev;
	struct Block *end = in->next;

	while (start && start->free && totalSize < newSize)
	{
		totalSize += BLOCK_SIZE + in->prev->size;
		start = in->prev;
	}

	while (end && end->free && totalSize < newSize)
	{
		totalSize += BLOCK_SIZE + end->size;
		end = end->next;
	}

	if (totalSize < newSize || (mode == MEM_MODE_EFFICIENCY && totalSize > newSize))
		return NULL;

	if (!start)
		start = in;

	if (!end)
		end = in;

	struct Block *b = start->next;
	while (b)
	{
		start->size += BLOCK_SIZE + b->size;
		b = b->next;
	}

	start->next = end->next;
	start->free = 0;

	if (!start->next)
		last = start;

	return start;
}

/// Will defragment adjacent memory blocks of the given memory block.
/// @param [in] in The memory block to defragment for.
void defragment_adjacent_blocks(struct Block *in)
{
	// Go in reverse order to find free memory blocks.
	while (in->prev && in->prev->free)
		in = in->prev;

	// Go in forward order to find memory blocks and merge them.
	while (in->next && in->next->free)
	{
		in->size += BLOCK_SIZE + in->next->size;
		in->next = in->next->next;
	}
}

/// 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.
void fragment_block(struct Block *in, size_t s)
{
	// Create the new block in the remainder space
	struct Block *newBlock = (struct Block *)((char *)(in + 1) + s);
	newBlock->size = in->size - BLOCK_SIZE - s;
	newBlock->prev = in;
	newBlock->next = in->next;
	newBlock->free = 1;

	// Update the current block to reflect the reduced size
	in->size = s;
	in->next = newBlock;
	in->free = 0;
}

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

		if (current->size >= s)
			break;

		current = current->next;
	}

	return current;
}

/// Finds the first memory block with more ore equal to the given size and fragments it if necessary.
/// @param [in] s The memory aligned size to look for.
/// @returns The matching available memory block.
struct Block *find_balanced(const size_t s)
{
	const size_t totalSize = BLOCK_SIZE + s;
	struct Block *current = first;
	while (current)
	{
		if (!current->free || current->size < s)
		{
			current = current->next;
			continue;
		}

		if (current->size == s)
			break;

		const size_t remainder = current->size - totalSize;
		if (!(remainder % ALIGNMENT))
		{
			fragment_block(current, s);
			break;
		}

		current = current->next;
	}

	return current;
}

/// Finds the best fitting memory block with the given size and fragments it if necessary.
/// @param [in] s The memory aligned size to look for.
/// @returns The matching available memory block.
struct Block *find_best_fit(const size_t s)
{
	const size_t totalSize = BLOCK_SIZE + s;
	unsigned char fragment = 0;
	struct Block *best = NULL;
	struct Block *current = first;

	while (current)
	{
		if (!current->free || current->size < s)
		{
			current = current->next;
			continue;
		}

		if (current->size == s)
		{
			fragment = 0;
			best = current;
			break;
		}

		if (best)
		{
			if (current->size < best->size && !((current->size - totalSize) % ALIGNMENT))
			{
				fragment = 1;
				best = current;
			}
		}
		else
		{
			if (!((current->size - totalSize) % ALIGNMENT))
			{
				fragment = 1;
				best = current;
			}
		}

		current = current->next;
	}

	if (best && fragment)
		fragment_block(best, s);

	return best;
}

void *malloc(size_t size)
{
	if (!size)
		return NULL;

	size = ALIGN16(size);

	struct Block *b;

	if (first)
	{
		switch (mode)
		{
			case MEM_MODE_SPEED:
				b = find_first(size);
				break;
			case MEM_MODE_EFFICIENCY:
				b = find_best_fit(size);
				break;
			case MEM_MODE_BALANCED:
				b = find_balanced(size);
				break;
		}

		if (!b)
		{
			b = extend_heap(size);
			if (!b)
				return NULL;
		}
	}
	else
	{
		b = extend_heap(size);
		if (!b)
			return NULL;

		first = b;
		last = b;
	}

	return b + 1;
}

void *realloc(void *ptr, size_t newSize)
{
	if (!newSize)
		return NULL;

	if (!ptr)
		return malloc(newSize);

	newSize = ALIGN16(newSize);

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

	size_t totalSize = BLOCK_SIZE + newSize;
	size_t remainder = b->size - totalSize;

	if (b->size < newSize && !extend_block(b, newSize))
		return NULL;
	else if (b->size > newSize && !(remainder % ALIGNMENT))
		fragment_block(b, newSize);
	else
		return ptr;

	return b + 1;
}

void free(void *ptr)
{
	if (!ptr)
		return;

	struct Block *b = (struct Block *)ptr - 1;
	if (b->free)
	{
		fprintf(stderr, "Double free detected for the address %p.\n", ptr);
		return;
	}

	b->free = 1;

	if (mode != MEM_MODE_SPEED)
		defragment_adjacent_blocks(b);

	if (!b->next)
	{
		if (!b->prev)
		{
			first = NULL;
			last = NULL;
		}
		else
			last = b;

		syscall(SYS_brk, b);
	}
}

size_t get_heap_size()
{
	if (!first || !last)
		return 0;

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

void check_memory()
{
	if (!first || !last)
		return;

	struct Block *current = first;
	while (current)
	{
		if (current->free)
			printf("Memory at the address %p has been marked free, but still persists.", current + 1);
		else
			printf("Memory at the address %p has not been marked free. Possible memory leak?", current + 1);

		current = current->next;
	}
}