#include "Valu.h"
#include "verilated.h"
#include "tester.hpp"

#include <stdint.h>
#include <iostream>
#include <random>

#ifdef TRACE
#include "verilated_vcd_c.h"
#endif

struct state {
    VerilatedContext *ctx;
    Valu *valu;
    #ifdef TRACE
    VerilatedVcdC *trace;
    #endif
};

struct alu_testcase {
    state *state;

    std::string name;
    // Inputs
    uint32_t A, B;
    uint8_t op;

    // Outputs
    uint64_t O; // {O_hi, O_lo}
    std::optional<bool> overflow, zero;

    std::optional<unsigned int> max_cycles;
};

std::string fmt_hex(uint32_t n) {
    char hex[100];
    if (n < 0x100)
        snprintf(hex, sizeof hex, "0x%02x", n);
    else if (n < 0x10000)
        snprintf(hex, sizeof hex, "0x%04x", n);
    else
        snprintf(hex, sizeof hex, "0x%08x", n);

    return hex;
}

void posedge(state *state) {
    #ifdef TRACE
        state->ctx->timeInc(1);
        state->valu->CLK = 1;
        state->valu->eval();
        state->trace->dump(state->ctx->time());

        state->ctx->timeInc(1);
        state->valu->CLK = 0;
        state->valu->eval();
        state->trace->dump(state->ctx->time());
    #else
        state->ctx->timeInc(1);
        state->valu->CLK = 1;
        state->valu->eval();

        state->ctx->timeInc(1);
        state->valu->CLK = 0;
        state->valu->eval();
    #endif
}

void test_op(Tester *tester, alu_testcase test) {
    Tester subtester(tester, test.name);

    posedge(test.state);

    // assign inputs
    test.state->valu->op = test.op;
    test.state->valu->A = test.A;
    test.state->valu->B = test.B;
    test.state->valu->EN = 1;
    posedge(test.state);
    test.state->valu->EN = 0;

    int max_cycles = test.max_cycles.has_value() ? *test.max_cycles : 10000;

    int n_cycles = 1;
    for (; !test.state->valu->RDY && n_cycles < 10 + max_cycles * 2; n_cycles++) {
        posedge(test.state);
    }

    char rdy_after_s[100];
    snprintf(rdy_after_s, sizeof rdy_after_s, "RDY = 1 (after %d cycle(s))", n_cycles);
    subtester.assert_eq(rdy_after_s, test.state->valu->RDY, 1);

    if (test.max_cycles.has_value()) {
        if (n_cycles <= test.max_cycles) {
            char n_cycles_s[100];
            snprintf(n_cycles_s, sizeof n_cycles_s, "Finished within %d cycle(s) (was: %d)", max_cycles, n_cycles);
            subtester.assert_eq(n_cycles_s, n_cycles, n_cycles);
        } else {
            subtester.assert_eq("Finished within correct number of cycles", n_cycles, max_cycles);
        }
    }

    std::string o_name("O == ");
    o_name.append(fmt_hex(test.O));

    uint64_t O = (((uint64_t) test.state->valu->O_hi) << 32) | (uint64_t) test.state->valu->O_lo;
    subtester.assert_eq(o_name, O, test.O);

    if (test.overflow.has_value()) {
        if (*test.overflow)
            subtester.assert_eq("overflow flag set", test.state->valu->Fflow, 1);
        else
            subtester.assert_eq("no overflow flag", test.state->valu->Fflow, 0);
    }
    if (test.zero.has_value()) {
        if (*test.zero)
            subtester.assert_eq("zero flag set", test.state->valu->Fzero, 1);
        else
            subtester.assert_eq("no zero flag", test.state->valu->Fzero, 0);
    }
}

