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

#define __WRAPPER
#include "pu.c"

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

void test_opcode_0() {
    printf("Testing opcode 0 (hlt)... ");
    reset();
    m[0] = 0;  // hlt instruction
    m[7] = 0;  // pc starts at 0
    pu();
    printf("PASS\n");
}

void test_opcode_1() {
    printf("Testing opcode 1 (const in r0)... ");
    reset();
    m[0] = 1;   // const instruction
    m[1] = 42;  // constant value
    m[2] = 0;   // hlt
    m[7] = 0;   // pc starts at 0
    pu();
    printf("r0=%d (expected 42): %s\n", m[0], m[0] == 42 ? "PASS" : "FAIL");
}

void test_opcode_2() {
    printf("Testing opcode 4 (r0 in r2 if acc non-zero)... ");
    reset();
    m[0] = 123;  // r0 = 123
    a = 1;       // acc = 1 (non-zero)
    m[10] = 4;   // opcode 4
    m[11] = 0;   // hlt
    m[7] = 10;   // pc starts at 10
    pu();
    printf("r2=%d (expected 123): %s\n", m[2], m[2] == 123 ? "PASS" : "FAIL");
}

void test_opcode_2_zero_acc() {
    printf("Testing opcode 4 with zero acc (should not store)... ");
    reset();
    m[0] = 123;  // r0 = 123
    m[2] = 999;  // r2 = 999 initially
    a = 0;       // acc = 0
    m[10] = 4;   // opcode 4
    m[11] = 0;   // hlt
    m[7] = 10;   // pc starts at 10
    pu();
    printf("r2=%d (expected 999): %s\n", m[2], m[2] == 999 ? "PASS" : "FAIL");
}

void test_opcode_10() {
    printf("Testing opcode 12 (r2 in r0)... ");
    reset();
    m[2] = 456;  // r2 = 456
    m[10] = 12;  // opcode 12
    m[11] = 0;   // hlt
    m[7] = 10;   // pc starts at 10
    pu();
    printf("r0=%d (expected 456): %s\n", m[0], m[0] == 456 ? "PASS" : "FAIL");
}

void test_opcode_18() {
    printf("Testing opcode 18 (r0 in m[r1])... ");
    reset();
    m[0] = 789;   // r0 = 789
    m[1] = 100;   // r1 = 100 (address)
    m[10] = 18;   // opcode 18
    m[11] = 0;    // hlt
    m[7] = 10;    // pc starts at 10
    pu();
    printf("m[100]=%d (expected 789): %s\n", m[100], m[100] == 789 ? "PASS" : "FAIL");
}

void test_opcode_19() {
    printf("Testing opcode 19 (acc in r0)... ");
    reset();
    a = 555;      // acc = 555
    m[10] = 19;   // opcode 19
    m[11] = 0;    // hlt
    m[7] = 10;    // pc starts at 10
    pu();
    printf("r0=%d (expected 555): %s\n", m[0], m[0] == 555 ? "PASS" : "FAIL");
}

void test_opcode_20() {
    printf("Testing opcode 20 (acc += r0)... ");
    reset();
    a = 10;       // acc = 10
    m[0] = 5;     // r0 = 5
    m[10] = 20;   // opcode 20
    m[11] = 0;    // hlt
    m[7] = 10;    // pc starts at 10
    pu();
    printf("acc=%d (expected 15): %s\n", a, a == 15 ? "PASS" : "FAIL");
}

void test_opcode_21() {
    printf("Testing opcode 21 (acc -= r0)... ");
    reset();
    a = 20;       // acc = 20
    m[0] = 7;     // r0 = 7
    m[10] = 21;   // opcode 21
    m[11] = 0;    // hlt
    m[7] = 10;    // pc starts at 10
    pu();
    printf("acc=%d (expected 13): %s\n", a, a == 13 ? "PASS" : "FAIL");
}

void test_opcode_22() {
    printf("Testing opcode 22 (acc = !acc)... ");
    reset();
    a = 5;        // acc = 5 (non-zero)
    m[10] = 22;   // opcode 22
    m[11] = 0;    // hlt
    m[7] = 10;    // pc starts at 10
    pu();
    printf("acc=%d (expected 0): %s\n", a, a == 0 ? "PASS" : "FAIL");
}

void test_opcode_23() {
    printf("Testing opcode 23 (acc = 0)... ");
    reset();
    a = 999;      // acc = 999
    m[10] = 23;   // opcode 23
    m[11] = 0;    // hlt
    m[7] = 10;    // pc starts at 10
    pu();
    printf("acc=%d (expected 0): %s\n", a, a == 0 ? "PASS" : "FAIL");
}

// Helper to compare output array with expected values
int check_output(int *O, int *expected, int len, const char *test_name) {
    int i;
    for (i = 0; i < len; i++) {
        if (O[i] != expected[i]) {
            printf("%s: output[%d]=%d (expected %d): FAIL\n", test_name, i, O[i], expected[i]);
            return 0;
        }
    }
    printf("%s: PASS\n", test_name);
    return 1;
}

// Reset output array
void reset_output(int *O, int size) {
    memset(O, 0, size * sizeof(int));
}

void test_hlt() {
    printf("Testing hlt (opcode 0)... ");
    int O[10];
    reset_output(O, 10);
    int expected[] = {0};
    int ret = as("hlt", O);
    printf("return=%d (expected 0): %s\n", ret, ret == 0 ? check_output(O, expected, 1, "hlt") ? "PASS" : "FAIL" : "FAIL");
}

