#include "config.h"
#include "error.h"
#include "memory.h"
#include "plat.h"

#include <stddef.h>
#include <stdint.h>
#if use_system_malloc
#include <stdlib.h>
#endif

static Heap global_heap;
static char* buffer;

static int stack_idx;

void memory_init(void* memory) {
	buffer = (char*)memory;
	stack_idx = 0;
	init_heap(
		&global_heap,
		stack_alloc(memory_heap_size),
		memory_heap_size
	);
}

static void* stack_alloc_impl(int size) {
	char* a;
#ifdef DEBUG
	if (stack_idx + size > memory_size) {
		print_err("Stack exhausted.\n");
		pbreak(error_out_of_memory);
	}
#endif
	a = &buffer[stack_idx];
	stack_idx += size;
	return a;
}

void* stack_alloc(int size) {
	return stack_alloc_aligned(
		size,
		allocation_default_alignment
	);
}

static uintptr_t align_address(
	uintptr_t address,
	size_t align
) {
	size_t m;
	m = align - 1;
	return (address + m) & ~m;
}

void* stack_alloc_aligned(int size, unsigned align) {
	char* p;
	size += (int)align + 1;
	p = stack_alloc_impl(size);
	p = (char*)align_address((uintptr_t)p, (size_t)align);
	return p;
}

void init_heap(Heap* heap, char* buf, int size) {
	heap->buffer = buf;
	heap->size = size;
	heap->blocks = 1;
	heap->usage = 0;
	((int*)buf)[0] = size << 1;
}

static void* heap_alloc_impla(Heap* h, int size) {
	int* header, * nh, s, f, i, as, p, m;
	i = 0;
	as = size + sizeof *header;
	p = 0;
	m = sizeof *header;
	h->usage += as;
	printf("%d\n", h->usage);
	for (i = 0; i < h->blocks; i++) {
		header = (int*)(h->buffer + p);
		s = header[0];
		f = !(s & 1);
		s >>= 1;
		if (f) {
			if (as < s - m) {
				s -= as;
				header[0] = (s << 1);
				nh = (int*)(h->buffer + p + s);
				nh[0] = (as << 1) | 1;
				h->blocks++;
				return nh + 1;
			} else if (s == as) {
				header[0] |= 1;
				return header + 1;
			}
		}
		p += s;
	}
	return 0;
}

static void heap_defrag(Heap* h) {
	int* header, *h2, i, j, s, f, d, n, ns, p;
	p = n = 0;
	for (i = 0; i < h->blocks; i++) {
		header = (int*)(h->buffer + p);
		s = header[0];
		f = !(s & 1);
		s >>= 1;
		if (f) {
			ns = s;
			p += s;
			d = 0;
			for (j = i + 1; j < h->blocks; j++) {
				h2 = (int*)(h->buffer + p);
				s = h2[0];
				f = !(s & 1);
				if (!f) { break; }
				s >>= 1;
				p += s;
				ns += s;
				d++;
			}
			i = j - 1;
			n += d;
			if (d)
				header[0] = ns << 1;
		} else
			p += s;
	}
	h->blocks -= n;
}

static void* heap_alloc_impl(Heap* h, int size) {
	void* ptr = heap_alloc_impla(h, size);
	if (!ptr) {
		heap_defrag(h);
		ptr = heap_alloc_impla(h, size);
		if (!ptr) {
			print_err("Heap exhausted.\n");
			pbreak(error_out_of_memory);
		}
		return ptr;
	}
	return ptr;
}

static void heap_free_impl(Heap* h, void* ptr) {
	int* header, s;
	header = &((int*)ptr)[-1];
	s = header[0] >> 1;
	h->usage -= s;
	*header &= INT32_MAX << 1;
}

void heap_print_blocks(const Heap* h) {
	const int* header;
	int s, f, p, i;
	p = 0;
	print("%8s %8s %s\n", "offset", "size", "free");
	for (i = 0; i < h->blocks; i++) {
		header = (const int*)(h->buffer + p);
		s = header[0];
		f = !(s & 1);
		s >>= 1;
		print("%8d %8d %d\n", p, s, f);
		p += s;
	}
}

void* heap_alloc_aligned(Heap* h, int size, int align) {
	unsigned char* p, * a;
	ptrdiff_t shift;
	size += (int)align;
	p = heap_alloc_impl(h, size);
	if (!p) { return 0; }
	a = (unsigned char*)align_address((uintptr_t)p, align);
	a += align * (unsigned)(p == a);
	shift = a - p;
	a[-1] = shift & 0xff;
	return a;
}

void heap_free_aligned(Heap* h, void* ptr) {
	unsigned char* a;
	ptrdiff_t shift;
	a = ptr;
	shift = a[-1];
	shift += 256 * shift == 0;
	a -= shift;
	heap_free_impl(h, a);
}

void* heap_alloc(Heap* h, int size) {
	return heap_alloc_aligned(
		h,
		size,
		allocation_default_alignment
	);
}

void heap_free(Heap* h, void* ptr) {
	heap_free_aligned(h, ptr);
}

void* galloc(int size) {
#if use_system_malloc
	return malloc(size);
#else
	return heap_alloc(&global_heap, size);
#endif
}

void gfree(void* ptr) {
#if use_system_malloc
	free(ptr);
#else
	heap_free(&global_heap, ptr);
#endif
}

int get_stack_usage() {
	return stack_idx;
}

int get_heap_usage() {
	return global_heap.usage;
}

void init_arena(Arena* arena, char* buffer, int size) {
	arena->buffer = buffer;
	arena->size = size;
	arena->ptr = 0;
}

static void* arena_alloc_impl(Arena* arena, int size) {
	char* a;
#ifdef DEBUG
	if (arena->ptr + size > arena->size) {
		print_err("Arena exhausted.\n");
		pbreak(error_out_of_memory);
	}
#endif
	a = &arena->buffer[arena->ptr];
	arena->ptr += size;
	return a;
}

void* arena_alloc(Arena* arena, int size) {
	return arena_alloc_aligned(
		arena,
		size,
		allocation_default_alignment
	);
}

void* arena_alloc_aligned(
	Arena* arena,
	int size,
	int align
) {
	char* p;
	size += (int)align + 1;
	p = arena_alloc_impl(arena, size);
	p = (char*)align_address((uintptr_t)p, (size_t)align);
	return p;
}