#include <verilated.h>
#include "VCore.h"
#include "memory.h"
#include <cstdio>
#include <cstdlib>

#define MAX_CYCLES 100000

int main(int argc, char** argv) {
    if (argc < 2) {
        fprintf(stderr, "Usage: %s <test.elf>\n", argv[0]);
        return 1;
    }

    Verilated::commandArgs(argc, argv);

    VCore* core = new VCore;
    Memory* mem = new Memory();

    if (!mem->load_elf(argv[1])) {
        fprintf(stderr, "Failed to load test binary\n");
        return 1;
    }

    // Reset
    core->reset = 1;
    core->clock = 0;
    core->eval();
    core->clock = 1;
    core->eval();
    core->clock = 0;
    core->reset = 0;
    core->eval();

    uint64_t cycle = 0;
    bool test_done = false;
    int exit_code = 0;

    while (cycle < MAX_CYCLES && !test_done) {
        // Handle instruction memory interface
        if (core->io_imem_req_valid) {
            uint64_t pc = core->io_imem_req_bits;
            core->io_imem_resp_valid = 1;
            core->io_imem_resp_bits_0 = mem->read32(pc);
            core->io_imem_resp_bits_1 = mem->read32(pc + 4);
        } else {
            core->io_imem_resp_valid = 0;
        }

        // Handle data memory interface
        if (core->io_dmem_req_valid) {
            uint64_t addr = core->io_dmem_req_bits_addr;

            // Check for tohost write
            if (core->io_dmem_req_bits_isStore && addr == TOHOST_ADDR) {
                uint64_t tohost = core->io_dmem_req_bits_data;
                if (tohost == 1) {
                    printf("[%lu] TEST PASSED\n", cycle);
                    test_done = true;
                    exit_code = 0;
                } else if (tohost & 1) {
                    printf("[%lu] TEST FAILED: error code %lu\n", cycle, tohost >> 1);
                    test_done = true;
                    exit_code = 1;
                }
            }

            if (core->io_dmem_req_bits_isStore) {
                switch (core->io_dmem_req_bits_size) {
                    case 0: mem->write8(addr, core->io_dmem_req_bits_data & 0xff); break;
                    case 1: mem->write16(addr, core->io_dmem_req_bits_data & 0xffff); break;
                    case 2: mem->write32(addr, core->io_dmem_req_bits_data & 0xffffffff); break;
                    default: mem->write64(addr, core->io_dmem_req_bits_data); break;
                }
                core->io_dmem_resp_valid = 0;
            } else {
                core->io_dmem_resp_bits = mem->read64(addr);
                core->io_dmem_resp_valid = 1;
            }
        } else {
            core->io_dmem_resp_valid = 0;
        }

        // Clock cycle
        core->clock = 0;
        core->eval();
        core->clock = 1;
        core->eval();

        cycle++;
    }

    if (!test_done) {
        printf("[%lu] TEST TIMEOUT\n", cycle);
        exit_code = 2;
    }

    delete core;
    delete mem;
    return exit_code;
}