#include "test.h"

#include "util.h"
#include "parser.h"

#include <setjmp.h>
#include <stdio.h>
#include <string.h>
#include <stdlib.h>

static jmp_buf s_testJmp;
static char s_failMsg[1024];
static char* s_logOutput = NULL;
static const char* s_currentTestName = NULL;
static char* s_testSource = NULL;

static Module* s_currentModule = NULL;
static TokenStream* s_currentTokenStream = NULL;

void fail(const char* msg) {
    if (msg) {
        strncpy(s_failMsg, msg, sizeof(s_failMsg) - 1);
        s_failMsg[sizeof(s_failMsg) - 1] = '\0';
    } else {
        s_failMsg[0] = '\0';
    }
    longjmp(s_testJmp, 1);
}

char* read_file_content(const char* filepath) {
    FILE* f;
    long size;
    char* content;

    f = fopen(filepath, "r");
    if (!f) return NULL;
    fseek(f, 0, SEEK_END);
    size = ftell(f);
    fseek(f, 0, SEEK_SET);
    content = malloc(size + 1);
    if (!content) {
        fclose(f);
        return NULL;
    }
    fread(content, 1, size, f);
    content[size] = '\0';
    fclose(f);
    return content;
}

void assert_not_null(void* ptr, const char* msg) {
    if (ptr == NULL) {
        fail(msg);
    }
}

void assert_string(const char* expected, String actual, const char* msg) {
    if (expected == NULL || actual.data == NULL || strlen(expected) != actual.length || strncmp(expected, actual.data, actual.length) != 0) {
        fail(msg);
    }
}

void assert_str(const char* expected, const char* actual, const char* msg) {
    if (expected == NULL || actual == NULL || strcmp(expected, actual) != 0) {
        fail(msg);
    }
}

TokenStream* test_get_tokenstream(void) {
    if (s_currentTokenStream == NULL) {
        char* filepath = NULL;

        filepath = format_string("v0/tests/%s.c2", s_currentTestName);
        if (!filepath) {
            fail("out of memory");
            return NULL;
        }
        
        if (s_testSource) free(s_testSource);
        s_testSource = read_file_content(filepath);
        if (!s_testSource) {
            puts(filepath);
            free(filepath);
            fail("could not read test source file");
            return NULL;
        }
        s_currentTokenStream = tokenstream_open(filepath, s_testSource);
        free(filepath);
    }
    return s_currentTokenStream;
}

Module* test_get_ast(void) {
    if (s_currentModule == NULL) {
        s_currentModule = parser_parse(test_get_tokenstream());
    }
    return s_currentModule;
}

void assert_log(const char* expected, const char* msg) {
    assert_str(expected, s_logOutput, msg);
}

void assert_log_file(const char* msg) {
    char* filepath = format_string("v0/tests/%s.log", s_currentTestName);
    const char* generate;
    char* content;
    if (!filepath) {
        fail("out of memory");
        return;
    }
    generate = getenv("GENERATE_GOLDEN");
    if (generate && strcmp(generate, "1") == 0) {
        FILE* f = fopen(filepath, "w");
        if (!f) {
            free(filepath);
            fail("could not open golden file for writing");
            return;
        }
        fputs(s_logOutput ? s_logOutput : "", f);
        fclose(f);
        free(filepath);
        return;
    }

    content = read_file_content(filepath);
    if (!content) {
        free(filepath);
        fail("could not open golden file for reading");
        return;
    }

    bool match = strcmp(content, s_logOutput ? s_logOutput : "") == 0;
    free(content);
    free(filepath);
    if (!match) {
        fail(msg);
    }
}

void assert_int(int expected, int actual, const char* msg) {
    if (expected != actual) {
        char* buf = format_string("%s (expected %d, got %d)", msg, expected, actual);
        if (buf) {
            fail(buf);
            free(buf);
        } else {
            fail("out of memory");
        }
    }
}

void assert_true(bool condition, const char* msg) {
    if (!condition) {
        fail(msg);
    }
}

