#include "../lib/priv.h"

#include <assert.h>
#include <math.h>
#include <stdio.h>
#include <string.h>
#include <stdlib.h>

#include "lua.h"
#include "lauxlib.h"
#include "lualib.h"

#define EQ(a, b) (memcmp(a, b) == 0)

static void run_assembly_tests(void);
static void run_assembly_test(lua_State* L);
static void run_assembly_test_code(lua_State* L, bool debug, bool decompile);
static void run_assembly_test_template(lua_State* L, bool debug, bool decompile);

int main(void)
{
    {
        printf("## Values\n");
        assert(value_type(create_value_integer(42)) == TT_INTEGER);
        assert(value_integer(create_value_integer(-42)) == -42);
        assert(fabsf(value_real(create_value_real(42.4f)) - 42.4f) < 0.00001f);
        assert(value_idx(create_value_idx(TT_FUNCTION, 42)) == 42);
    }

    {
        printf("## Stack\n");

        Stack* s = stack_new();

        stack_push(s, create_value_integer(10));
        stack_push(s, create_value_integer(20));
        stack_push(s, create_value_integer(30));

        VALUE v;
        assert(stack_size(s) == 3);
        assert(stack_at(s, 0, &v) == T_OK); assert(value_integer(v) == 10);
        assert(stack_at(s, 1, &v) == T_OK); assert(value_integer(v) == 20);
        assert(stack_at(s, -1, &v) == T_OK); assert(value_integer(v) == 30);
        assert(stack_at(s, -2, &v) == T_OK); assert(value_integer(v) == 20);

        assert(stack_at(s, 3, &v) == T_ERR_STACK_ACCESS_OUT_OF_RANGE);
        assert(stack_at(s, -4, &v) == T_ERR_STACK_ACCESS_OUT_OF_RANGE);

        assert(stack_set(s, 1, create_value_integer(99)) == T_OK);
        assert(stack_at(s, 1, &v) == T_OK); assert(value_integer(v) == 99);
        assert(stack_at(s, -2, &v) == T_OK); assert(value_integer(v) == 99);

        assert(stack_pop(s, NULL) == T_OK);
        assert(stack_pop(s, NULL) == T_OK);
        assert(stack_at(s, -1, &v) == T_OK); assert(value_integer(v) == 10);
        assert(stack_pop(s, NULL) == T_OK);
        assert(stack_size(s) == 0);

        assert(stack_pop(s, NULL) == T_ERR_STACK_UNDERFLOW);

        stack_destroy(s);
    }

    {
        printf("## Stack with frame pointer\n");

        Stack* s = stack_new();

        stack_push(s, create_value_integer(10));
        stack_push(s, create_value_integer(20));
        stack_push_fp(s);
        stack_push(s, create_value_integer(30));
        stack_push(s, create_value_integer(40));
        stack_push(s, create_value_integer(50));

        VALUE v;
        assert(stack_size(s) == 3);
        assert(stack_at(s, 0, &v) == T_OK); assert(value_integer(v) == 30);
        assert(stack_at(s, 1, &v) == T_OK); assert(value_integer(v) == 40);
        assert(stack_at(s, -1, &v) == T_OK); assert(value_integer(v) == 50);
        assert(stack_at(s, -2, &v) == T_OK); assert(value_integer(v) == 40);

        assert(stack_set(s, -2, create_value_integer(99)) == T_OK);
        assert(stack_at(s, -2, &v) == T_OK); assert(value_integer(v) == 99);

        assert(stack_at(s, 3, &v) == T_ERR_STACK_ACCESS_OUT_OF_RANGE);
        assert(stack_at(s, -4, &v) == T_ERR_STACK_ACCESS_OUT_OF_RANGE);

        stack_pop_fp(s);

        assert(stack_size(s) == 2);
        assert(stack_at(s, 0, &v) == T_OK); assert(value_integer(v) == 10);
        assert(stack_at(s, 1, &v) == T_OK); assert(value_integer(v) == 20);
        assert(stack_at(s, -1, &v) == T_OK); assert(value_integer(v) == 20);
        assert(stack_at(s, -2, &v) == T_OK); assert(value_integer(v) == 10);

        assert(stack_at(s, 2, &v) == T_ERR_STACK_ACCESS_OUT_OF_RANGE);
        assert(stack_at(s, -3, &v) == T_ERR_STACK_ACCESS_OUT_OF_RANGE);

        stack_destroy(s);
    }

    {
        printf("## Arrays\n");

        Array* a = array_new();
        assert(array_len(a) == 0);

        array_set(a, 1, create_value_integer(40));
        assert(array_len(a) == 2);
        assert(value_type(array_get(a, 0)) == TT_NIL);
        assert(value_type(array_get(a, 1)) == TT_INTEGER);

        array_append(a, create_value_integer(50));
        assert(array_len(a) == 3);
        assert(value_integer(array_get(a, 2)) == 50);

        array_set(a, 2, create_value_integer(60));
        assert(array_len(a) == 3);
        assert(value_integer(array_get(a, 2)) == 60);

        array_destroy(a);
    }

    {
        printf("## Table - integer index\n");

        Heap* h = heap_new();
        Table* t = table_new(h);

        table_set(t, create_value_integer(10), create_value_integer(100));
        table_set(t, create_value_integer(20), create_value_integer(200));

        VALUE v;
        assert(table_get(t, create_value_integer(10), &v) == T_OK); assert(value_integer(v) == 100);
        assert(table_get(t, create_value_integer(20), &v) == T_OK); assert(value_integer(v) == 200);

        table_del(t, create_value_integer(20));
        assert(table_get(t, create_value_integer(10), &v) == T_OK);
        assert(table_get(t, create_value_integer(20), &v) == T_ERR_TABLE_KEY_NOT_FOUND);

        table_destroy(t);
        heap_destroy(h);
    }

    {
        printf("## Table - string index\n");

        Heap* h = heap_new();
        Table* t = table_new(h);

        VALUE key1 = create_value_idx(TT_STRING, heap_add_string(h, "key1"));
        VALUE key2 = create_value_idx(TT_STRING, heap_add_string(h, "key2"));

        table_set(t, key1, create_value_integer(100));
        table_set(t, key2, create_value_integer(200));

        VALUE key1b = create_value_idx(TT_STRING, heap_add_string(h, "key1"));
        VALUE key2b = create_value_idx(TT_STRING, heap_add_string(h, "key2"));

        VALUE v;
        assert(table_get(t, key1b, &v) == T_OK); assert(value_integer(v) == 100);
        assert(table_get(t, key2b, &v) == T_OK); assert(value_integer(v) == 200);

        table_del(t, key2b);
        assert(table_get(t, key1b, &v) == T_OK);
        assert(table_get(t, key2b, &v) == T_ERR_TABLE_KEY_NOT_FOUND);

        table_destroy(t);
        heap_destroy(h);
    }

    {
        printf("## Heap - strings\n");

        Heap* h = heap_new();

        HEAP_KEY key1 = heap_add_string(h, "hello");
        HEAP_KEY key2 = heap_add_string(h, "world");

        const char* value;
        assert(heap_get_string(h, key1, &value) == T_OK); assert(strcmp(value, "hello") == 0);
        assert(heap_get_string(h, key2, &value) == T_OK); assert(strcmp(value, "world") == 0);
        assert(heap_get_string(h, 1000, &value) == T_ERR_HEAP_KEY_NOT_FOUND);

        heap_destroy(h);
    }

    {
        printf("## Heap - string GC\n");

        Stack* s = stack_new();
        Heap* h = heap_new();

        stack_push(s, create_value_idx(TT_STRING, heap_add_string(h, "item1")));
        stack_push(s, create_value_idx(TT_STRING, heap_add_string(h, "item2")));
        stack_push(s, create_value_idx(TT_STRING, heap_add_string(h, "item3")));

        size_t v_sz;
        VALUE* v_idx;

        assert(heap_size(h) == 3);
        v_sz = stack_collectable_array(s, &v_idx);
        heap_gc(h, v_idx, v_sz);
        free(v_idx);
        assert(heap_size(h) == 3);

        stack_pop(s, NULL);

        assert(heap_size(h) == 3);
        v_sz = stack_collectable_array(s, &v_idx);
        heap_gc(h, v_idx, v_sz);
        free(v_idx);
        assert(heap_size(h) == 2);

        stack_pop(s, NULL);
        v_sz = stack_collectable_array(s, &v_idx);
        heap_gc(h, v_idx, v_sz);
        free(v_idx);
        assert(heap_size(h) == 1);

        heap_destroy(h);
        stack_destroy(s);
    }

    {
        printf("## Bytecode\n");
        const char* assembly_code =
                ".const\n"
                "    0: 3.14\n"
                "    1: \"Hello world\"\n"
                "\n"
                ".func 0\n"
                "    pushi   2   ; this is a comment\n"
                "    pushi   -3\n"
                "    sum\n"
                "    ret\n"
                ".func 1\n"
                "    pushi   5000\n"
                "    ret";

        uint8_t* bytecode; size_t bytecode_sz;
        assert(code_assemble(assembly_code, &bytecode, &bytecode_sz) == T_OK);

        Code* code = code_new();

        assert(code_load_bytecode(code, bytecode, bytecode_sz) == T_OK);

        assert(code_n_consts(code) == 2);
        assert(code_const_type(code, 0) == TC_REAL);
        assert(code_const_type(code, 1) == TC_STRING);
        assert(code_const_real(code, 0) > 3.13f && code_const_real(code, 0) < 3.15f);
        assert(strcmp(code_const_string(code, 1), "Hello world") == 0);
        assert(code_n_functions(code) == 2);
        assert(code_function_sz(code, 0) == 6);
        assert(code_function_sz(code, 1) == 4);

        uint32_t addr = 0;
        Instruction inst = code_next_instruction(code, 0, addr);
        assert(inst.operator == TO_PUSHI);
        assert(inst.operand == 2);
        assert(inst.sz == 2);
        addr += inst.sz;

        inst = code_next_instruction(code, 0, addr);
        assert(inst.operator == TO_PUSHI);
        assert(inst.operand == -3);
        addr += inst.sz;

        inst = code_next_instruction(code, 0, addr);
        assert(inst.operator == TO_SUM);
        assert(inst.operand == 0);
        addr += inst.sz;

        inst = code_next_instruction(code, 1, 0);
        assert(inst.operator == TO_PUSHI);
        assert(inst.operand == 5000);
        assert(inst.sz == 3);

        code_destroy(code);
        free(bytecode);
    }

    {
        printf("## Bytecode - labels\n");
        const char* assembly_code =
                ".func 0\n"
                "    jmp    @my_label\n"
                "    pushi  \n"
                "@my_label:\n"
                "    ret";

        uint8_t* bytecode; size_t bytecode_sz;
        assert(code_assemble(assembly_code, &bytecode, &bytecode_sz) == T_OK);

        Code* code = code_new();
        assert(code_load_bytecode(code, bytecode, bytecode_sz) == T_OK);

        Instruction inst = code_next_instruction(code, 0, 0);
        assert(inst.operator == TO_JMP);
        assert(inst.operand == 4);
        assert(inst.sz == 3);

        code_destroy(code);
        free(bytecode);
    }

    {
        printf("## VM - Basic\n");

        TycheVM* T = tyc_new();

        tyc_pushinteger(T, 2);
        tyc_pushinteger(T, 3);
        assert(tyc_expr(T, TX_SUM) == T_OK);
        int32_t result; assert(tyc_tointeger(T, -1, &result) == T_OK);
        assert(result == 5);

        tyc_destroy(T);
    }

    {
        printf("## Assembly tests\n");
        run_assembly_tests();
    }
}

