#include <gtest/gtest.h>
#include <dlfcn.h>

extern "C" {
#include "utils/defs.h"
#include "wrt/lib/dl.h"
#include "wasm_exec_env.h"

extern bool signature_check(const char *signature, int *arg_count, bool *has_ret);
}

DLContext   dl_context;
WASMExecEnv exec_env;

void init() {
    dl_context.mem.malloc = malloc;
    dl_context.mem.free   = free;
    dl_context.symTable   = nullptr;
    memset(dl_context.hnd, 0, sizeof(dl_context.hnd));
    memset(dl_context.sym, 0, sizeof(dl_context.sym));
    exec_env.attachment = &dl_context;
}

#define check_symbol_consistency(sym, postfix)                       \
    do {                                                             \
        ASSERT_GT((sym), 0);                                         \
        ASSERT_LT((sym), DL_MAX_SYMBOLS);                            \
        ASSERT_NE(dl_context.sym[(sym)-1].symbol, nullptr);          \
        ASSERT_NE(dl_context.sym[(sym)-1].cif, nullptr);             \
        ASSERT_GT(dl_context.sym[(sym)-1].backref, 0);               \
        unsigned ref##postfix = dl_context.sym[(sym)-1].backref;     \
        ASSERT_GT(ref##postfix, 0);                                  \
        ASSERT_LE(ref##postfix, DL_MAX_HANDLES);                     \
        ASSERT_NE(dl_context.hnd[ref##postfix - 1].handle, nullptr); \
    } while (0)

#define check_hnd_empty(postfix)                     \
    do {                                             \
        for (auto &hnd##postfix : dl_context.hnd) {  \
            ASSERT_EQ(hnd##postfix.handle, nullptr); \
        }                                            \
    } while (0)

#define check_sym_empty(postfix)                     \
    do {                                             \
        for (auto &sym##postfix : dl_context.sym) {  \
            ASSERT_EQ(sym##postfix.symbol, nullptr); \
            ASSERT_EQ(sym##postfix.cif, nullptr);    \
            ASSERT_EQ(sym##postfix.backref, 0);      \
        }                                            \
    } while (0)

#define check_empty(postfix)      \
    do {                          \
        check_sym_empty(postfix); \
        check_hnd_empty(postfix); \
    } while (0)

TEST(DLTest, InitAndFree) {
    // use valgrind !
    auto *context = (DLContext *)malloc(sizeof(DLContext));
    ASSERT_NE(context, nullptr);
    context->mem.malloc = malloc;
    context->mem.free   = free;
    dl_init(context);

    EXPECT_NE(context->symTable, nullptr);
    for (auto &hnd : context->hnd) {
        EXPECT_EQ(hnd.handle, nullptr);
    }
    for (auto &sym : context->sym) {
        EXPECT_EQ(sym.symbol, nullptr);
        EXPECT_EQ(sym.cif, nullptr);
        EXPECT_EQ(sym.backref, 0);
    }

    dl_free(context);
    free(context);
}

TEST(DLTest, OpenAndClose) {
    init();

    int handle = dl_open(&exec_env, "./SimpleLib.so", RTLD_LAZY);
    ASSERT_GT(handle, 0);
    int result = dl_close(&exec_env, handle);
    ASSERT_EQ(result, WRT_OK);
    check_empty(0);

    handle = dl_open(&exec_env, "not_exists", RTLD_LAZY);
    ASSERT_EQ(handle, WRT_ERROR);
    check_empty(1);

    handle = dl_open(&exec_env, "", RTLD_LAZY);
    ASSERT_EQ(handle, WRT_ERROR);
    check_empty(2);

    handle = dl_open(&exec_env, nullptr, RTLD_LAZY);
    ASSERT_EQ(handle, WRT_ERROR);
    check_empty(3);
}

TEST(DLTest, ManyOpen) {
    init();

    int handles[DL_MAX_HANDLES];
    for (int &handle : handles) {
        handle = dl_open(&exec_env, "./SimpleLib.so", RTLD_LAZY);
        ASSERT_GT(handle, 0);
    }

    {
        int handle = dl_open(&exec_env, "./SimpleLib.so", RTLD_LAZY);
        ASSERT_EQ(handle, WRT_ERROR);
    }

    for (int handle : handles) {
        int result = dl_close(&exec_env, handle);
        ASSERT_EQ(result, WRT_OK);
    }

    check_empty(0);
}

TEST(DLTest, SignatureCheck) {
    typedef struct {
        const char *test;
        bool        result;
        int         arg_count;
        bool        has_ret;
    } Case;

    const static Case cases[] = {
        {"", false, 0, false},           {"()", true, 0, false},    {"(i)", true, 1, false},
        {"()i", true, 0, true},          {"(x)", false, 0, false},  {"()x", false, 0, false},
        {"(i)f", true, 1, true},         {"(i)x", false, 0, false}, {"(x)f", false, 0, false},
        {"i", false, 0, false},          {"i()", false, 0, false},  {"i(i)", false, 0, false},
        {"(i)()", false, 0, false},      {"(ii", false, 0, false},  {"(iiffdd)l", true, 6, true},
        {"(ffddii)ll", false, 0, false}, {"()ii", false, 0, false}, {"(i", false, 0, false},
        {"(i(i))i", false, 0, false},    {"i)i", false, 0, false},  {"((i)i", false, 0, false},
        {"(i))i", false, 0, false},      {"())", false, 0, false}};

    for (const auto &cur_case : cases) {
        int  arg_count = 0;
        bool has_ret   = false;
        bool ok        = signature_check(cur_case.test, &arg_count, &has_ret);

        EXPECT_EQ(ok, cur_case.result);
        if (ok) {
            EXPECT_EQ(arg_count, cur_case.arg_count);
            EXPECT_EQ(has_ret, cur_case.has_ret);
        }
    }
}

TEST(DLTest, GetCloseSym) {
    init();

    int hnd = dl_open(&exec_env, "./SimpleLib.so", RTLD_LAZY);
    {
        int sym = dl_sym(&exec_env, hnd, "fib_fast", "(i)i");
        check_symbol_consistency(sym, 0);
    }

    {
        int sym = dl_sym(&exec_env, hnd, "fib_fast", "(v)i");
        ASSERT_EQ(sym, WRT_ERROR);
    }

    {
        dl_close_sym(&exec_env, hnd, 1);
        check_sym_empty(0);
    }

    dl_close(&exec_env, hnd);
    check_empty(0);
}