#include "runner.h"
#include "../colla.h"
#include <stdio.h>

UNIT_TEST(highlight_init) {
    arena_t arena = arena_make(ARENA_MALLOC, KB(4));
    
    // Create a basic configuration
    hl_config_t config = {0};
    
    // Define custom colors
    config.colors[HL_COLOR_NORMAL] = strv_init("normal");
    config.colors[HL_COLOR_KEYWORDS] = strv_init("keyword");
    config.colors[HL_COLOR_STRING] = strv_init("string");
    config.colors[HL_COLOR_COMMENT] = strv_init("comment");
    
    // Initialize highlighter
    hl_ctx_t *ctx = hl_init(&arena, &config);
    ASSERT(ctx != NULL);
    
    arena_cleanup(&arena);
}

UNIT_TEST(highlight_basic) {
    arena_t arena = arena_make(ARENA_MALLOC, KB(4));
    
    // Create a configuration
    hl_config_t config = {0};
    
    // Define custom colors for HTML output
    config.colors[HL_COLOR_NORMAL] = strv_init("color:black");
    config.colors[HL_COLOR_KEYWORDS] = strv_init("color:blue");
    config.colors[HL_COLOR_STRING] = strv_init("color:green");
    config.colors[HL_COLOR_COMMENT] = strv_init("color:gray");
    config.colors[HL_COLOR_NUMBER] = strv_init("color:purple");
    
    // Set HTML output flag
    config.flags = HL_FLAG_HTML;
    
    // Initialize highlighter
    hl_ctx_t *ctx = hl_init(&arena, &config);
    ASSERT(ctx != NULL);
    
    // Sample C code to highlight
    strview_t code = strv_init(
        "// This is a comment\n"
        "int main() {\n"
        "    printf(\"Hello, World!\\n\");\n"
        "    return 0;\n"
        "}\n"
    );
    
    // Highlight the code
    str_t highlighted = hl_highlight(&arena, ctx, code);
    
    // Verify the output
    ASSERT(!str_is_empty(highlighted));
    
    // We can't easily test the exact output without parsing HTML,
    // but we can check that key strings are present
    const char *html_start = "<span";
    ASSERT(strstr(highlighted.buf, html_start) != NULL);
    
    // Check if comment coloring is present
    ASSERT(strstr(highlighted.buf, "// This is a comment") != NULL);
    
    // Check if string highlighting is present
    ASSERT(strstr(highlighted.buf, "Hello, World!") != NULL);
    
    arena_cleanup(&arena);
}

UNIT_TEST(highlight_custom_keywords) {
    arena_t arena = arena_make(ARENA_MALLOC, KB(4));
    
    // Create configuration
    hl_config_t config = {0};
    
    // Define colors
    config.colors[HL_COLOR_NORMAL] = strv_init("normal");
    config.colors[HL_COLOR_KEYWORDS] = strv_init("keyword");
    config.colors[HL_COLOR_CUSTOM_TYPES] = strv_init("custom-type");
    
    // Define custom keywords
    hl_keyword_t custom_keywords[] = {
        {.keyword = strv_init("MyClass"), .color = HL_COLOR_CUSTOM_TYPES},
        {.keyword = strv_init("custom_func"), .color = HL_COLOR_FUNC}
    };
    
    config.extra_kwrds = custom_keywords;
    config.kwrds_count = 2;
    
    // Initialize highlighter
    hl_ctx_t *ctx = hl_init(&arena, &config);
    ASSERT(ctx != NULL);
    
    // Test code with custom keywords
    strview_t code = strv_init(
        "MyClass obj;\n"
        "custom_func(obj);\n"
    );
    
    // Highlight the code
    str_t highlighted = hl_highlight(&arena, ctx, code);
    
    // Verify output contains our code
    ASSERT(!str_is_empty(highlighted));
    ASSERT(strstr(highlighted.buf, "MyClass") != NULL);
    ASSERT(strstr(highlighted.buf, "custom_func") != NULL);
    
    arena_cleanup(&arena);
}

UNIT_TEST(highlight_add_keyword) {
    arena_t arena = arena_make(ARENA_MALLOC, KB(4));
    
    // Create configuration
    hl_config_t config = {0};
    
    // Define colors
    config.colors[HL_COLOR_NORMAL] = strv_init("normal");
    config.colors[HL_COLOR_KEYWORDS] = strv_init("keyword");
    config.colors[HL_COLOR_CUSTOM_TYPES] = strv_init("custom-type");
    
    // Initialize highlighter
    hl_ctx_t *ctx = hl_init(&arena, &config);
    ASSERT(ctx != NULL);
    
    // Add a custom keyword after initialization
    hl_keyword_t new_keyword = {
        .keyword = strv_init("NewKeyword"),
        .color = HL_COLOR_CUSTOM_TYPES
    };
    
    hl_add_keyword(&arena, ctx, &new_keyword);
    
    // Test code with the new keyword
    strview_t code = strv_init("NewKeyword x;\n");
    
    // Highlight the code
    str_t highlighted = hl_highlight(&arena, ctx, code);
    
    // Verify output contains our code
    ASSERT(!str_is_empty(highlighted));
    ASSERT(strstr(highlighted.buf, "NewKeyword") != NULL);
    
    arena_cleanup(&arena);
}

UNIT_TEST(highlight_symbol_table) {
    arena_t arena = arena_make(ARENA_MALLOC, KB(4));
    
    // Create configuration
    hl_config_t config = {0};
    
    // Define colors
    config.colors[HL_COLOR_NORMAL] = strv_init("normal");
    config.colors[HL_COLOR_SYMBOL] = strv_init("symbol");
    
    // Initialize highlighter
    hl_ctx_t *ctx = hl_init(&arena, &config);
    ASSERT(ctx != NULL);
    
    // Set '@' as a symbol
    hl_set_symbol_in_table(ctx, '@', true);
    
    // Test code with the symbol
    strview_t code = strv_init("@decorator\n");
    
    // Highlight the code
    str_t highlighted = hl_highlight(&arena, ctx, code);
    
    // Verify output contains our code
    ASSERT(!str_is_empty(highlighted));
    ASSERT(strstr(highlighted.buf, "@") != NULL);
    
    arena_cleanup(&arena);
}