#define _LARGEFILE64_SOURCE

#include <stdlib.h>
#include <stdio.h>
#include <unistd.h>
#include <errno.h>
#include <assert.h>

#include <sys/mman.h>
#include <sys/types.h>
#include <sys/stat.h>
#include <fcntl.h>

#include <algorithm>
#include <iostream>

#include "va2pa.h"

/* The following is taken from /usr/src/linux/fs/proc/task_mmu.c
 *
 * /proc/pid/pagemap - an array mapping virtual pages to pfns
 *
 * For each page in the address space, this file contains one 64-bit entry
 * consisting of the following:
 *
 * Bits 0-55  page frame number (PFN) if present
 * Bits 0-4   swap type if swapped
 * Bits 5-55  swap offset if swapped
 * Bits 55-60 page shift (page size = 1<<page shift)
 * Bit  61    reserved for future use
 * Bit  62    page swapped
 * Bit  63    page present
 *
 * If the page is not present but in swap, then the PFN contains an
 * encoding of the swap file number and the page's offset into the
 * swap. Unmapped pages return a null PFN. This allows determining
 * precisely which pages are mapped (or in swap) and comparing mapped
 * pages between processes.
 *
 * Efficient users of this interface will use /proc/pid/maps to
 * determine which areas of memory are actually mapped and llseek to
 * skip over unmapped regions.
 */

// entries in /proc/pid/pagemap are 64 bits
#define PAGEMAP_ENTRY_SIZE 8

unsigned va_to_pa(pfn * map, void * area, size_t size)
{
	uintptr_t start = ((uintptr_t) area) >> PAGE_SIZE_SHIFT;
	size_t pages = size >> PAGE_SIZE_SHIFT;

	int pagemap_fd = open("/proc/self/pagemap", O_RDONLY | O_LARGEFILE);

	if (pagemap_fd == -1) {
		perror("open");
		return -1;
	}

	off64_t pagemap_pos = start * PAGEMAP_ENTRY_SIZE;

	if (lseek64(pagemap_fd, pagemap_pos, SEEK_SET) != pagemap_pos) {
		perror("lseek64");
		if (close(pagemap_fd) == -1) {
			perror("close");
		}
		return -1;
	}

	uint64_t entry;

	while (pages-- > 0) {
		if (read(pagemap_fd, &entry, PAGEMAP_ENTRY_SIZE) != PAGEMAP_ENTRY_SIZE) {
			perror("read");
			if (close(pagemap_fd) == -1) {
				perror("close");
			}
			return -1;
		}
		//printf("present: %c shift: %d pfn: %llx\n", (entry & (uint64_t) 1 << 63) ? 'y' : 'n', (int)((entry >> 55) & ((uint64_t) -1 >> (64-5))), entry & ((uint64_t) -1 >> (64-55)));
		*map++ = (pfn) (entry & ((uint64_t) -1 >> (64-55)));
	}

	if (close(pagemap_fd) == -1) {
		perror("close");
	}
	return 0;
}

