/*
 * Copyright (C) 2014 - 2019 Intel Corporation.
 * All rights reserved.
 *
 * Redistribution and use in source and binary forms, with or without
 * modification, are permitted provided that the following conditions are met:
 * 1. Redistributions of source code must retain the above copyright notice(s),
 *    this list of conditions and the following disclaimer.
 * 2. Redistributions in binary form must reproduce the above copyright notice(s),
 *    this list of conditions and the following disclaimer in the documentation
 *    and/or other materials provided with the distribution.
 *
 * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDER(S) ``AS IS'' AND ANY EXPRESS
 * OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED WARRANTIES OF
 * MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE DISCLAIMED.  IN NO
 * EVENT SHALL THE COPYRIGHT HOLDER(S) BE LIABLE FOR ANY DIRECT, INDIRECT,
 * INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT
 * LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR
 * PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF
 * LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE
 * OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF
 * ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
 */

#include <hbwmalloc.h>
#include <memkind.h>
#include <memkind/internal/memkind_private.h>
#include <memkind/internal/memkind_hbw.h>
#include <memkind/internal/memkind_log.h>

#include <stdlib.h>
#include <stdio.h>
#include <pthread.h>
#include <errno.h>
#include <numa.h>
#include <numaif.h>
#include <unistd.h>
#include <stdint.h>

static hbw_policy_t hbw_policy_g = HBW_POLICY_PREFERRED;
static pthread_once_t hbw_policy_once_g = PTHREAD_ONCE_INIT;

static void hbw_policy_bind_init(void)
{
    hbw_policy_g = HBW_POLICY_BIND;
}

static void hbw_policy_bind_all_init(void)
{
    hbw_policy_g = HBW_POLICY_BIND_ALL;
}

static void hbw_policy_preferred_init(void)
{
    hbw_policy_g = HBW_POLICY_PREFERRED;
}

static void hbw_policy_interleave_init(void)
{
    hbw_policy_g = HBW_POLICY_INTERLEAVE;
}

// This function is intended to be called once per pagesize
// Getting kind should be done using hbw_get_kind() defined below
static memkind_t hbw_choose_kind(hbw_pagesize_t pagesize)
{
    memkind_t result = NULL;

    hbw_set_policy(hbw_policy_g);

    int policy = hbw_get_policy();

    // PREFERRED policy have separate handling cause it can fallback
    // to non-HBW kinds in case of HBW absence
    if (policy != HBW_POLICY_PREFERRED ) {
        switch (pagesize) {
            case HBW_PAGESIZE_2MB:
                if(policy == HBW_POLICY_BIND_ALL) {
                    result = MEMKIND_HBW_ALL_HUGETLB;
                } else {
                    result = MEMKIND_HBW_HUGETLB;
                }
                break;
            case HBW_PAGESIZE_1GB:
            case HBW_PAGESIZE_1GB_STRICT:
                result = MEMKIND_HBW_GBTLB;
                break;
            default:
                if (policy == HBW_POLICY_BIND) {
                    result = MEMKIND_HBW;
                } else if (policy == HBW_POLICY_BIND_ALL) {
                    result = MEMKIND_HBW_ALL;
                } else {
                    result = MEMKIND_HBW_INTERLEAVE;
                }
                break;
        }
    } else if (memkind_check_available(MEMKIND_HBW) == 0) {
        switch (pagesize) {
            case HBW_PAGESIZE_2MB:
                result = MEMKIND_HBW_PREFERRED_HUGETLB;
                break;
            case HBW_PAGESIZE_1GB:
            case HBW_PAGESIZE_1GB_STRICT:
                result = MEMKIND_HBW_PREFERRED_GBTLB;
                break;
            default:
                result = MEMKIND_HBW_PREFERRED;
                break;
        }
    } else {
        switch (pagesize) {
            case HBW_PAGESIZE_2MB:
                result = MEMKIND_HUGETLB;
                break;
            case HBW_PAGESIZE_1GB:
            case HBW_PAGESIZE_1GB_STRICT:
                result = MEMKIND_GBTLB;
                break;
            default:
                result = MEMKIND_DEFAULT;
                break;
        }
    }
    return result;
}

static memkind_t pagesize_kind[HBW_PAGESIZE_MAX_VALUE];
static inline memkind_t hbw_get_kind(hbw_pagesize_t pagesize)
{
    if(pagesize_kind[pagesize] == NULL) {
        pagesize_kind[pagesize] = hbw_choose_kind(pagesize);
    }
    return pagesize_kind[pagesize];
}


MEMKIND_EXPORT hbw_policy_t hbw_get_policy(void)
{
    return hbw_policy_g;
}