int main(int argc, char **argv) {
    bool DO_AUTO = false;

    VerilatedContext *vctx = new VerilatedContext;
    Verilated::traceEverOn(true);
    Valu *valu = new Valu(vctx);

    #ifdef TRACE
        if (argc != 2) {
            std::cout << "Run with argument for destination!" << std::endl;
            return 1;
        }

        VerilatedVcdC *trace = new VerilatedVcdC;
        valu->trace(trace, 99);
        trace->open(argv[1]);
        std::cout << "(writing trace to " << argv[1] << ")" << std::endl;

        state state = {
            .ctx = vctx,
            .valu = valu,
            .trace = trace,
        };
    #else
        state state = {
            .ctx = vctx,
            .valu = valu,
        };
    #endif

    Tester alu_t("alu");
    {
        Tester add_t(&alu_t, "add", true);

        test_op(&add_t, {
            .state = &state,
            .name = "0x2137+0x1234",
            .A = 0x2137, .B = 0x1234, .op = 0b000,
            .O = 0x336b, .overflow = false,
        });

        test_op(&add_t, {
            .state = &state,
            .name = "0x09+0x10",
            .A = 0x09, .B = 0x10, .op = 0b000,
            .O = 0x19, .overflow = false,
        });

        test_op(&add_t, {
            .state = &state,
            .name = "0x5555+0x5555",
            .A = 0x5555, .B = 0x5555, .op = 0b000,
            .O = 0xaaaa, .overflow = false,
        });

        test_op(&add_t, {
            .state = &state,
            .name = "0xfffffffe+0x1",
            .A = 0xfffffffe, .B = 0x1, .op = 0b000,
            .O = 0xffffffff, .overflow = false, .zero = false,
        });
        test_op(&add_t, {
            .state = &state,
            .name = "0xffffffff+0x1",
            .A = 0xffffffff, .B = 0x1, .op = 0b000,
            .O = 0x0, .overflow = true, .zero = true,
        });
        test_op(&add_t, {
            .state = &state,
            .name = "0xffffffff+0x2",
            .A = 0xffffffff, .B = 0x2, .op = 0b000,
            .O = 0x1, .overflow = true, .zero = false,
        });
        test_op(&add_t, {
            .state = &state,
            .name = "0x0+0x0",
            .A = 0x0, .B = 0x0, .op = 0b000,
            .O = 0x0, .overflow = false, .zero = true,
        });
    }

    {
        Tester sub_t(&alu_t, "sub", true);

        test_op(&sub_t, {
            .state = &state,
            .name = "0x2137-0x0420",
            .A = 0x2137, .B = 0x0420, .op = 0b001,
            .O = 0x1d17, .overflow = false,
        });

        test_op(&sub_t, {
            .state = &state,
            .name = "0x0-0x1",
            .A = 0x0, .B = 0x1, .op = 0b001,
            .O = 0xffffffff, .overflow = true,
        });

        test_op(&sub_t, {
            .state = &state,
            .name = "0x100-0x0200",
            .A = 0x100, .B = 0x200, .op = 0b001,
            .O = 0xffffff00, .overflow = true,
        });

        test_op(&sub_t, {
            .state = &state,
            .name = "0x21-0x9",
            .A = 0x21, .B = 0x9, .op = 0b001,
            .O = 0x18, .overflow = false, .zero = false,
        });

        test_op(&sub_t, {
            .state = &state,
            .name = "0x20-0x20",
            .A = 0x20, .B = 0x20, .op = 0b001,
            .O = 0x0, .overflow = false, .zero = true,
        });
    }

    {
        Tester bitwise_t(&alu_t, "bitwise", true);
        // 0x3 = 0b0011, 0x5 = 0b0101

        test_op(&bitwise_t, {
            .state = &state,
            .name = "0x3&0x5",
            .A = 0x3, .B = 0x5, .op = 0b100,
            .O = 0x1, .zero = false,
        });

        test_op(&bitwise_t, {
            .state = &state,
            .name = "0x3|0x5",
            .A = 0x3, .B = 0x5, .op = 0b101,
            .O = 0x7, .zero = false,
        });

        test_op(&bitwise_t, {
            .state = &state,
            .name = "0x3^0x5",
            .A = 0x3, .B = 0x5, .op = 0b110,
            .O = 0x6, .zero = false,
        });

        test_op(&bitwise_t, {
            .state = &state,
            .name = "~0xa5a5a5a5",
            .A = 0xa5a5a5a5, .B = 0x0, .op = 0b111,
            .O = 0x5a5a5a5a, .zero = false,
        });

        test_op(&bitwise_t, {
            .state = &state,
            .name = "~0xffffffff",
            .A = 0xffffffff, .B = 0x0, .op = 0b111,
            .O = 0x0, .zero = true,
        });
    }

    if (DO_AUTO) {
        Tester auto_t(&alu_t, "auto", true);

        std::default_random_engine eng;
        std::uniform_int_distribution<uint32_t> op_gen(0, 5);
        std::uniform_int_distribution<uint32_t> gen(0, 0xffffffff);

        for (int i = 0; i < 100; i++) {
            uint32_t A = gen(eng);
            uint32_t B = gen(eng);
            std::string name;

            switch (op_gen(eng)) {
            case 0: // Add
                name.append(fmt_hex(A));
                name.append("+");
                name.append(fmt_hex(B));

                test_op(&auto_t, {
                    .state = &state,
                    .name = name,
                    .A = A, .B = B, .op = 0b000,
                    .O = A + B, .overflow = (A + B < A), .zero = (A + B == 0),
                });

                break;
            case 1: // Subtract
                name.append(fmt_hex(A));
                name.append("-");
                name.append(fmt_hex(B));

                test_op(&auto_t, {
                    .state = &state,
                    .name = name,
                    .A = A, .B = B, .op = 0b001,
                    .O = A - B, .overflow = (B > A), .zero = (A == B),
                });

                break;
            case 2: // Bitwise AND
                name.append(fmt_hex(A));
                name.append("&");
                name.append(fmt_hex(B));

                test_op(&auto_t, {
                    .state = &state,
                    .name = name,
                    .A = A, .B = B, .op = 0b100,
                    .O = A & B, .overflow = 0, .zero = ((A & B) == 0),
                });

                break;
            case 3: // Bitwise OR
                name.append(fmt_hex(A));
                name.append("|");
                name.append(fmt_hex(B));

                test_op(&auto_t, {
                    .state = &state,
                    .name = name,
                    .A = A, .B = B, .op = 0b101,
                    .O = A | B, .overflow = 0, .zero = ((A | B) == 0),
                });

                break;
            case 4: // Bitwise XOR
                name.append(fmt_hex(A));
                name.append("^");
                name.append(fmt_hex(B));

                test_op(&auto_t, {
                    .state = &state,
                    .name = name,
                    .A = A, .B = B, .op = 0b110,
                    .O = A ^ B, .overflow = 0, .zero = ((A ^ B) == 0),
                });

                break;
            case 5: // Bitwise NOT
                name.append("~");
                name.append(fmt_hex(A));

                test_op(&auto_t, {
                    .state = &state,
                    .name = name,
                    .A = A, .B = B, .op = 0b111,
                    .O = ~A, .overflow = 0, .zero = (A == 0xffffffff),
                });

                break;
            }
        }
    }
    #ifdef TRACE
    state.trace->close();
    #endif
}