void test_con() {
    printf("Testing con 123 (opcode 1, 123)... ");
    int O[10];
    reset_output(O, 10);
    int expected[] = {1, 123};
    int ret = as("con 123", O);
    printf("return=%d (expected 0): %s\n", ret, ret == 0 ? check_output(O, expected, 2, "con") ? "PASS" : "FAIL" : "FAIL");
}

void test_ar1() {
    printf("Testing ar1 (opcode 2)... ");
    int O[10];
    reset_output(O, 10);
    int expected[] = {2};
    int ret = as("ar1", O);
    printf("return=%d (expected 0): %s\n", ret, ret == 0 ? check_output(O, expected, 1, "ar1") ? "PASS" : "FAIL" : "FAIL");
}

void test_ar8() {
    printf("Testing ar8 (opcode 9)... ");
    int O[10];
    reset_output(O, 10);
    int expected[] = {9};
    int ret = as("ar8", O);
    printf("return=%d (expected 0): %s\n", ret, ret == 0 ? check_output(O, expected, 1, "ar8") ? "PASS" : "FAIL" : "FAIL");
}

void test_r1a() {
    printf("Testing r1a (opcode 10)... ");
    int O[10];
    reset_output(O, 10);
    int expected[] = {10};
    int ret = as("r1a", O);
    printf("return=%d (expected 0): %s\n", ret, ret == 0 ? check_output(O, expected, 1, "r1a") ? "PASS" : "FAIL" : "FAIL");
}

void test_r8a() {
    printf("Testing r8a (opcode 17)... ");
    int O[10];
    reset_output(O, 10);
    int expected[] = {17};
    int ret = as("r8a", O);
    printf("return=%d (expected 0): %s\n", ret, ret == 0 ? check_output(O, expected, 1, "r8a") ? "PASS" : "FAIL" : "FAIL");
}

void test_mem() {
    printf("Testing mem (opcode 18)... ");
    int O[10];
    reset_output(O, 10);
    int expected[] = {18};
    int ret = as("mem", O);
    printf("return=%d (expected 0): %s\n", ret, ret == 0 ? check_output(O, expected, 1, "mem") ? "PASS" : "FAIL" : "FAIL");
}

void test_aca() {
    printf("Testing aca (opcode 19)... ");
    int O[10];
    reset_output(O, 10);
    int expected[] = {19};
    int ret = as("aca", O);
    printf("return=%d (expected 0): %s\n", ret, ret == 0 ? check_output(O, expected, 1, "aca") ? "PASS" : "FAIL" : "FAIL");
}

void test_add() {
    printf("Testing add (opcode 20)... ");
    int O[10];
    reset_output(O, 10);
    int expected[] = {20};
    int ret = as("add", O);
    printf("return=%d (expected 0): %s\n", ret, ret == 0 ? check_output(O, expected, 1, "add") ? "PASS" : "FAIL" : "FAIL");
}

void test_sub() {
    printf("Testing sub (opcode 21)... ");
    int O[10];
    reset_output(O, 10);
    int expected[] = {21};
    int ret = as("sub", O);
    printf("return=%d (expected 0): %s\n", ret, ret == 0 ? check_output(O, expected, 1, "sub") ? "PASS" : "FAIL" : "FAIL");
}

void test_not() {
    printf("Testing not (opcode 22)... ");
    int O[10];
    reset_output(O, 10);
    int expected[] = {22};
    int ret = as("not", O);
    printf("return=%d (expected 0): %s\n", ret, ret == 0 ? check_output(O, expected, 1, "not") ? "PASS" : "FAIL" : "FAIL");
}

void test_clr() {
    printf("Testing clr (opcode 23)... ");
    int O[10];
    reset_output(O, 10);
    int expected[] = {23};
    int ret = as("clr", O);
    printf("return=%d (expected 0): %s\n", ret, ret == 0 ? check_output(O, expected, 1, "clr") ? "PASS" : "FAIL" : "FAIL");
}

void test_multi_instruction() {
    printf("Testing multi-instruction (hlt con 123 ar1 clr)... ");
    int O[10];
    reset_output(O, 10);
    int expected[] = {0, 1, 123, 2, 23};
    int ret = as("hlt con 123 ar1 clr", O);
    printf("return=%d (expected 0): %s\n", ret, ret == 0 ? check_output(O, expected, 5, "multi-instruction") ? "PASS" : "FAIL" : "FAIL");
}

void test_whitespace() {
    printf("Testing whitespace handling (hlt  \\t\\n con 456 \\n clr)... ");
    int O[10];
    reset_output(O, 10);
    int expected[] = {0, 1, 456, 23};
    int ret = as("hlt  \t\n con 456 \n clr", O);
    printf("return=%d (expected 0): %s\n", ret, ret == 0 ? check_output(O, expected, 4, "whitespace") ? "PASS" : "FAIL" : "FAIL");
}

void test_error() {
    printf("Testing error case (invalid instruction)... ");
    int O[10];
    reset_output(O, 10);
    int ret = as("xyz", O);
    printf("return=%d (expected 1): %s\n", ret, ret == 1 ? "PASS" : "FAIL");
}

int main() {
    printf("Running CPU unit tests...\n\n");
    
    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("\n\nRunning assembler unit tests...\n\n");
    
    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_error();
    
    
    return 0;
}