MEMKIND_EXPORT int hbw_set_policy(hbw_policy_t mode)
{
    switch(mode) {
        case HBW_POLICY_PREFERRED:
            pthread_once(&hbw_policy_once_g, hbw_policy_preferred_init);
            break;
        case HBW_POLICY_BIND:
            pthread_once(&hbw_policy_once_g, hbw_policy_bind_init);
            break;
        case HBW_POLICY_BIND_ALL:
            pthread_once(&hbw_policy_once_g, hbw_policy_bind_all_init);
            break;
        case HBW_POLICY_INTERLEAVE:
            pthread_once(&hbw_policy_once_g, hbw_policy_interleave_init);
            break;
        default:
            return EINVAL;
    }

    if (mode != hbw_policy_g) {
        return EPERM;
    }

    return 0;
}

MEMKIND_EXPORT int hbw_check_available(void)
{
    return  (memkind_check_available(MEMKIND_HBW) == 0) ? 0 : ENODEV;
}

static inline void hbw_touch_page(void *addr)
{
    volatile char *temp_ptr = (volatile char *) addr;
    char value = temp_ptr[0];
    temp_ptr[0] = value;
}

MEMKIND_EXPORT int hbw_verify_memory_region(void *addr, size_t size, int flags)
{
    /*
     * if size is invalid, flags have unsupported bit set or if addr is NULL.
     */
    if (addr == NULL || size == 0 || flags & ~HBW_TOUCH_PAGES) {
        return EINVAL;
    }

    /*
     * 4KB is the smallest pagesize. When pagesize is bigger, pages are verified more than once
     */
    const size_t page_size = sysconf(_SC_PAGESIZE);
    const size_t page_mask = ~(page_size-1);

    /*
     * block size should be power of two to enable compiler optimizations
     */
    const unsigned block_size = 64;

    char *end = addr + size;
    char *aligned_beg = (char *)((uintptr_t)addr & page_mask);
    nodemask_t nodemask;
    struct bitmask expected_nodemask = {NUMA_NUM_NODES, nodemask.n};

    memkind_hbw_all_get_mbind_nodemask(NULL, expected_nodemask.maskp,
                                       expected_nodemask.size);

    while(aligned_beg < end) {
        int nodes[block_size];
        void *pages[block_size];
        int i = 0, page_count = 0;
        char *iter_end = aligned_beg + block_size*page_size;

        if (iter_end > end) {
            iter_end = end;
        }

        while (aligned_beg < iter_end) {
            if (flags & HBW_TOUCH_PAGES) {
                hbw_touch_page(aligned_beg);
            }
            pages[page_count++] = aligned_beg;
            aligned_beg += page_size;
        }

        if (move_pages(0, page_count, pages, NULL, nodes, MPOL_MF_MOVE)) {
            return EFAULT;
        }

        for (i = 0; i < page_count; i++) {
            /*
             * negative value of nodes[i] indicates that move_pages could not establish
             * page location, e.g. addr is not pointing to valid virtual mapping
             */
            if(nodes[i] < 0) {
                return -1;
            }
            /*
             * if nodes[i] is not present in expected_nodemask then
             * physical memory backing page is not hbw
             */
            if (!numa_bitmask_isbitset(&expected_nodemask, nodes[i])) {
                return -1;
            }
        }
    }

    return 0;
}

MEMKIND_EXPORT void *hbw_malloc(size_t size)
{
    return memkind_malloc(hbw_get_kind(HBW_PAGESIZE_4KB), size);
}

MEMKIND_EXPORT void *hbw_calloc(size_t num, size_t size)
{
    return memkind_calloc(hbw_get_kind(HBW_PAGESIZE_4KB), num, size);
}

MEMKIND_EXPORT int hbw_posix_memalign(void **memptr, size_t alignment,
                                      size_t size)
{
    return memkind_posix_memalign(hbw_get_kind(HBW_PAGESIZE_4KB), memptr, alignment,
                                  size);
}

MEMKIND_EXPORT int hbw_posix_memalign_psize(void **memptr, size_t alignment,
                                            size_t size,
                                            hbw_pagesize_t pagesize)
{
    if (pagesize == HBW_PAGESIZE_1GB_STRICT && size % (1 << 30)) {
        return EINVAL;
    }

    if((pagesize == HBW_PAGESIZE_2MB ||
        pagesize == HBW_PAGESIZE_1GB_STRICT ||
        pagesize == HBW_PAGESIZE_1GB) &&
       hbw_get_policy() == HBW_POLICY_INTERLEAVE) {

        log_err("HBW_POLICY_INTERLEAVE is unsupported with used page size!");
        return EINVAL;
    }

    return memkind_posix_memalign(hbw_get_kind(pagesize), memptr, alignment, size);
}

MEMKIND_EXPORT void *hbw_realloc(void *ptr, size_t size)
{
    return memkind_realloc(hbw_get_kind(HBW_PAGESIZE_4KB), ptr, size);
}

MEMKIND_EXPORT void hbw_free(void *ptr)
{
    memkind_free(0, ptr);
}

MEMKIND_EXPORT size_t hbw_malloc_usable_size(void *ptr)
{
    return memkind_malloc_usable_size(0, ptr);
}