#include "string.h"
#include "memory.h"
#include "unistd.h"
#include "types.h"
#include "stdlib.h"

#define IN_USE_BIT 0x80000000
#define IN_USE(x) ((size_t)(x) & IN_USE_BIT)
#define LENGTH(x) ((size_t)(x) & 0x7FFFFFFF)

void *malloc(size_t n) {
	// We only deal in 4 byte chunks.
	n += n % 4;
	
	// Start searching for free space at the first allocation
	size_t *p = VIRTUAL_HEAP_START;
	size_t *heap_top = sbrk(0);
	size_t *end, *new, *nend;
	size_t len;
	
	for (; p < heap_top; p += (len / 4) + 3) {
		len = LENGTH(*p); // Length of the chunk
		if (!IN_USE(*p)) {
//			printf("   UNUSED\n");
			// Found an unused chunk
			if (len >= n && len < n + 12) {
//				printf(" kmalloc: Found 1\n");
				// This chunk's size is close enough
				// Set its in use bits
				*p = (size_t)(len | IN_USE_BIT);
				*(p + (n / 4) + 1) = (size_t)(len | IN_USE_BIT);
				// And return a pointer to it
//				printf(" kmalloc: Found suitable chunk\n");
				return (p + 1);
			} else if (LENGTH(*p) >= n + 12) {
//				printf(" kmalloc: Found 2\n");
				// Found a suitable chunk, the chunk needs to be split
				end = p + (len / 4) + 1; // Pointer to the end of the chunk to be split

				nend = p + (n / 4) + 1; // Pointer to the end tag of the new chunk

				*p = (size_t)(n | IN_USE_BIT); // New chunk's beginning tag
				*nend = (size_t)(n | IN_USE_BIT); // New chunk's end tag
				*(nend + 1) = MALLOC_GUARD; // Add magic
				
				// Update remaining chunk's tags:
				*end = len - (12 + n); // End tag
				*(nend + 2) = len - (12 + n); // Beginning tag
				
//				printf(" kmalloc: Resized chunk\n");
				// Return pointer
				return (p + 1);
			}
		}
	}
	
	if (p >= heap_top) {
		// No suitable chunk found... allocate moar
		// n bytes plus 12 more for lengths and magic
		new = sbrk(n + 12);
//		printf(" kmalloc: New chunk 0x%x\n", new);
		*new = (size_t)(n | IN_USE_BIT);
		*(new + (n / 4) + 1) = (size_t)(n | IN_USE_BIT);
		*(new + (n / 4) + 2) = MALLOC_GUARD;

		return (new + 1);
	}
}

void free(void *cp) {
	size_t *p = (size_t *)cp - 1;
	
	// [LEN] .data .data [LEN] [MAG] [LEN] .data [LEN] [MAG] split!
	// [LEN] .data .data .data .data .data .data [LEN] [MAG] ....
	// p                                         endp

	size_t *heap_top = sbrk(0);
	size_t len = LENGTH(*p);
	size_t *endp = p + len / 4 + 1;

	// Sanity checks
	if ((p > (size_t *)VIRTUAL_HEAP_START && *(p - 1) != MALLOC_GUARD)
		|| (endp < heap_top && *(endp + 1) != MALLOC_GUARD))
	{
//		printf("%x %x %x %x\n!", VIRTUAL_HEAP_START, p, heap_top, endp + 1);
		panic("This process should really be killed, it's naughty and corrupts heap");
	}

	// Clear the in use bits for both tags
	*p &= ~IN_USE_BIT;
	*endp &= ~IN_USE_BIT;
	
	// Check if we can combine some chunks
	
	size_t *next = p + (len / 4) + 2;
	size_t alen;
	
	if (next < heap_top - 4) {
		// This chunk is not the highest one
		// Check if we can combine with the next chunk
//		printf(" kfree: Combined with next chunk\n");
		if (!IN_USE(*next)) {
			// Next chunk is unused; combine!
			alen = LENGTH(*next);
			// Chunk lengths plus 12 for the tags and magic that disappear between
			*p = len + alen + 12;
		}
	}
	
	if (p > (size_t *)VIRTUAL_HEAP_START + 1) {
		// This chunk is not the lowest one
		size_t prevlen = LENGTH(*(p - 2));
		size_t *prev = p - prevlen / 4 - 3; // Previous chunk's beginning tag
//		printf(" kfree: Combined with previous chunk\n");
		if (!IN_USE(*prev)) {
			alen = prevlen;
			*prev = len + alen + 12;
		}
	}

//	printf(" kfree: Freed %d bytes\n", len);
}

void dumpheap() {
	printf("Heap dump:\n");
	size_t *p = VIRTUAL_HEAP_START;
	size_t *heap_top = sbrk(0);
	size_t len;

	printf("  heap is 0x%x - 0x%x\n  walking the heap...\n", p, heap_top);
	
	for (; p < heap_top; p += (len / 4) + 3) {
		len = LENGTH(*p);
		printf("  chunk at 0x%x, length %d, %sin use\n", p, len, IN_USE(*p) ? "" : "not ");
	}
	printf("End.\n");
}

void *calloc(size_t n) {
	void *r = malloc(n);
	memset(r, 0, n);
	return r;
}

void *realloc(void *p, size_t n) {
	free(p);
	p = malloc(n);
	return p;
}