/*
 * Copyright (C) 2017 - 2020 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 <memkind/internal/memkind_hbw.h>

#include "common.h"
#include "allocator_perf_tool/GTestAdapter.hpp"
#include "numa.h"

#include <memory>

namespace TestPolicy
{
    typedef std::unique_ptr<struct bitmask, decltype(&numa_free_nodemask)>
    unique_bitmask_ptr;

    unique_bitmask_ptr make_nodemask_ptr()
    {
        return unique_bitmask_ptr(numa_allocate_nodemask(), numa_free_nodemask);
    }

    unique_bitmask_ptr make_cpumask_ptr()
    {
        return unique_bitmask_ptr(numa_allocate_cpumask(), numa_free_nodemask);

    }

    int get_num_of_pages(const size_t size, const size_t page_size)
    {
        size_t pages_number = size / page_size;
        pages_number += size % page_size ? 1 : 0;
        return pages_number;
    }

    std::vector<void *> get_address_of_pages(const void *ptr,
                                             const size_t pages_number, const size_t page_size)
    {
        std::vector<void *> address(pages_number);
        const size_t page_mask = ~(page_size-1);
        address[0] = (void *)((uintptr_t)ptr &
                              page_mask); //aligned address of first page
        for (size_t page_num = 1; page_num < pages_number; page_num++) {
            address[page_num] = (char *)address[page_num-1] + page_size;
        }
        return address;
    }

    void record_page_association(const void *ptr, const size_t size,
                                 const size_t page_size)
    {
        size_t pages_number = get_num_of_pages(size, page_size);
        std::vector<void *> address = get_address_of_pages(ptr, pages_number,
                                                           page_size);

        int max_node_id = numa_max_node();
        std::vector<int> nodes(pages_number);
        std::vector<int> pages_on_node(max_node_id+1);

        if (move_pages(0, pages_number, address.data(), NULL, nodes.data(),
                       MPOL_MF_MOVE)) {
            fprintf(stderr, "Error: move_pages() returned %s\n", strerror(errno));
            return;
        }

        for (size_t i = 0; i < pages_number; i++) {
            if (nodes[i] < 0) {
                fprintf(stderr,"Error: status of page %p is %d\n", address[i], nodes[i]);
                return;
            } else {
                pages_on_node[nodes[i]]++;
            }
        }

        for (size_t i = 0; i < (size_t)max_node_id + 1; i++) {
            if (pages_on_node[i] > 0) {
                char buffer[1024];
                snprintf(buffer, sizeof(buffer), "Node%zd", i);
                GTestAdapter::RecordProperty(buffer, pages_on_node[i]);
            }
        }
    }

    void check_numa_nodes(unique_bitmask_ptr &expected_bitmask, int policy,
                          void *ptr, size_t size)
    {

        const size_t page_size = sysconf(_SC_PAGESIZE);
        size_t pages_number = get_num_of_pages(size, page_size);
        std::vector<void *> address = get_address_of_pages(ptr, pages_number,
                                                           page_size);
        unique_bitmask_ptr returned_bitmask = make_nodemask_ptr();
        int status = -1;

        for (size_t page_num = 0; page_num < address.size(); page_num++) {
            ASSERT_EQ(0, get_mempolicy(&status, returned_bitmask->maskp,
                                       returned_bitmask->size, address[page_num], MPOL_F_ADDR));
            ASSERT_EQ(policy, status);
            switch(policy) {
                case MPOL_INTERLEAVE:
                    EXPECT_TRUE(numa_bitmask_equal(expected_bitmask.get(), returned_bitmask.get()));
                    break;
                case MPOL_DEFAULT:
                    break;
                case MPOL_BIND:
                case MPOL_PREFERRED:
                    for(int i=0; i < numa_num_possible_nodes(); i++) {
                        if(numa_bitmask_isbitset(returned_bitmask.get(), i)) {
                            EXPECT_TRUE(numa_bitmask_isbitset(expected_bitmask.get(), i));
                        }
                    }
                    break;
                default:
                    assert(!"Unknown policy\n");
            }
        }
    }

    void check_hbw_numa_nodes(int policy, void *ptr, size_t size)
    {
        unique_bitmask_ptr expected_bitmask = make_nodemask_ptr();

        memkind_hbw_all_get_mbind_nodemask(NULL, expected_bitmask->maskp,
                                           expected_bitmask->size);
        check_numa_nodes(expected_bitmask, policy, ptr, size);
    }

    void check_all_numa_nodes(int policy, void *ptr, size_t size)
    {
        if (policy != MPOL_INTERLEAVE && policy != MPOL_DEFAULT) return;

        unique_bitmask_ptr expected_bitmask = make_nodemask_ptr();

        for(int i=0; i < numa_num_configured_nodes(); i++) {
            numa_bitmask_setbit(expected_bitmask.get(), i);
        }

        check_numa_nodes(expected_bitmask, policy, ptr, size);
    }

    std::set<int> get_regular_numa_nodes(void)
    {
        struct bitmask *cpu_mask = numa_allocate_cpumask();
        std::set<int> regular_nodes;

        const int MAXNODE_ID = numa_max_node();
        for (int id = 0; id <= MAXNODE_ID; ++id) {
            numa_node_to_cpus(id, cpu_mask);

            if (numa_bitmask_weight(cpu_mask) != 0) {
                regular_nodes.insert(id);
            }
        }
        numa_free_cpumask(cpu_mask);

        return regular_nodes;
    }
}