static void run_assembly_tests(void)
{
    lua_State* L = luaL_newstate();
    luaL_openlibs(L);

    int r = luaL_loadfile(L, "./test/code-tests.lua");
    assert(r == LUA_OK);
    lua_call(L, 0, 1);
    assert(lua_istable(L, -1));

    size_t len = (size_t) luaL_len(L, -1);
    for (size_t i = 0; i < len; ++i) {
        lua_geti(L, -1, (int)i + 1);
        run_assembly_test(L);
        lua_pop(L, 1);
    }

    lua_close(L);
}

static void run_assembly_test(lua_State* L)
{
    // print test name
    lua_getfield(L, -1, "name");
    printf("   - %s\n", lua_tostring(L, -1));
    lua_pop(L, 1);

    // debug?
    lua_getfield(L, -1, "debug");
    bool debug = lua_isboolean(L, -1) && lua_toboolean(L, -1);
    lua_pop(L, 1);

    // decompile
    lua_getfield(L, -1, "decompile");
    bool decompile = lua_isboolean(L, -1) && lua_toboolean(L, -1);
    lua_pop(L, 1);

    // has code?
    lua_getfield(L, -1, "code");
    if (!lua_isnil(L, -1)) {
        lua_pop(L, 1);
        run_assembly_test_code(L, debug, decompile);
        return;
    } else {
        lua_pop(L, 1);
    }

    // has template
    lua_getfield(L, -1, "template");
    if (!lua_isnil(L, -1)) {
        lua_pop(L, 1);
        run_assembly_test_template(L, debug, decompile);
    } else {
        lua_pop(L, 1);
    }
}

