#include "pu2.out.c"
#include <stdio.h>
#include <string.h>

/* ANSI colors */
#define CLR_GREEN "\x1b[32m"
#define CLR_RED   "\x1b[31m"
#define CLR_YEL   "\x1b[33m"
#define CLR_RST   "\x1b[0m"

/* test counters */
static int total_tests  = 0;
static int passed_tests = 0;

#define static_len(G) (sizeof(G) / sizeof((G)[ 0 ]))

/* helper: load program at pc and set PC (m[7]) */
void load_prog(int pc, int *prog, int len) {
    int i;
    a = b = 0;
    u = 1;
    /* clear memory (keeps tests isolated) */
    for (i = 0; i < static_len(m); ++i) m[ i ] = 0;
    for (i = 0; i < len; ++i) m[ pc + i ] = prog[ i ];
    m[ 7 ] = pc;
}

/* unified expect function (increments counters and prints colored output) */
void expect_int(const char *name, int got, int want) {
    ++total_tests;
    if (got == want) {
        ++passed_tests;
        printf("%sPASS%s: %s == %d\n", CLR_GREEN, CLR_RST, name, want);
    } else {
        printf("%sFAIL%s: %s == %d (want %d)\n", CLR_RED, CLR_RST, name, got, want);
    }
}

int outbuf[ 512 ];

void zero_Y(void) {
    int ii;
    for (ii = 0; ii < (1 << 10); ++ii) Y[ ii ] = 0;
}

void set_outbuf(void) {
    memset(outbuf, 0, sizeof outbuf);
    O = outbuf;
    zero_Y();
}

/* small helper to run a CPU test with a program loaded at 1024 */
void run_cpu_test(int *prog, int len, void (*setup)(void), const char *name, int (*check)(void)) {
    load_prog(1024, prog, len);
    if (setup) setup();
    pu2();
    if (check) {
        /* check should call expect_int internally or return 1 on success */
        check();
    }
    (void) name; /* kept for future extension */
}

/* example setups and checks for clarity */
int check_a_is_42(void) {
    expect_int("a (load imm 42)", a, 42);
    return 0;
}

int main(void) {
    /* CPU tests */
    {
        int prog[] = { 1, 42, 0, 0 }; /* [opcode, imm, halt, dummy] */
        load_prog(1024, prog, sizeof prog / sizeof prog[ 0 ]);
        a = b = v = c = 0;
        pu2();
        expect_int("a (load imm 42)", a, 42);
    }

    {
        int prog[] = { 2, 0, 0, 0 };
        load_prog(1024, prog, 4);
        /* place value at memory index 500 and set b=500 */
        m[ 500 ] = 99;
        b        = 500;
        a        = 0;
        pu2();
        expect_int("a (a = m[b])", a, 99);
    }

    {
        int prog[] = { 3, 0, 0, 0 };
        load_prog(1024, prog, 4);
        a = 600;
        b = 1234;
        pu2();
        expect_int("m[600] (m[a]=b)", m[ 600 ], 1234);
    }

    {
        int prog[] = { 4, 0, 0, 0 };
        load_prog(1024, prog, 4);
        a = 7;
        b = 11;
        v = 0;
        pu2();
        expect_int("a (after swap)", a, 11);
        expect_int("b (after swap)", b, 7);
    }

    {
        int prog[] = { 5, 0, 0, 0 };
        load_prog(1024, prog, 4);
        a = 5;
        b = 8;
        pu2();
        expect_int("a (a += b)", a, 13);
    }

    {
        int prog[] = { 6, 0, 0, 0 };
        load_prog(1024, prog, 4);
        a = 20;
        b = 3;
        pu2();
        expect_int("a (a -= b)", a, 17);
    }

    {
        int prog[] = { 7, 0, 0, 0 };
        load_prog(1024, prog, 4);
        a = 5;
        b = 2;
        pu2();
        expect_int("a (5 && 2)", a, 1);
        load_prog(1024, prog, 4);
        a = 0;
        b = 3;
        pu2();
        expect_int("a (0 && 3)", a, 0);
    }

    {
        int prog[] = { 8, 0, 0, 0 };
        load_prog(1024, prog, 4);
        a = 1;
        b = 1;
        pu2();
        expect_int("a (!a)", a, 0);
        load_prog(1024, prog, 4);
        a = 0;
        b = 1;
        pu2();
        expect_int("a (!a)", a, 1);
    }

    // .. TODO: implement test case for ACC instruction ..

    {
        int prog[] = { 1, 9, 4, 0, 3, 0, 0, 0 };
        /* sequence: load imm 9 into a, swap, store b into m[a], halt */
        load_prog(1024, prog, sizeof prog / sizeof prog[ 0 ]);
        a = 2;
        b = 87; /* after load & swap: a=87, b=9 ; then m[87]=9 */
        pu2();
        expect_int("m[87] (from multi-op)", m[ 87 ], 9);
    }

    /* Assembler (as2) tests: use outbuf */
    {
        int *o;

        set_outbuf();
        as2("hlt");
        o = outbuf;
        expect_int("hlt:opcode", o[ 0 ], 0);
        expect_int("hlt:imm", o[ 1 ], 0);

        set_outbuf();
        as2("imm 42");
        o = outbuf;
        expect_int("imm:opcode", o[ 0 ], 1);
        expect_int("imm:imm", o[ 1 ], 42);

        set_outbuf();
        as2("add sub");
        o = outbuf;
        expect_int("add:opcode", o[ 0 ], 5);
        expect_int("add:imm", o[ 1 ], 0);
        expect_int("sub:opcode", o[ 2 ], 6);
        expect_int("sub:imm", o[ 3 ], 0);

        set_outbuf();
        as2("{0 hlt} @0");
        o = outbuf;
        expect_int("macro-hlt:opcode", o[ 0 ], 0);
        expect_int("macro-hlt:imm", o[ 1 ], 0);

        set_outbuf();
        as2("{0 imm 7} @0");
        o = outbuf;
        expect_int("macro-imm:opcode", o[ 0 ], 1);
        expect_int("macro-imm:imm", o[ 1 ], 7);

        set_outbuf();
        as2("{0 imm 3 add} @0");
        o = outbuf;
        expect_int("macro-mix:imm-op", o[ 0 ], 1);
        expect_int("macro-mix:imm-val", o[ 1 ], 3);
        expect_int("macro-mix:add-op", o[ 2 ], 5);
        expect_int("macro-mix:add-imm", o[ 3 ], 0);
    }

    /* Summary */
    printf(
        "\n%sTest summary%s: %d total, %d passed, %d failed\n",
        CLR_YEL,
        CLR_RST,
        total_tests,
        passed_tests,
        total_tests - passed_tests);

    if (passed_tests == total_tests) {
        printf("%sALL TESTS PASSED%s\n", CLR_GREEN, CLR_RST);
        return 0;
    } else {
        printf("%sSOME TESTS FAILED%s\n", CLR_RED, CLR_RST);
        return 1;
    }
}