#include "socks/queue.h"
#include <stdint.h>
#include <stdio.h>
#include <stddef.h>
#include <stdlib.h>
#include <inttypes.h>
#include <time.h>
#include <assert.h>
#include <sys/mman.h>
#include <socks/types.h>
#include <socks/util.h>
#include <socks/memblock.h>
#include <socks/vm.h>

/* we're working with 512MiB of simulated system RAM */
#define MEMORY_SIZE_MB 512

#define ALLOC_START_MB 16
#define ALLOC_END_MB 32

#define MEMPTR(offset) ((uintptr_t)system_memory + (offset))
#define MB_TO_BYTES(v) ((size_t)(v) * 0x100000)

#define PHYS_TO_VIRT(p) ((void *)((uintptr_t)system_memory + (p)))
#define VIRT_TO_PHYS(p) ((void *)((p) - (uintptr_t)system_memory))

struct mem_map_region {
	phys_addr_t base;
	phys_addr_t limit;
	enum { REGION_FREE, REGION_RESERVED } status;
};

static struct mem_map_region mem_map[] = {
	{ .base = 0x00000000, .limit = 0x0000ffff, .status = REGION_RESERVED },
	{ .base = 0x00010000, .limit = 0x0004ffff, .status = REGION_FREE },
	{ .base = 0x00050000, .limit = 0x0005ffff, .status = REGION_RESERVED },
	{ .base = 0x00060000, .limit = 0x000fffff, .status = REGION_FREE },
	{ .base = 0x00100000, .limit = 0x001fffff, .status = REGION_RESERVED },
	{ .base = 0x00200000, .limit = 0x005fffff, .status = REGION_FREE },
	{ .base = 0x00600000, .limit = 0x007fffff, .status = REGION_RESERVED },
	{ .base = 0x00800000, .limit = MB_TO_BYTES(MEMORY_SIZE_MB) - 1, .status = REGION_FREE },
};

extern void tmp_set_vaddr_base(void *);

/* virtual address of where system memory is mapped */
static void *system_memory = NULL;

static void print_free_pages(vm_zone_t *z)
{
	printf(" * %s:\n", z->z_info.zd_name);

	for (int i = VM_PAGE_MIN_ORDER; i <= VM_PAGE_MAX_ORDER; i++) {
		if (queue_length(&z->z_free_pages[i]) == 0) {
			continue;
		}

		char size_str[64];
		data_size_to_string(vm_page_order_to_bytes(i), size_str, sizeof size_str);

		printf("  - %u pages with size %s (order-%u)\n", queue_length(&z->z_free_pages[i]), size_str, i);
	}
}

static void print_all_pages(void)
{
	for (phys_addr_t i = 0; i < UINTPTR_MAX; ) {
		vm_page_t *pg = vm_page_get(i);
		if (!pg) {
			break;
		}

		vm_zone_t *z = vm_page_get_zone(pg);
		printf(" * %08" PRIxPTR ": %s order-%u (%zu bytes) %s\n",
			i,
			z ? z->z_info.zd_name : "[none]",
			pg->p_order,
			vm_page_order_to_bytes(pg->p_order),
			pg->p_flags & VM_PAGE_RESERVED ? "reserved" : "free");
		i += vm_page_order_to_bytes(pg->p_order);
	}
}

