// Fundamental Boolean and Null types for no-std environment
#define true  1
#define false 0
#define NULL  ((void*)0)

#include "duffle.amd64.win32.h"

// --- Semantic Tags (Using X-Macros & Enum_) ---
typedef Enum_(U4, STag) {
#define Tag_Entries() \
    X(Define,  "Define",  0x003333FF, ":") /* RED */ \
    X(Call,    "Call",    0x0033FF33, "~") /* GREEN */ \
    X(Data,    "Data",    0x00FFFF33, "$") /* CYAN */ \
    X(Imm,     "Imm",     0x0033FFFF, "^") /* YELLOW */ \
    X(Comment, "Comment", 0x00888888, ".") /* DIM */ \
    X(Format,  "Format",  0x00444444, " ") /* INVISIBLE/FORMAT */

#define X(n, s, c, p) tmpl(STag, n),
    Tag_Entries()
#undef X
    STag_Count,
};

// Helper array to fetch Hex colors for UI rendering based on STag
global U4 tag_colors[] = {
#define X(n, s, c, p) c,
    Tag_Entries()
#undef X
};

// Helper array to fetch the text prefix based on STag
global const char* tag_prefixes[] = {
#define X(n, s, c, p) p,
    Tag_Entries()
#undef X
};

// Helper array to fetch the full name of the STag
global const char* tag_names[] = {
#define X(n, s, c, p) s,
    Tag_Entries()
#undef X
};

// Token Packing: 28 bits payload | 4 bits tag
#define PACK_TOKEN(tag, val) (((U4)(tag) << 28) | ((U4)(val) & 0x0FFFFFFF))
#define UNPACK_TAG(token)    (((token) >> 28) & 0x0F)
#define UNPACK_VAL(token)    ((token) & 0x0FFFFFFF)

// 2-Character Mapped Dictionary Helper
#define ID2(a, b) (((U4)(a) << 8) | (U4)(b))

#define TOKENS_PER_ROW 8

// The Tape Drive (Using FArena from duffle)
global FArena tape_arena;
global U8     cursor_idx = 0;

// Executable Code Arena (The JIT)
global FArena code_arena;

// VM State: 2-Reg Stack + Global Memory
global U8 vm_rax = 0; // Top
global U8 vm_rdx = 0; // Next
global U8 vm_globals[16] = {0};

#pragma clang optimize off
void* memset(void* dest, int c, U8 count) {
    U1* bytes = (U1*)dest;
    while (count--) *bytes++ = (U1)c;
    return dest;
}

void* memcpy(void* dest, const void* src, U8 count) {
    U1* d = (U1*)dest;
    const U1* s = (const U1*)src;
    while (count--) *d++ = *s++;
    return dest;
}
#pragma clang optimize on

IA_ void scatter(U4 token) { 
    if (tape_arena.used + sizeof(U4) <= tape_arena.capacity) {
        U4*r ptr = C_(U4*r, tape_arena.start + tape_arena.used);
        ptr[0] = token;
        tape_arena.used += sizeof(U4);
    }
}

IA_ void u64_to_hex(U8 val, char* buf, S4 chars) {
    static const char hex_chars[] = "0123456789ABCDEF";
    for(S1 i = chars - 1; i >= 0; --i) { buf[i] = hex_chars[val & 0xF]; val >>= 4; }
}

// --- Minimal x86-64 Emitter ---
internal void emit8(U1 b) { 
    if (code_arena.used + 1 <= code_arena.capacity) {
        U1* ptr = (U1*)(code_arena.start + code_arena.used);
        *ptr = b;
        code_arena.used += 1;
    }
}

internal void emit32(U4 val) { 
    if (code_arena.used + 4 <= code_arena.capacity) {
        U4* ptr = (U4*)(code_arena.start + code_arena.used);
        *ptr = val;
        code_arena.used += 4;
    }
}

