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

#define __WRAPPER
#include "pu.out.c"

/* ─── color & statistics ──────────────────────────────────────────────────── */
#define COLOR_RED   "\x1b[1;31m"
#define COLOR_BLUE  "\033[1;34m"
#define COLOR_SUM  "\033[1;33m"
#define COLOR_RESET "\x1b[0m"

static int tests_run    = 0;
static int tests_passed = 0;

/* PASS if cond true, else FAIL; color and bump counters */
#define PRINT_RESULT(cond)                                                                         \
    do {                                                                                           \
        tests_run++;                                                                               \
        if (cond) {                                                                                \
            tests_passed++;                                                                        \
            printf(COLOR_BLUE "PASS" COLOR_RESET "\n");                                            \
        } else {                                                                                   \
            printf(COLOR_RED "FAIL" COLOR_RESET "\n");                                             \
        }                                                                                          \
    } while (0)
/* ───────────────────────────────────────────────────────────────────────── */

void reset() {
    memset(m, 0, sizeof(m));
    a = 0;
    c = 0;
}

void test_opcode_0() {
    printf("Testing opcode 0 (hlt)... ");
    reset();
    m[ 0 ] = 0;
    m[ 7 ] = 0;
    pu();
    PRINT_RESULT(1);
}

void test_opcode_1() {
    printf("Testing opcode 1 (const in r0)... ");
    reset();
    m[ 0 ] = 1;
    m[ 1 ] = 42;
    m[ 2 ] = 0;
    m[ 7 ] = 0;
    pu();
    PRINT_RESULT(m[ 0 ] == 42);
}

void test_opcode_2() {
    printf("Testing opcode 4 (r0 in r2 if acc zero)... ");
    reset();
    m[ 0 ]  = 123;
    a       = 0;
    m[ 10 ] = 4;
    m[ 11 ] = 0;
    m[ 7 ]  = 10;
    pu();
    PRINT_RESULT(m[ 2 ] == 123);
}

void test_opcode_2_zero_acc() {
    printf("Testing opcode 4 with zero acc (should not store)... ");
    reset();
    m[ 0 ]  = 123;
    m[ 2 ]  = 999;
    a       = 1;
    m[ 10 ] = 4;
    m[ 11 ] = 0;
    m[ 7 ]  = 10;
    pu();
    PRINT_RESULT(m[ 2 ] == 999);
}

void test_opcode_10() {
    printf("Testing opcode 12 (r2 in r0)... ");
    reset();
    m[ 2 ]  = 456;
    m[ 10 ] = 12;
    m[ 11 ] = 0;
    m[ 7 ]  = 10;
    pu();
    PRINT_RESULT(m[ 0 ] == 456);
}

void test_opcode_18() {
    printf("Testing opcode 18 (r0 in m[r1])... ");
    reset();
    m[ 0 ]  = 789;
    m[ 1 ]  = 100;
    m[ 10 ] = 18;
    m[ 11 ] = 0;
    m[ 7 ]  = 10;
    pu();
    PRINT_RESULT(m[ 100 ] == 789);
}

void test_opcode_19() {
    printf("Testing opcode 19 (acc in r0)... ");
    reset();
    a       = 555;
    m[ 10 ] = 19;
    m[ 11 ] = 0;
    m[ 7 ]  = 10;
    pu();
    PRINT_RESULT(m[ 0 ] == 555);
}

void test_opcode_20() {
    printf("Testing opcode 20 (acc += r0)... ");
    reset();
    a       = 10;
    m[ 0 ]  = 5;
    m[ 10 ] = 20;
    m[ 11 ] = 0;
    m[ 7 ]  = 10;
    pu();
    PRINT_RESULT(a == 15);
}

void test_opcode_21() {
    printf("Testing opcode 21 (acc -= r0)... ");
    reset();
    a       = 20;
    m[ 0 ]  = 7;
    m[ 10 ] = 21;
    m[ 11 ] = 0;
    m[ 7 ]  = 10;
    pu();
    PRINT_RESULT(a == 13);
}