int memory_test(void)
{
	srand(time(NULL));
	system_memory = mmap(
			NULL,
			MB_TO_BYTES(MEMORY_SIZE_MB),
			PROT_READ | PROT_WRITE,
			MAP_PRIVATE | MAP_ANONYMOUS,
			-1, 0);

	if (system_memory == MAP_FAILED) {
		perror("mmap");
		fprintf(stderr, "cannot allocate simulated system RAM buffer\n");
		return -1;
	}

	phys_addr_t pmem_base = UINTPTR_MAX, pmem_limit = 0;
	size_t nr_mem_map_entries = sizeof mem_map / sizeof mem_map[0];

	for (size_t i = 0; i < nr_mem_map_entries; i++) {
		if (mem_map[i].base < pmem_base) {
			pmem_base = mem_map[i].base;
		}

		if (mem_map[i].limit > pmem_limit) {
			pmem_limit = mem_map[i].limit;
		}
	}

	tmp_set_vaddr_base(system_memory);
	memblock_add(pmem_base, pmem_limit + 1);

	for (size_t i = 0; i < nr_mem_map_entries; i++) {
		if (mem_map[i].status == REGION_RESERVED) {
			memblock_reserve(mem_map[i].base, mem_map[i].limit - mem_map[i].base + 1);
		}
	}

	printf("allocated %u MiB (0x%zx bytes) of memory to act as system RAM at %p\n", MEMORY_SIZE_MB, MB_TO_BYTES(MEMORY_SIZE_MB), system_memory);

	printf("sizeof(vm_page_t) = %zu bytes\n", sizeof(vm_page_t));

	uintptr_t voffset = (uintptr_t)system_memory;

	memblock_init(MB_TO_BYTES(ALLOC_START_MB) + voffset, MB_TO_BYTES(ALLOC_END_MB) + voffset, voffset);

	printf("memblock heap initialised in 0x%zx-0x%zx\n", MB_TO_BYTES(ALLOC_START_MB), MB_TO_BYTES(ALLOC_END_MB));

	for (int i = 0; i < 4; i++) {
		int size = 512 + (rand() % 16384);

		phys_addr_t alloc = memblock_alloc_phys(size);
		printf("allocated %d bytes at 0x%" PRIxPTR "\n", size, alloc);
	}

	vm_zone_descriptor_t zones[] = {
		{ .zd_id = VM_ZONE_DMA, .zd_name = "dma", .zd_base = 0x0, .zd_limit = MB_TO_BYTES(16) - 1 },
		{ .zd_id = VM_ZONE_NORMAL, .zd_name = "normal", .zd_base = MB_TO_BYTES(16), .zd_limit = MB_TO_BYTES(1024) - 1 },
		{ .zd_id = VM_ZONE_HIGHMEM, .zd_name = "highmem", .zd_base = MB_TO_BYTES(1024), .zd_limit = UINTPTR_MAX },
	};

	vm_bootstrap(zones, sizeof zones / sizeof zones[0]);

	printf("memory regions:\n");

	memblock_iter_t it;
	for_each_mem_range(&it, 0, 0x100000) {
		printf("\t%08" PRIxPTR "-%08" PRIxPTR "\n",
				it.it_base,
				it.it_limit);
	}

	printf("reserved regions:\n");
	for_each_reserved_mem_range(&it, 0, 0x100000) {
		printf("\t%08" PRIxPTR "-%08" PRIxPTR " (%s)\n",
				it.it_base,
				it.it_limit,
				it.it_status == MEMBLOCK_ALLOC ? "allocated" : "reserved");
	}

	printf("free regions:\n");
	for_each_free_mem_range(&it, 0, ULLONG_MAX) {
		printf("\t%08" PRIxPTR "-%08" PRIxPTR "\n",
				it.it_base,
				it.it_limit);
	}

	vm_pg_data_t *pg_data = vm_pg_data_get(0);
	printf("free pages:\n");
	for (int i = VM_ZONE_MIN; i <= VM_ZONE_MAX; i++) {
		print_free_pages(&pg_data->pg_zones[i]);
	}

	printf("all pages:\n");
	print_all_pages();

	vm_page_t *pg = vm_page_alloc(VM_PAGE_128K, 0);
	printf("allocated 128K at 0x%lx\n", vm_page_get_paddr(pg));

	vm_page_t *a, *b;
	if (vm_page_split(pg, &a, &b) == 0) {
		printf("split page into two 64K pages at 0x%lx and 0x%lx:\n", vm_page_get_paddr(a), vm_page_get_paddr(b));

		assert(a->p_flags & VM_PAGE_HEAD);
		assert(b->p_flags & VM_PAGE_HEAD);

		printf("first page block:\n");
		vm_page_foreach (a, i) {
			printf("  0x%lx: order:%u, flags:0x%x\n", vm_page_get_paddr(i), i->p_order, i->p_flags);
			assert(i->p_flags & VM_PAGE_HUGE);
			assert((i->p_flags & VM_PAGE_RESERVED) == 0);
		}
		printf("second page block:\n");
		vm_page_foreach (b, i) {
			printf("  0x%lx: order:%u, flags:0x%x\n", vm_page_get_paddr(i), i->p_order, i->p_flags);
			assert(i->p_flags & VM_PAGE_HUGE);
			assert((i->p_flags & VM_PAGE_RESERVED) == 0);
		}

		pg = vm_page_merge(a, b);
		if (pg) {
			char size_str[64];
			data_size_to_string(vm_page_order_to_bytes(pg->p_order), size_str, sizeof size_str);
			printf("merged pages 0x%lx and 0x%lx to single page of size %s:\n", vm_page_get_paddr(a), vm_page_get_paddr(b), size_str);

			size_t block_sz = 0;
			vm_page_foreach (pg, i) {
				printf("  0x%lx: order:%u, flags:0x%x\n", vm_page_get_paddr(i), i->p_order, i->p_flags);
				assert(i->p_flags & VM_PAGE_HUGE);
				assert((i->p_flags & VM_PAGE_RESERVED) == 0);
				block_sz += VM_PAGE_SIZE;
			}

			assert(block_sz == vm_page_order_to_bytes(pg->p_order));

			vm_page_free(pg);
		} else {
			printf("cannot merge pages 0x%lx and 0x%lx\n", vm_page_get_paddr(a), vm_page_get_paddr(b));
		}
	}

	pg = vm_page_alloc(VM_PAGE_128K, 0);
	printf("allocated 128K at 0x%lx\n", vm_page_get_paddr(pg));

	if (vm_page_split(pg, &a, &b) == 0) {
		assert(a->p_order == VM_PAGE_64K);
		assert(b->p_order == VM_PAGE_64K);

		printf("split 128K block into two 64K blocks\n");
		vm_page_free(a);
		vm_page_free(b);

		/* if these conditions are true, the two blocks were successfully
		   merged after being freed. */
		if (a->p_order == VM_PAGE_128K && b->p_order == VM_PAGE_128K) {
			printf("two 64K blocks were merged into one 128K block after free\n");
		} else {
			printf("two 64K blocks were NOT merged into one 128K block after free!\n");
		}
	}

	munmap(system_memory, MB_TO_BYTES(MEMORY_SIZE_MB));
	return 0;
}