#include <stdbool.h>
#include <limits.h>
#include <string.h>
#include <socks/memblock.h>

#define MIN(a, b) ((a) < (b) ? (a) : (b))
#define MAX(a, b) ((a) > (b) ? (a) : (b))

#define ITER(a, b) ((uint64_t)(a) | ((uint64_t)(b) << 32))
#define ITER_END ULLONG_MAX
#define IDX_A(idx) ((idx) & 0xFFFFFFFF)
#define IDX_B(idx) (((idx) >> 32) & 0xFFFFFFFF)

/* the maximum possible value for a pointer type.
 * Note that any pointers returned by the memblock API will still
 * be bounded by the defined memory regions, and not by this constant. */
#define ADDR_MAX (~(uintptr_t)0)

static memblock_region_t init_memory_regions[MEMBLOCK_INIT_MEMORY_REGION_COUNT];
static memblock_region_t init_reserved_regions[MEMBLOCK_INIT_RESERVED_REGION_COUNT];

memblock_t memblock = {
	.memory.regions = init_memory_regions,
	.memory.count = 0,
	.memory.max = MEMBLOCK_INIT_MEMORY_REGION_COUNT,
	.memory.name = "memory",

	.reserved.regions = init_reserved_regions,
	.reserved.count = 0,
	.reserved.max = MEMBLOCK_INIT_RESERVED_REGION_COUNT,
	.reserved.name = "reserved",
};

static int memblock_insert_region(memblock_type_t *type, memblock_region_t *to_add)
{
	unsigned int i = 0;

	for (i = 0; i < type->count; i++) {
		const memblock_region_t *cur = &type->regions[i];

		if (cur->base >= to_add->limit) {
			break;
		}
	}

	memblock_region_t *src = &type->regions[i];
	memblock_region_t *dst = &type->regions[i + 1];
	unsigned int count = type->count - i;

	memmove(dst, src, count * sizeof *src);

	*src = *to_add;
	type->count++;

	return 0;
}

static int memblock_remove_region(memblock_type_t *type, unsigned int i)
{
	if (i >= type->count) {
		return -1;
	}

	memblock_region_t *src = &type->regions[i + 1];
	memblock_region_t *dst = &type->regions[i];
	unsigned int count = type->count - i;

	memmove(dst, src, count * sizeof *src);
	type->count--;
	return 0;
}

int memblock_add_range(memblock_type_t *type, uintptr_t base, size_t size)
{
	if (size == 0) {
		return 0;
	}

	uintptr_t limit = base + size - 1;

	if (type->count == 0) {
		type->regions[0].base = base;
		type->regions[0].limit = limit;
		type->count++;
		return 0;
	}

	memblock_region_t new_region = { .base = base, .limit = limit };
	bool add_new = true;

	for (unsigned int i = 0; i < type->count; i++) {
		memblock_region_t *cur_region = &type->regions[i];

		/* case 1: the region being added and the current region have no connection what-so-ever (no overlaps) */
		if (cur_region->limit + 1 < new_region.base || cur_region->base > new_region.limit) {
			continue;
		}

		/* case 2: the region being added matches a region already in the list. */
		if (cur_region->base == new_region.base && cur_region->limit == new_region.limit) {
			/* nothing needs to be done */
			add_new = false;
			break;
		}


		/* case 3: the region being added completely contains a region already in the list. */
		if (cur_region->base > new_region.base && cur_region->limit <= new_region.limit) {
			memblock_remove_region(type, i);

			/* after memblock_remove_region(), a different region will have moved into the array slot referenced by i.
			 * decrementing i means we'll stay at the current index and process this region. */
			i--;
			continue;
		}


		/* case 4: the region being added meets or partially overlaps a region already in the list. */

		/* there can be an overlap at the beginning and the end of the region being added,
		 * anything else is either a full overlap (case 3) or not within the region being added at all.
		 * to handle this, remove the region that's already in the list and extend the region being added to cover it */
		if (new_region.base > cur_region->base || new_region.base == cur_region->limit - 1) {
			/* the new region overlaps the END of the current region, change the base of the new region to match that of the current region. */
			new_region.base = cur_region->base;
		} else if (new_region.base < cur_region->base || new_region.limit + 1 == cur_region->base){
			/* the new region overlaps the BEGINNING of the current region, change the limit of the new region to match that of the current region. */
			new_region.limit = cur_region->limit;
		} else {
			continue;
		}

		/* with the new region updated to include the current region, we can remove the current region from the list */
		memblock_remove_region(type, i);
		i--;
	}

	if (add_new) {
		memblock_insert_region(type, &new_region);
	}

	return 0;
}

int memblock_add(uintptr_t base, size_t size)
{
	return memblock_add_range(&memblock.memory, base, size);
}

int memblock_reserve(uintptr_t base, size_t size)
{
	return memblock_add_range(&memblock.reserved, base, size);
}

uintptr_t memblock_alloc(size_t size)
{
	return 0;
}

void __next_memory_region(memblock_iter_t *it, memblock_type_t *type_a, memblock_type_t *type_b, uintptr_t start, uintptr_t end)
{
	unsigned int idx_a = IDX_A(it->idx);
	unsigned int idx_b = IDX_B(it->idx);

	for (; idx_a < type_a->count; idx_a++) {
		memblock_region_t *m = &type_a->regions[idx_a];
		if (!type_b) {
			it->base = m->base;
			it->limit = m->limit;

			it->idx = ITER(idx_a + 1, idx_b);
			return;
		}

		uintptr_t m_start = m->base;
		uintptr_t m_end = m->limit;

		if (m_end < start) {
			/* we haven't reached the requested memory range yet */
			continue;
		}

		if (m_start > end) {
			/* we have gone past the requested memory range and can now stop */
			break;
		}

		for (; idx_b < type_b->count + 1; idx_b++) {
			memblock_region_t *r = &type_b->regions[idx_b];

			/* r_start and r_end delimit the region of memory between the current and previous reserved regions.
			 * if we have gone past the last reserved region, these variables delimit the range between the end
			 * of the last reserved region and the end of memory. */
			uintptr_t r_start = idx_b > 0 ? r[-1].limit + 1 : 0;
			uintptr_t r_end = idx_b < type_b->count ? r->base - 1 : ADDR_MAX;

			if (r_start >= m_end) {
				/* we've gone past the end of the current memory region, and need to go to the next one */
				break;
			}

			/* we've already gone past this free memory region. move to the next one */
			if (m_start >= r_end) {
				continue;
			}

			/* we want the area that is overlapped by both
			 *	region M (m_start - m_end) : The region defined as system memory.
			 *	region R (r_start - r_end) : The region defined as free / outside of any reserved regions.
			 */
			it->base = MAX(m_start, r_start);
			it->limit = MIN(m_end, r_end);

			/* further limit the region to the intersection between the region itself and the
			 * specified iteration bounds */
			it->base = MAX(it->base, start);
			it->limit = MIN(it->limit, end);

			/* whichever region is smaller, increment the pointer for that type, so we can
			 * compare the larger region with the next region of the incremented type. */
			if (m_end <= r_end) {
				idx_a++;
			} else {
				idx_b++;
			}

			/* store the position for the next iteration */
			it->idx = ITER(idx_a, idx_b);
			return;
		}
	}

	/* ULLONG_MAX signals the end of the iteration */
	it->idx = ITER_END;
}