#include "../src/encoder/symbols.h"
#include "../src/ast.h"
#include "../src/error.h"
#include "../src/lexer.h"
#include "../src/parser/parser.h"
#include "munit.h"
#include <string.h>

void symbols_setup_test(ast_node_t **node, tokenlist_t **list, char *path) {
    lexer_t *lex = &(lexer_t){};
    lexer_open(lex, path);
    tokenlist_alloc(list);
    tokenlist_fill(*list, lex);
    parse_result_t result = parse((*list)->head);
    lexer_close(lex);

    *node = result.node;
}

MunitResult test_symbol_table_alloc(const MunitParameter params[], void *data) {
    (void)params;
    (void)data;

    symbol_table_t *table = nullptr;
    error_t *err = symbol_table_alloc(&table);

    munit_assert_ptr_not_null(table);
    munit_assert_ptr_null(err);
    munit_assert_size(table->cap, ==, 64); // Default capacity
    munit_assert_size(table->len, ==, 0);
    munit_assert_ptr_not_null(table->symbols);

    symbol_table_free(table);
    return MUNIT_OK;
}

MunitResult test_symbol_table_lookup_empty(const MunitParameter params[], void *data) {
    (void)params;
    (void)data;

    symbol_table_t *table = nullptr;
    symbol_table_alloc(&table);

    symbol_t *symbol = symbol_table_lookup(table, "nonexistent");
    munit_assert_ptr_null(symbol);

    symbol_table_free(table);
    return MUNIT_OK;
}

MunitResult test_symbol_add_reference(const MunitParameter params[], void *data) {
    (void)params;
    (void)data;
    ast_node_t *root;
    tokenlist_t *list;
    symbol_table_t *table = nullptr;
    symbols_setup_test(&root, &list, "tests/input/symbols.asm");
    symbol_table_alloc(&table);

    ast_node_t *reference = root->children[3]->children[1]->children[0]->children[0];
    munit_assert_int(reference->id, ==, NODE_LABEL_REFERENCE);
    munit_assert_size(table->len, ==, 0);

    error_t *err = symbol_table_update(table, reference);
    munit_assert_null(err);
    munit_assert_size(table->len, ==, 1);

    symbol_t *symbol = symbol_table_lookup(table, "test");
    munit_assert_not_null(symbol);
    munit_assert_int(SYMBOL_REFERENCE, ==, symbol->kind);
    munit_assert_ptr_equal(reference, symbol->node);
    munit_assert_string_equal(symbol->name, "test");

    symbol_table_free(table);
    ast_node_free(root);
    tokenlist_free(list);
    return MUNIT_OK;
}

MunitResult test_symbol_add_label(const MunitParameter params[], void *data) {
    (void)params;
    (void)data;
    ast_node_t *root;
    tokenlist_t *list;
    symbol_table_t *table = nullptr;
    symbols_setup_test(&root, &list, "tests/input/symbols.asm");
    symbol_table_alloc(&table);

    ast_node_t *label = root->children[2];
    munit_assert_int(label->id, ==, NODE_LABEL);
    munit_assert_size(table->len, ==, 0);

    error_t *err = symbol_table_update(table, label);
    munit_assert_null(err);
    munit_assert_size(table->len, ==, 1);

    symbol_t *symbol = symbol_table_lookup(table, "test");
    munit_assert_not_null(symbol);
    munit_assert_int(SYMBOL_LOCAL, ==, symbol->kind);
    munit_assert_ptr_equal(label, symbol->node);
    munit_assert_string_equal(symbol->name, "test");

    symbol_table_free(table);
    ast_node_free(root);
    tokenlist_free(list);
    return MUNIT_OK;
}

MunitResult test_symbol_add_import(const MunitParameter params[], void *data) {
    (void)params;
    (void)data;
    ast_node_t *root;
    tokenlist_t *list;
    symbol_table_t *table = nullptr;
    symbols_setup_test(&root, &list, "tests/input/symbols.asm");
    symbol_table_alloc(&table);

    ast_node_t *import_directive = root->children[0]->children[1];
    munit_assert_int(import_directive->id, ==, NODE_IMPORT_DIRECTIVE);
    munit_assert_size(table->len, ==, 0);

    error_t *err = symbol_table_update(table, import_directive);
    munit_assert_null(err);
    munit_assert_size(table->len, ==, 1);

    symbol_t *symbol = symbol_table_lookup(table, "test");
    munit_assert_not_null(symbol);
    munit_assert_int(SYMBOL_IMPORT, ==, symbol->kind);
    munit_assert_ptr_equal(import_directive, symbol->node);
    munit_assert_string_equal(symbol->name, "test");

    symbol_table_free(table);
    ast_node_free(root);
    tokenlist_free(list);
    return MUNIT_OK;
}

MunitResult test_symbol_add_export(const MunitParameter params[], void *data) {
    (void)params;
    (void)data;
    ast_node_t *root;
    tokenlist_t *list;
    symbol_table_t *table = nullptr;
    symbols_setup_test(&root, &list, "tests/input/symbols.asm");
    symbol_table_alloc(&table);

    ast_node_t *export_directive = root->children[1]->children[1];
    munit_assert_int(export_directive->id, ==, NODE_EXPORT_DIRECTIVE);
    munit_assert_size(table->len, ==, 0);

    error_t *err = symbol_table_update(table, export_directive);
    munit_assert_null(err);
    munit_assert_size(table->len, ==, 1);

    symbol_t *symbol = symbol_table_lookup(table, "test");
    munit_assert_not_null(symbol);
    munit_assert_int(SYMBOL_EXPORT, ==, symbol->kind);
    munit_assert_ptr_equal(export_directive, symbol->node);
    munit_assert_string_equal(symbol->name, "test");

    symbol_table_free(table);
    ast_node_free(root);
    tokenlist_free(list);
    return MUNIT_OK;
}

MunitTest symbols_tests[] = {
    {"/table_alloc",        test_symbol_table_alloc,        nullptr, nullptr, MUNIT_TEST_OPTION_NONE, nullptr},
    {"/table_lookup_empty", test_symbol_table_lookup_empty, nullptr, nullptr, MUNIT_TEST_OPTION_NONE, nullptr},
    {"/add_reference",      test_symbol_add_reference,      nullptr, nullptr, MUNIT_TEST_OPTION_NONE, nullptr},
    {"/add_label",          test_symbol_add_label,          nullptr, nullptr, MUNIT_TEST_OPTION_NONE, nullptr},
    {"/add_import",         test_symbol_add_import,         nullptr, nullptr, MUNIT_TEST_OPTION_NONE, nullptr},
    {"/add_export",         test_symbol_add_export,         nullptr, nullptr, MUNIT_TEST_OPTION_NONE, nullptr},
    {nullptr,               nullptr,                        nullptr, nullptr, MUNIT_TEST_OPTION_NONE, nullptr}
};