internal void compile_word(U4 tag, U4 val) {
    if (tag == tmpl(STag, Data)) {
        // mov rdx, rax
        emit8(0x48); emit8(0x89); emit8(0xC2);
        // mov rax, imm32
        emit8(0x48); emit8(0xC7); emit8(0xC0); emit32(val);
    } else if (tag == tmpl(STag, Imm) || tag == tmpl(STag, Call)) {
        if (val == ID2('S','W')) { // SWAP: xchg rax, rdx
            emit8(0x48); emit8(0x87); emit8(0xC2);
        } else if (val == ID2('M','*')) { // MULT: imul rax, rdx
            emit8(0x48); emit8(0x0F); emit8(0xAF); emit8(0xC2);
        } else if (val == ID2('+',' ')) { // ADD: add rax, rdx
            emit8(0x48); emit8(0x01); emit8(0xD0);
        } else if (val == ID2('@',' ')) { // FETCH: mov rax, QWORD PTR [rcx + rax*8]
            emit8(0x48); emit8(0x8B); emit8(0x04); emit8(0xC1);
        } else if (val == ID2('-','1')) { // DEC: dec rax
            emit8(0x48); emit8(0xFF); emit8(0xC8);
        } else if (val == ID2('!',' ')) { // STORE: mov QWORD PTR [rcx + rax*8], rdx
            emit8(0x48); emit8(0x89); emit8(0x14); emit8(0xC1);
        } else if (val == ID2('R','0')) { // RET_IF_ZERO: test rax, rax; jnz +9; epilogue; ret
            emit8(0x48); emit8(0x85); emit8(0xC0); // test rax, rax
            emit8(0x75); emit8(0x09);              // jnz skip_ret (+9 bytes)
            emit8(0x48); emit8(0x89); emit8(0x41); emit8(0x70); // mov [rcx+112], rax
            emit8(0x48); emit8(0x89); emit8(0x51); emit8(0x78); // mov [rcx+120], rdx
            emit8(0xC3);                           // ret
        } else if (val == ID2('F','S')) { // F_STEP (Inlined Compiler Macro)
            compile_word(tmpl(STag, Data), 0);  
            compile_word(tmpl(STag, Imm),  ID2('@',' ')); 
            compile_word(tmpl(STag, Imm),  ID2('R','0')); 

            compile_word(tmpl(STag, Data), 1);  
            compile_word(tmpl(STag, Imm),  ID2('@',' ')); 
            compile_word(tmpl(STag, Data), 0);  
            compile_word(tmpl(STag, Imm),  ID2('@',' ')); 
            compile_word(tmpl(STag, Imm),  ID2('M','*')); 
            compile_word(tmpl(STag, Data), 1);  
            compile_word(tmpl(STag, Imm),  ID2('!',' ')); 

            compile_word(tmpl(STag, Data), 0);  
            compile_word(tmpl(STag, Imm),  ID2('@',' ')); 
            compile_word(tmpl(STag, Imm),  ID2('-','1')); 
            compile_word(tmpl(STag, Data), 0);  
            compile_word(tmpl(STag, Imm),  ID2('!',' ')); 
        }
    }
}

internal void compile_and_run_tape(void) {
    farena_reset(&code_arena); 
    
    // Prologue: Load VM state from globals[14] and [15]
    emit8(0x48); emit8(0x8B); emit8(0x41); emit8(0x70); // mov rax, [rcx+112]
    emit8(0x48); emit8(0x8B); emit8(0x51); emit8(0x78); // mov rdx, [rcx+120]

    // Compile the selected tokens
    U4*r tape_ptr = C_(U4*r, tape_arena.start);
    for (U8 i = 0; i <= cursor_idx; i++) {
        compile_word(UNPACK_TAG(tape_ptr[i]), UNPACK_VAL(tape_ptr[i]));
    }

    // Epilogue: Save VM state back to globals
    emit8(0x48); emit8(0x89); emit8(0x41); emit8(0x70); // mov [rcx+112], rax
    emit8(0x48); emit8(0x89); emit8(0x51); emit8(0x78); // mov [rcx+120], rdx
    emit8(0xC3); // ret

    // Cast code arena to function pointer and CALL it!
    typedef void (*JIT_Func)(U8* globals_ptr);
    JIT_Func func = (JIT_Func)code_arena.start;
    func(vm_globals);
    
    // Read state for UI
    vm_rax = vm_globals[14];
    vm_rdx = vm_globals[15];
}

#define MS_WM_CHAR 0x0102
#define MS_VK_RETURN 0x0D
#define MS_VK_BACK 0x08
#define MS_VK_TAB 0x09
#define MS_VK_SPACE 0x20