static void run_assembly_test_code(lua_State* L, bool debug, bool decompile)
{
    TycheVM* T = tyc_new();
    tyc_debug_to_console(T, debug);

    // load code
    uint8_t* bytecode; size_t bytecode_sz;
    lua_getfield(L, -1, "code");
    assert(code_assemble(lua_tostring(L, -1), &bytecode, &bytecode_sz) == T_OK);
    lua_pop(L, 1);

    // run code
    assert(tyc_load_bytecode(T, bytecode, bytecode_sz) == T_OK);
    if (decompile)
        tyc_assembly_decompile(T);
    assert(tyc_call(T, 0) == T_OK);

    // check stack size
    lua_getfield(L, -1, "expected_stack_size");
    if (!lua_isnil(L, -1))
        assert(tyc_stack_size(T) == (size_t) lua_tointeger(L, -1));
    lua_pop(L, 1);

    // check stack top
    lua_getfield(L, -1, "expected_stack_top");
    if (lua_isinteger(L, -1)) {
        TYC_TYPE type; assert(tyc_type(T, -1, &type) == T_OK); assert(type == TT_INTEGER);
        int32_t v; assert(tyc_tointeger(T, -1, &v) == T_OK); assert(v == lua_tointeger(L, -1));
    } else if (!lua_isnil(L, -1)) {
        abort();
    }
    lua_pop(L, 1);

    // cleanup
    free(bytecode);
    tyc_destroy(T);
}