void test_opcode_22() {
    printf("Testing opcode 22 (acc = !acc)... ");
    reset();
    a       = 5;
    m[ 10 ] = 22;
    m[ 11 ] = 0;
    m[ 7 ]  = 10;
    pu();
    PRINT_RESULT(a == 0);
}

void test_opcode_23() {
    printf("Testing opcode 23 (acc = 0)... ");
    reset();
    a       = 999;
    m[ 10 ] = 23;
    m[ 11 ] = 0;
    m[ 7 ]  = 10;
    pu();
    PRINT_RESULT(a == 0);
}

/* ─── assembler support ───────────────────────────────────────────────────── */

i;
int check_output(int* O, int* expected, int len, const char* test_name) {
    for (i = 0; i < len; i++) {
        if (O[ i ] != expected[ i ]) {
            printf(
                "%s: output[%d]=%d (expected %d): " COLOR_RED "FAIL" COLOR_RESET "\n",
                test_name,
                i,
                O[ i ],
                expected[ i ]);
            return 0;
        }
    }
    return 1;
}

void reset_output(int* O, int size) { memset(O, 0, size * sizeof(int)); }

void test_hlt() {
    printf("Testing hlt (opcode 0)... ");
    int O[ 10 ], ok;
    reset_output(O, 10);
    int   exp[] = { 0 };
    char* ret   = as("hlt", O);
    ok          = (ret == 0 && check_output(O, exp, 1, "hlt"));
    PRINT_RESULT(ok);
}

void test_con() {
    printf("Testing con 123 (opcode 1, 123)... ");
    int O[ 10 ], ok;
    reset_output(O, 10);
    int   exp[] = { 1, 123 };
    char* ret   = as("con 123", O);
    ok          = (ret == 0 && check_output(O, exp, 2, "con"));
    PRINT_RESULT(ok);
}

void test_ar1() {
    printf("Testing ar1 (opcode 2)... ");
    int O[ 10 ], ok;
    reset_output(O, 10);
    int   exp[] = { 2 };
    char* ret   = as("ar1", O);
    ok          = (ret == 0 && check_output(O, exp, 1, "ar1"));
    PRINT_RESULT(ok);
}

void test_ar8() {
    printf("Testing ar8 (opcode 9)... ");
    int O[ 10 ], ok;
    reset_output(O, 10);
    int   exp[] = { 9 };
    char* ret   = as("ar8", O);
    ok          = (ret == 0 && check_output(O, exp, 1, "ar8"));
    PRINT_RESULT(ok);
}

void test_r1a() {
    printf("Testing r1a (opcode 10)... ");
    int O[ 10 ], ok;
    reset_output(O, 10);
    int   exp[] = { 10 };
    char* ret   = as("r1a", O);
    ok          = (ret == 0 && check_output(O, exp, 1, "r1a"));
    PRINT_RESULT(ok);
}

void test_r8a() {
    printf("Testing r8a (opcode 17)... ");
    int O[ 10 ], ok;
    reset_output(O, 10);
    int   exp[] = { 17 };
    char* ret   = as("r8a", O);
    ok          = (ret == 0 && check_output(O, exp, 1, "r8a"));
    PRINT_RESULT(ok);
}

void test_mem() {
    printf("Testing mem (opcode 18)... ");
    int O[ 10 ], ok;
    reset_output(O, 10);
    int   exp[] = { 18 };
    char* ret   = as("mem", O);
    ok          = (ret == 0 && check_output(O, exp, 1, "mem"));
    PRINT_RESULT(ok);
}

void test_aca() {
    printf("Testing aca (opcode 19)... ");
    int O[ 10 ], ok;
    reset_output(O, 10);
    int   exp[] = { 19 };
    char* ret   = as("aca", O);
    ok          = (ret == 0 && check_output(O, exp, 1, "aca"));
    PRINT_RESULT(ok);
}