// --- Window Procedure ---
S8 win_proc(void* hwnd, U4 msg, U8 wparam, S8 lparam) {
    U8 tape_count = tape_arena.used / sizeof(U4);
    U4*r tape_ptr = C_(U4*r, tape_arena.start);

    switch (msg) {
        case MS_WM_CHAR: {
            U4 t = tape_ptr[cursor_idx];
            U4 tag = UNPACK_TAG(t);
            U4 val = UNPACK_VAL(t);
            U1 c = (U1)wparam;
            
            // Skip control characters in WM_CHAR (handled in KEYDOWN)
            if (c < 32) break;

            if (tag == tmpl(STag, Data)) {
                // Hex input
                U4 digit = 16;
                if (c >= '0' && c <= '9') digit = c - '0';
                if (c >= 'a' && c <= 'f') digit = c - 'a' + 10;
                if (c >= 'A' && c <= 'F') digit = c - 'A' + 10;
                if (digit < 16) {
                    val = ((val << 4) | digit) & 0x0FFFFFFF;
                    tape_ptr[cursor_idx] = PACK_TOKEN(tag, val);
                }
            } else {
                // 2-Char String input (Shift left and insert new char)
                val = ((val << 8) | c) & 0xFFFF;
                tape_ptr[cursor_idx] = PACK_TOKEN(tag, val);
            }
            
            vm_rax = 0; vm_rdx = 0; mem_zero(u8_(vm_globals), sizeof(vm_globals));
            compile_and_run_tape();
            ms_invalidate_rect(hwnd, NULL, true);
            return 0;
        }
        case MS_WM_KEYDOWN: {
            if (wparam == MS_VK_RIGHT && cursor_idx < tape_count - 1) cursor_idx++;
            if (wparam == MS_VK_LEFT && cursor_idx > 0) cursor_idx--;
            
            if (wparam == MS_VK_UP) {
                U8 line_start = cursor_idx;
                while (line_start > 0 && UNPACK_TAG(tape_ptr[line_start - 1]) != tmpl(STag, Format)) line_start--;
                if (line_start > 0) {
                    U8 col = cursor_idx - line_start;
                    U8 prev_line_start = line_start - 1;
                    while (prev_line_start > 0 && UNPACK_TAG(tape_ptr[prev_line_start - 1]) != tmpl(STag, Format)) prev_line_start--;
                    U8 prev_line_len = (line_start - 1) - prev_line_start;
                    cursor_idx = prev_line_start + (col < prev_line_len ? col : prev_line_len);
                }
            }
            if (wparam == MS_VK_DOWN) {
                U8 line_start = cursor_idx;
                while (line_start > 0 && UNPACK_TAG(tape_ptr[line_start - 1]) != tmpl(STag, Format)) line_start--;
                U8 col = cursor_idx - line_start;
                
                U8 next_line_start = cursor_idx;
                while (next_line_start < tape_count && UNPACK_TAG(tape_ptr[next_line_start]) != tmpl(STag, Format)) next_line_start++;
                if (next_line_start < tape_count) {
                    next_line_start++; // Skip the newline
                    U8 next_line_end = next_line_start;
                    while (next_line_end < tape_count && UNPACK_TAG(tape_ptr[next_line_end]) != tmpl(STag, Format)) next_line_end++;
                    U8 next_line_len = next_line_end - next_line_start;
                    cursor_idx = next_line_start + (col < next_line_len ? col : next_line_len);
                }
            }
            
            if (wparam == MS_VK_TAB) {
                // Cycle Color Tag
                U4 t = tape_ptr[cursor_idx];
                U4 tag = (UNPACK_TAG(t) + 1) % STag_Count;
                tape_ptr[cursor_idx] = PACK_TOKEN(tag, UNPACK_VAL(t));
            } else if (wparam == MS_VK_BACK) {
                // Delete Token and move cursor left
                if (tape_count > 0) {
                    for (U8 i = cursor_idx; i < tape_count - 1; i++) {
                        tape_ptr[i] = tape_ptr[i+1];
                    }
                    tape_arena.used -= sizeof(U4);
                    if (cursor_idx > 0) cursor_idx--;
                }
            } else if (wparam == MS_VK_SPACE || wparam == MS_VK_RETURN) {
                // Insert New Token (Pre-append at cursor)
                if (tape_arena.used + sizeof(U4) <= tape_arena.capacity) {
                    for (U8 i = tape_count; i > cursor_idx; i--) {
                        tape_ptr[i] = tape_ptr[i-1];
                    }
                    if (wparam == MS_VK_RETURN) {
                        tape_ptr[cursor_idx] = PACK_TOKEN(tmpl(STag, Format), 0xA); // Newline
                        cursor_idx++; // Move past newline
                    } else {
                        tape_ptr[cursor_idx] = PACK_TOKEN(tmpl(STag, Comment), ID2(' ',' '));
                        cursor_idx++; // Move past space
                    }
                    tape_arena.used += sizeof(U4);
                }
            }
            
            // Interaction: Reset VM and compile up to cursor
            vm_rax = 0; vm_rdx = 0;
            mem_zero(u8_(vm_globals), sizeof(vm_globals));
            
            compile_and_run_tape();
            
            ms_invalidate_rect(hwnd, NULL, true);
            return 0;
        }
        case MS_WM_PAINT: {
            MS_PAINTSTRUCT ps;
            void* hdc = ms_begin_paint(hwnd, &ps);
            void* hFont = ms_create_font_a(20, 0, 0, 0, 400, 0, 0, 0, 0, 0, 0, 0, 0, "Consolas");
            void* hOldFont = ms_select_object(hdc, hFont);

            ms_set_bk_color(hdc, 0x001E1E1E);
            S4 start_x = 40, start_y = 60, spacing_x = 110, spacing_y = 35;
            S4 x = start_x, y = start_y;

            U4*r tape_ptr = C_(U4*r, tape_arena.start);

            // Render Tokens
            for (U8 i = 0; i < tape_count; i++) {
                if (x >= start_x + (TOKENS_PER_ROW * spacing_x)) {
                    x = start_x; y += spacing_y;
                }

                if (i == cursor_idx) {
                    void* hBrush = ms_get_stock_object(2); // GRAY_BRUSH
                    ms_select_object(hdc, hBrush);
                    ms_rectangle(hdc, x - 5, y - 2, x + 95, y + 22);
                }

                U4 t = tape_ptr[i];
                U4 tag = UNPACK_TAG(t);
                U4 val = UNPACK_VAL(t);

                if (tag == tmpl(STag, Format) && val == 0xA) {
                    ms_set_text_color(hdc, 0x00444444);
                    ms_text_out_a(hdc, x, y, "  \\n  ", 6);
                    x = start_x;
                    y += spacing_y;
                } else {
                    U4 color = tag_colors[tag];
                    const char* prefix = tag_prefixes[tag];

                    ms_set_text_color(hdc, color);
                    
                    char val_str[9];
                    if (tag == tmpl(STag, Data)) {
                        u64_to_hex(val, val_str, 6);
                        val_str[6] = '\0';
                    } else {
                        // Decode 2-character dictionary ID
                        val_str[0] = (char)((val >> 8) & 0xFF);
                        val_str[1] = (char)(val & 0xFF);
                        val_str[2] = ' '; val_str[3] = ' '; val_str[4] = ' '; val_str[5] = ' ';
                        val_str[6] = '\0';
                    }

                    char out_buf[10];
                    out_buf[0] = prefix[0];
                    out_buf[1] = ' ';
                    mem_copy(u8_(out_buf + 2), u8_(val_str), 6);
                    out_buf[8] = '\0';
                    
                    ms_text_out_a(hdc, x, y, out_buf, 8);
                    x += spacing_x;
                }
            }

            ms_set_text_color(hdc, 0x00AAAAAA);
            ms_text_out_a(hdc, 40, 20, "x86-64 Machine Code Emitter | 2-Reg Stack + Global Tape | Factorial", 68);
            
            // HUD: Display Current Token Meaning
            if (tape_count > 0 && cursor_idx < tape_count) {
                U4 cur_tag = UNPACK_TAG(tape_ptr[cursor_idx]);
                const char* tag_name = tag_names[cur_tag];
                U4 cur_color = tag_colors[cur_tag];
                
                char semantics_str[64] = "Current Tag: ";
                U4 name_len = 0;
                while (tag_name[name_len]) {
                    semantics_str[13 + name_len] = tag_name[name_len];
                    name_len++;
                }
                semantics_str[13 + name_len] = '\0';
                
                ms_set_text_color(hdc, cur_color);
                ms_text_out_a(hdc, 400, 20, semantics_str, 13 + name_len);
            }
            
            // Render VM State
            ms_set_text_color(hdc, 0x00FFFFFF);
            char jit_str[32] = "JIT Size: 0x000 bytes";
            u64_to_hex(code_arena.used, jit_str + 12, 3);
            ms_text_out_a(hdc, 40, 480, jit_str, 21);
            
            char state_str[64] = "RAX: 00000000 | RDX: 00000000";
            u64_to_hex(vm_rax, state_str + 5, 8);
            u64_to_hex(vm_rdx, state_str + 21, 8);
            ms_set_text_color(hdc, 0x0033FF33);
            ms_text_out_a(hdc, 40, 510, state_str, 29);

            ms_set_text_color(hdc, 0x00FFFFFF);
            ms_text_out_a(hdc, 400, 480, "Global Memory (Contiguous Array):", 33);
            for (int i=0; i<4; i++) {
                char glob_str[32] = "[0]: 00000000";
                glob_str[1] = '0' + i;
                u64_to_hex(vm_globals[i], glob_str + 5, 8);
                ms_set_text_color(hdc, 0x00FFFF33);
                ms_text_out_a(hdc, 400, 510 + (i * 25), glob_str, 13);
            }

            ms_select_object(hdc, hOldFont);
            ms_end_paint(hwnd, &ps);
            return 0;
        }
        case MS_WM_DESTROY: { ms_post_quit_message(0); return 0; }
    }
    return ms_def_window_proc_a(hwnd, msg, wparam, lparam);
}