void assert_false(bool condition, const char* msg) {
    if (condition) {
        fail(msg);
    }
}

static void log_append(const char* msg) {
    size_t oldLen = s_logOutput ? strlen(s_logOutput) : 0;
    size_t newLen = oldLen + strlen(msg) + 1;
    char* newOutput = malloc(newLen);
    if (newOutput) {
        if (s_logOutput) {
            strcpy(newOutput, s_logOutput);
            free(s_logOutput);
        } else {
            newOutput[0] = '\0';
        }
        strcat(newOutput, msg);
        s_logOutput = newOutput;
    }
}

static void log_clear(void) {
    free(s_logOutput);
    s_logOutput = NULL;
}

typedef struct {
    const char* name;
    Test func;
} TestCase;

#include "test_token.c"
#include "test_parser.c"
#include "test_log.c"

static int s_totalTests;
static int s_greenTests;

#define TEST(name) {#name, name},

static TestCase s_tests[] = {
    TEST(test_log_error)
    TEST(test_log_on_line_variadic)
    TEST(test_log_on_line)
    TEST(test_parser_alias_array)
    TEST(test_parser_alias_simple)
    TEST(test_parser_bad_import_name)
    TEST(test_parser_bad_module_name)
    TEST(test_parser_imports)
    TEST(test_parser_missing_semicolon_import)
    TEST(test_parser_missing_semicolon_module)
    TEST(test_parser_module_name)
    TEST(test_parser_public_imports)
    TEST(test_parser_variable_simple)
    TEST(test_tokenstream_comma)
    TEST(test_tokenstream_info)
    TEST(test_tokenstream_keywords_and_symbols)
    TEST(test_tokenstream_open_fail)
    TEST(test_tokenstream_parentheses_and_brackets)
    TEST(test_tokenstream_primitive_types)
    TEST(test_tokenstream_simple_keyword)
    TEST(test_tokenstream_unknown_token)
    TEST(test_tokenstream_void_function_signature)
    TEST(test_tokenstream_whitespace_ignored)
};


int main(int argc, char** argv) {
    const char** failedTests;
    int failedCount;

    (void)argc;
    (void)argv;

    s_totalTests = sizeof(s_tests) / sizeof(s_tests[0]);
    s_greenTests = 0;

    // Allocate failed tests array dynamically to avoid VLAs
    failedTests = (const char**)malloc((s_totalTests + 1) * sizeof(const char*));
    failedCount = 0;

    for (int i = 0; i < s_totalTests; i++) {
        // Add 5 to strip the 'test_' prefix.
        s_currentTestName = s_tests[i].name + 5;
        log_set_output(log_append);
        printf("%s...", s_tests[i].name);
        fflush(stdout);
        s_failMsg[0] = '\0';

        if (setjmp(s_testJmp) == 0) {
            log_clear();
            if (s_testSource) {
                free(s_testSource);
                s_testSource = NULL;
            }
            s_tests[i].func();
            printf(" [OK]\n");
            s_greenTests++;
        } else {
            printf(" [FAIL]: %s\n", s_failMsg[0] ? s_failMsg : "");
            failedTests[failedCount++] = s_tests[i].name;
            
            // Log output on failure
            if (s_logOutput && s_logOutput[0]) {
                printf("%s\n", s_logOutput);
            }
        }

        // Free AST and TokenStream after each test
        if (s_currentModule) {
            parser_free(s_currentModule);
            s_currentModule = NULL;
        }
        if (s_currentTokenStream) {
            tokenstream_close(s_currentTokenStream);
            s_currentTokenStream = NULL;
        }
        fflush(stdout);
    }

    if (s_testSource) free(s_testSource);
    log_clear();

    if (failedCount > 0) {
        printf("\nFailed tests:\n");
        for (int j = 0; j < failedCount; j++) {
            printf("  - %s\n", failedTests[j]);
        }
    }

    printf("\n%d/%d tests passed.\n", s_greenTests, s_totalTests);
    free(failedTests);
    return failedCount > 0 ? 1 : 0;
}