static void run_assembly_test_template(lua_State* L, bool debug, bool decompile)
{
    lua_getfield(L, -1, "template");
    char* template = strdup(lua_tostring(L, -1));
    lua_pop(L, 1);

    lua_getfield(L, -1, "scenarios");
    assert(!lua_isnil(L, -1));

    long n_scenarios = luaL_len(L, -1);
    for (long i = 0; i < n_scenarios; ++i) {
        lua_geti(L, -1, (int)i + 1);

        lua_getfield(L, -1, "name");
        printf("     .. %s\n", lua_tostring(L, -1));
        lua_pop(L, 1);

        // format code
        luaL_dostring(L, "return string.format");
        assert(lua_isfunction(L, -1));
        lua_pushstring(L, template);

        lua_getfield(L, -3, "parameters");
        assert(!lua_isnil(L, -1));
        int n_params = (int) luaL_len(L, -1);
        for (int j = 0; j < n_params; ++j)
            lua_geti(L, -(j + 1), j + 1);
        lua_remove(L, -(n_params + 1));

        lua_call(L, n_params + 1, 1);
        char* formatted_code = strdup(lua_tostring(L, -1));
        lua_pop(L, 1);

        // run code
        TycheVM* T = tyc_new();
        tyc_debug_to_console(T, debug);
        uint8_t* bytecode; size_t bytecode_sz;
        assert(code_assemble(formatted_code, &bytecode, &bytecode_sz) == T_OK);
        assert(tyc_load_bytecode(T, bytecode, bytecode_sz) == T_OK);
        if (decompile)
            tyc_assembly_decompile(T);
        assert(tyc_call(T, 0) == T_OK);

        // check stack top
        lua_getfield(L, -1, "expected_stack_top");
        if (lua_isinteger(L, -1)) {
            TYC_TYPE type; assert(tyc_type(T, -1, &type) == T_OK); assert(type == TT_INTEGER);
            int32_t v; assert(tyc_tointeger(T, -1, &v) == T_OK); assert(v == lua_tointeger(L, -1));
        } else if (!lua_isnil(L, -1)) {
            abort();
        }
        lua_pop(L, 1);

        // cleanup
        free(bytecode);
        tyc_destroy(T);
        free(formatted_code);
        lua_pop(L, 1);
    }

    lua_pop(L, 1);
    free(template);
}