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

#define ALIGN16(s) (((s) + 15) & ~0x0F)
#define BLOCK_SIZE sizeof(struct block)
#define MINIMUM_BLOCK_SIZE (sizeof(struct block) + 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 = (struct block *)sbrk(0); // Get the current break

    // Attempt to extend the break by s bytes
    void *new_break = sbrk(BLOCK_SIZE + s);
    if (new_break == (void *)-1)
    {
        // Handle the error, e.g., by setting an error code or printing an error message
        fprintf(stderr, "Failed to extend heap by %zu bytes.\n", s);
        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)
{
	// Calculate the size of the new block, including the block header
	size_t newBlockSize = s + BLOCK_SIZE;

	// Check if the current block can be split
	if (in->size <= newBlockSize)
	{
		// Cannot split, return the original block
		return in;
	}

	// Calculate the size of the remainder block
	size_t remainderSize = in->size - newBlockSize;

	// Create the new block in the remainder space
	struct block *newBlock = (struct block *)((char *)(in + 1) + s);
	newBlock->size = remainderSize; // Subtract the size of the block header
	newBlock->prev = in;
	newBlock->next = in->next;
	newBlock->free = 1; // Set the new block as free

	// Update the current block to reflect the reduced size
	in->size = newBlockSize; // Subtract the size of the block header
	in->next = newBlock;

	// Insert the new block into the linked list of blocks
	if (newBlock->next)
	{
		newBlock->next->prev = newBlock;
	}

	return newBlock;
}

struct block *find_best_fit(size_t s)
{
	struct block *current = first;
	struct block *best_fit = NULL;
	while (current)
	{
		if (current->free && current->size >= s)
		{
			if (best_fit == NULL || current->size < best_fit->size)
			{
				best_fit = current;
			}
		}
		current = current->next;
	}
	return best_fit; // This return statement is now outside the while loop
}

/// 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)
{
	if (size == 0)
	{
		return NULL;
	}

	size = ALIGN16(size);				   // First align the requested size
	size_t total_size = size + BLOCK_SIZE; // Then add the size of the block header

	struct block *b;

	if (first)
	{
		b = find_best_fit(size);
		if (!b)
		{
			b = extend_heap(total_size);
			if (!b)
			{
				return NULL; // Check if heap extension failed
			}
		}
		else if (b->size > total_size + MINIMUM_BLOCK_SIZE)
		{
			b = fragment_block(b, size);
		}
	}
	else
	{
		b = extend_heap(total_size);
		if (!b)
		{
			return NULL;
		}
		first = b;
	}

	b->free = 0;				   // Mark the block as used
	return (char *)b + BLOCK_SIZE; // Return a pointer to the usable memory
}

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

	if (new_size == 0)
	{
		free(ptr);
		return NULL;
	}

	struct block *b = (struct block *)((char *)ptr - BLOCK_SIZE);
	if (b->size >= new_size + BLOCK_SIZE)
	{
		return ptr; // The block is already big enough
	}

	void *new_ptr = malloc(new_size);
	if (!new_ptr)
	{
		return NULL; // Allocation failed
	}
	memcpy(new_ptr, ptr, b->size - BLOCK_SIZE); // Copy old data to new block, excluding the header size
	free(ptr);									// Free the old block

	return new_ptr;
}

/// 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 *)((char *)ptr - BLOCK_SIZE);
	if (b->free)
	{
		fprintf(stderr, "Double free detected at block %p.\n", ptr);
		abort(); // Terminate the program immediately due to serious error
	}

	b->free = 1;

	// Coalesce free blocks
	while (b->prev && b->prev->free)
	{
		// Merge with previous block
		b->prev->size += BLOCK_SIZE + b->size;
		b->prev->next = b->next;
		b = b->prev;
	}

	// If there is a next block and it's free, merge with it
	if (b->next && b->next->free)
	{
		b->size += BLOCK_SIZE + b->next->size;
		b->next = b->next->next;
		if (b->next)
		{
			b->next->prev = b; // Update the prev pointer of the next block
		}
	}

	// After merging, update the 'last' pointer if necessary
	if (!b->next)
	{
		last = b;
	}

	// Check if we can shrink the heap
	if (b == last)
	{
		// Update 'last' to the previous block or NULL if there's no previous block
		last = b->prev;
		if (last)
		{
			last->next = NULL;
		}
		else
		{
			first = NULL;
		}
		// Reduce the program break to release the memory
		sbrk(0 - (b->size + BLOCK_SIZE));
	}
}

size_t get_heap_size()
{
	size_t total_size = 0;
	struct block *current = first;

	while (current != NULL)
	{
		total_size += current->size;
		current = current->next;
	}

	return total_size;
}

int main()
{
	int *a = (int *)malloc(sizeof(int));
	int *b = (int *)malloc(sizeof(int));
	int *c = (int *)malloc(sizeof(int)); // Allocate memory for c

	*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()); // Assuming get_heap_size is implemented

	free(a); // Free memory for a
	free(b); // Free memory for b
	free(c); // Free memory for c

	return 0;
}