int main(void) {
    // 1. Initialize Memory Arenas using WinAPI + FArena
    Slice tape_mem = slice_ut_(u8_(ms_virtual_alloc(NULL, 64 * 1024, MS_MEM_COMMIT | MS_MEM_RESERVE, MS_PAGE_READWRITE)), 64 * 1024);
    Slice code_mem = slice_ut_(u8_(ms_virtual_alloc(NULL, 64 * 1024, MS_MEM_COMMIT | MS_MEM_RESERVE, MS_PAGE_EXECUTE_READWRITE)), 64 * 1024);
    if (!tape_mem.ptr || !code_mem.ptr) ms_exit_process(1);
    
    farena_init(&tape_arena, tape_mem);
    farena_init(&code_arena, code_mem);

    // Bootstrap Robust Sample: Factorial State Machine
    scatter(PACK_TOKEN(tmpl(STag, Comment), ID2('I','N'))); // .IN
    scatter(PACK_TOKEN(tmpl(STag, Data),    5));      // $5
    scatter(PACK_TOKEN(tmpl(STag, Data),    0));      // $0 (Addr)
    scatter(PACK_TOKEN(tmpl(STag, Imm),     ID2('!',' ')));    // ^! 
    
    scatter(PACK_TOKEN(tmpl(STag, Data),    1));      // $1
    scatter(PACK_TOKEN(tmpl(STag, Data),    1));      // $1 (Addr)
    scatter(PACK_TOKEN(tmpl(STag, Imm),     ID2('!',' ')));    // ^! 

    scatter(PACK_TOKEN(tmpl(STag, Format),  0xA)); // Newline
    
    scatter(PACK_TOKEN(tmpl(STag, Imm),     ID2('F','S')));   // ^FS
    scatter(PACK_TOKEN(tmpl(STag, Imm),     ID2('F','S')));   
    scatter(PACK_TOKEN(tmpl(STag, Imm),     ID2('F','S')));   
    scatter(PACK_TOKEN(tmpl(STag, Imm),     ID2('F','S')));   
    scatter(PACK_TOKEN(tmpl(STag, Imm),     ID2('F','S')));   

    MS_WNDCLASSA wc;
    mem_fill(u8_(& wc), 0, sizeof(wc));
    wc.lpfnWndProc = win_proc;
    wc.hInstance = ms_get_stock_object(0);
    wc.lpszClassName = "ColorForthWindow";
    wc.hbrBackground = ms_get_stock_object(4);
    ms_register_class_a(&wc);

    void* hwnd = ms_create_window_ex_a(0, wc.lpszClassName, "Sourceless Global Memory Explorer", MS_WS_OVERLAPPEDWINDOW | MS_WS_VISIBLE, 100, 100, 1100, 750, NULL, NULL, wc.hInstance, NULL);
    MS_MSG msg;
    while (ms_get_message_a(&msg, NULL, 0, 0)) { ms_translate_message(&msg); ms_dispatch_message_a(&msg); }
    ms_exit_process(0);
    return 0;
}