void test_add() {
    printf("Testing add (opcode 20)... ");
    int O[ 10 ], ok;
    reset_output(O, 10);
    int   exp[] = { 20 };
    char* ret   = as("add", O);
    ok          = (ret == 0 && check_output(O, exp, 1, "add"));
    PRINT_RESULT(ok);
}

void test_sub() {
    printf("Testing sub (opcode 21)... ");
    int O[ 10 ], ok;
    reset_output(O, 10);
    int   exp[] = { 21 };
    char* ret   = as("sub", O);
    ok          = (ret == 0 && check_output(O, exp, 1, "sub"));
    PRINT_RESULT(ok);
}

void test_not() {
    printf("Testing not (opcode 22)... ");
    int O[ 10 ], ok;
    reset_output(O, 10);
    int   exp[] = { 22 };
    char* ret   = as("not", O);
    ok          = (ret == 0 && check_output(O, exp, 1, "not"));
    PRINT_RESULT(ok);
}

void test_clr() {
    printf("Testing clr (opcode 23)... ");
    int O[ 10 ], ok;
    reset_output(O, 10);
    int   exp[] = { 23 };
    char* ret   = as("clr", O);
    ok          = (ret == 0 && check_output(O, exp, 1, "clr"));
    PRINT_RESULT(ok);
}

void test_multi_instruction() {
    printf("Testing multi-instruction (hlt con 123 ar1 clr)... ");
    int O[ 10 ], ok;
    reset_output(O, 10);
    int   exp[] = { 0, 1, 123, 2, 23 };
    char* ret   = as("hlt con 123 ar1 clr", O);
    ok          = (ret == 0 && check_output(O, exp, 5, "multi-instruction"));
    PRINT_RESULT(ok);
}

void test_whitespace() {
    printf("Testing whitespace handling (hlt  \\t\\n con 456 \\n clr)... ");
    int O[ 10 ], ok;
    reset_output(O, 10);
    int   exp[] = { 0, 1, 456, 23 };
    char* ret   = as("hlt  \t\n con 456 \n clr", O);
    ok          = (ret == 0 && check_output(O, exp, 4, "whitespace"));
    PRINT_RESULT(ok);
}

void test_comment() {
    printf("Testing comment handling (should ignore line)... ");
    int O[ 10 ], ok;
    reset_output(O, 10);
    int   exp[] = { 1, 90 };
    char* ret   = as("# hello\ncon 90", O);
    ok          = (ret == 0 && check_output(O, exp, 2, "comment"));
    PRINT_RESULT(ok);
}

void test_error() {
    printf("Testing error case (invalid instruction)... ");
    int O[ 10 ], ok;
    reset_output(O, 10);
    char* ret = as("xyz", O);
    ok        = (ret != 0);
    PRINT_RESULT(ok);
}

int main() {
    printf(COLOR_SUM "Running CPU unit tests...\n\n" COLOR_RESET);

    test_opcode_0();
    test_opcode_1();
    test_opcode_2();
    test_opcode_2_zero_acc();
    test_opcode_10();
    test_opcode_18();
    test_opcode_19();
    test_opcode_20();
    test_opcode_21();
    test_opcode_22();
    test_opcode_23();

    printf(COLOR_SUM "\nRunning assembler unit tests...\n\n" COLOR_RESET);

    test_hlt();
    test_con();
    test_ar1();
    test_ar8();
    test_r1a();
    test_r8a();
    test_mem();
    test_aca();
    test_add();
    test_sub();
    test_not();
    test_clr();
    test_multi_instruction();
    test_whitespace();
    test_comment();
    test_error();

    printf(
        "\nTotal tests run: %d, passed: %d, failed: %d\n",
        tests_run,
        tests_passed,
        tests_run - tests_passed);

    return 0;
}