void * smart_alloc(size_t size, unsigned pagesets, bool executable)
{
	assert(size % PAGE_SIZE == 0);
	// how many pages we need
	size_t nNeedPages = size / PAGE_SIZE;
	// how many pages we need from each set, rounded up
	unsigned nPagesPerSet = (nNeedPages + pagesets - 1) / pagesets;
	// how many pages we will allocate and split to sets
	size_t nInitPages = std::max(size, (size_t) (pagesets * PAGE_SIZE)) * 2 / PAGE_SIZE;
	// array for storing counts of pages allocated in each set
	unsigned * setSizes = (unsigned *) malloc(pagesets * sizeof(unsigned));

	// allocate memory until we have enough to fill all page sets
	void * initPages = NULL;
	pfn * initPagesPFN = NULL;
	bool enoughPages = false;
	while (!enoughPages) {
		// how many bytes to allocate
		size_t initPagesSize = nInitPages * PAGE_SIZE;

		// get a memory area populated by physical pages, hopefully large around
		initPages = mmap(NULL, initPagesSize,
				PROT_READ | PROT_WRITE | (executable ? PROT_EXEC : 0),
				MAP_PRIVATE | MAP_ANONYMOUS | MAP_POPULATE,
				-1, 0);
		if (initPages == MAP_FAILED) {
			perror("smart_alloc: initPages = mmap()");
			return MAP_FAILED;
		}
		// determine pfn's of pages in initPages
		initPagesPFN = (pfn *) malloc(nInitPages * sizeof(pfn));

		if (va_to_pa(initPagesPFN, initPages, initPagesSize) != 0) {
			return MAP_FAILED;
		}
		// count how many pages we have in each set
		for (unsigned i = 0; i < pagesets; i++) {
			setSizes[i] = 0;
		}
		for (size_t i = 0; i < nInitPages; i++) {
			unsigned set = initPagesPFN[i] % pagesets;
			setSizes[set]++;
		}
		// check if there's enough
		enoughPages = true;
		for (unsigned i = 0; i < pagesets; i++) {
			//printf("%d: %d of %d %s\n", i, setSizes[i], nPagesPerSet, setSizes[i] >= nPagesPerSet ? "OK!" : "NOT ENOUGH!");
			// the are was not big enough, unmap it
			if (setSizes[i] < nPagesPerSet) {
				std::cout << "smart_alloc: not enough pages in a set, trying again with twice memory" << std::endl;
				if (munmap(initPages, initPagesSize) != 0) {
					perror("smart_alloc: munmap(not_enough_initPages)");
					return MAP_FAILED;
				}
				free(initPagesPFN);
				// try more memory next time
				nInitPages *= 2;
				enoughPages = false;
				break;
			}
		}
	}

	// now allocate the final area, with extra pages for realignment
	char * finalPages = (char *) mmap(NULL, (nNeedPages + pagesets - 1) * PAGE_SIZE,
			PROT_READ | PROT_WRITE | (executable ? PROT_EXEC : 0), MAP_PRIVATE | MAP_ANONYMOUS | MAP_NORESERVE, -1, 0);
	if (finalPages == MAP_FAILED) {
			perror("smart_alloc: finalPages = mmap()");
			return MAP_FAILED;
	}
	// realign the beginning
	char * newFinalPages = (char *) ((((uintptr_t)finalPages + (pagesets - 1) * PAGE_SIZE) / (pagesets * PAGE_SIZE)) * (pagesets * PAGE_SIZE));
	//printf("aligned from %p to %p\n", finalPages, newFinalPages);
	// unmap unneeded pages from the beginning
	if (newFinalPages > finalPages) {
		if (munmap(finalPages, newFinalPages - finalPages) != 0) {
			perror("smart_alloc: munmap(beginning)");
			return MAP_FAILED;
		}
	}
	// unmap unneeded pages from the end
	char * newFinalPagesEnd = newFinalPages + nNeedPages * PAGE_SIZE;
	char * finalPagesEnd = finalPages + (nNeedPages + pagesets - 1) * PAGE_SIZE;
	if (newFinalPagesEnd < finalPagesEnd) {
		if (munmap(newFinalPagesEnd, finalPagesEnd - newFinalPagesEnd) != 0) {
			perror("smart_alloc: munmap(end)");
			return MAP_FAILED;
		}
	}

	finalPages = newFinalPages;

	// from now on, setSizes will tell how many pages of each set we already remapped
	// which is also used to determine the address to remap to next
	for (unsigned i = 0; i < pagesets; i++) {
		setSizes[i] = 0;
	}

	// walk through the init pages and remap them according to their pfn
	char * from = (char *) initPages;
	for (size_t i = 0; i < nInitPages; i++) {
		unsigned set = initPagesPFN[i] % pagesets;
		// address to remap to
		void * target = (void *) (((uintptr_t) finalPages) + (setSizes[set] * pagesets + set) * PAGE_SIZE);
		// do we still need a page from this set?
		if ((uintptr_t) target < (uintptr_t) finalPages + nNeedPages * PAGE_SIZE) {
			// yes, remap it
			if (mremap(from, PAGE_SIZE, PAGE_SIZE, MREMAP_FIXED | MREMAP_MAYMOVE, target) != target) {
				perror("smart_alloc: mremap()");
				return MAP_FAILED;
			}
			// increase the counter of this set
			setSizes[set]++;
		} else {
			// no, unmap it
			if (munmap(from, PAGE_SIZE) != 0) {
				perror("smart_alloc: munmap(unusable_page)");
				return MAP_FAILED;
			}
		}
		from += PAGE_SIZE;
	}

	//mremap(finalPages, nNeedPages * PAGE_SIZE, nNeedPages * PAGE_SIZE, 0);

//	printf("init %p final %p\n", initPages, finalPages);

	// check that everything went as expected
	pfn * finalPagesPFN = (pfn *) malloc(nNeedPages * sizeof(pfn));
	if (va_to_pa(finalPagesPFN, finalPages, nNeedPages * PAGE_SIZE) != 0) {
		return MAP_FAILED;
	}

	for (size_t i = 0; i < nNeedPages; i++) {
		//printf("%p -> %x\n", (void *) ((uintptr_t) finalPages + i * PAGE_SIZE), finalPagesPFN[i]);
		assert((((uintptr_t) finalPages + i * PAGE_SIZE) >> PAGE_SIZE_SHIFT) % pagesets == finalPagesPFN[i] % pagesets);
	}

	free(finalPagesPFN);
	free(initPagesPFN);
	free(setSizes);

	return finalPages;
}