#define _GNU_SOURCE
#include "internal.h"
#include <sys/mman.h>
#include <numaif.h>
#include <errno.h>
#include <string.h>

#ifndef MAP_HUGE_2MB
#define MAP_HUGE_2MB (21 << MAP_HUGE_SHIFT)
#endif

#ifndef MAP_HUGE_SHIFT
#define MAP_HUGE_SHIFT 26
#endif

int region_alloc(size_t size, int numa_node, bool huge_pages, void **out_ptr)
{
    int flags = MAP_PRIVATE | MAP_ANONYMOUS;
    size_t alloc_size = size;

    if (huge_pages) {
        flags |= MAP_HUGETLB | MAP_HUGE_2MB;
        size_t huge_page_size = 2 * 1024 * 1024;
        alloc_size = (size + huge_page_size - 1) & ~(huge_page_size - 1);
    }

    void *ptr = mmap(NULL, alloc_size, PROT_READ | PROT_WRITE, flags, -1, 0);
    if (ptr == MAP_FAILED) {
        if (huge_pages && errno == ENOMEM) {
            flags = MAP_PRIVATE | MAP_ANONYMOUS;
            alloc_size = size;
            ptr = mmap(NULL, alloc_size, PROT_READ | PROT_WRITE, flags, -1, 0);
            if (ptr == MAP_FAILED)
                return ICEPICK_E_ALLOC;
        } else {
            return huge_pages ? ICEPICK_E_HUGEPAGE : ICEPICK_E_ALLOC;
        }
    }

    if (numa_node >= 0) {
        unsigned long nodemask = 1UL << numa_node;
        if (mbind(ptr, alloc_size, MPOL_BIND, &nodemask, sizeof(nodemask) * 8,
                  MPOL_MF_MOVE | MPOL_MF_STRICT) < 0) {
            munmap(ptr, alloc_size);
            return ICEPICK_E_NUMA;
        }
    }

    if (mlock(ptr, alloc_size) < 0) {
        munmap(ptr, alloc_size);
        return ICEPICK_E_PERMISSION;
    }

    madvise(ptr, alloc_size, MADV_DONTFORK);

    memset(ptr, 0, alloc_size);

    *out_ptr = ptr;
    return 0;
}

void region_free(void *ptr, size_t size)
{
    if (ptr) {
        munlock(ptr, size);
        munmap(ptr, size);
    }
}