// tb_ntt.cpp - Verilator C++ testbench for ntt_sync
//
// Reads test vectors from +VECTOR_FILE= plusarg.
// Format: "MODE COEFF0 COEFF1 ... COEFF255"
//   MODE: "F" = forward NTT, "I" = inverse NTT
//   COEFFx: 3-digit hex (12-bit, 000..CFF)
//
// Drives DUT with coefficients one per cycle, waits for output,
// prints "RESULT: COEFF0 COEFF1 ... COEFF255\n" to stdout.
//
// Clock: 10ns period. Reset: 2 cycles.

#include <iostream>
#include <fstream>
#include <string>
#include <sstream>
#include <vector>
#include <cstdlib>
#include <cstring>
#include <cstdint>
#include "Vntt_sync.h"
#include "verilated.h"

#define CLK_PERIOD_NS     10.0
#define TIMEOUT_CYCLES    500000

static vluint64_t main_time = 0;

double sc_time_stamp() {
    return main_time;
}

// Toggle clock: both edges + eval (one full cycle)
static void posedge(Vntt_sync* dut) {
    dut->clk = 1;
    main_time += (vluint64_t)(CLK_PERIOD_NS / 2.0);
    dut->eval();
    dut->clk = 0;
    main_time += (vluint64_t)(CLK_PERIOD_NS / 2.0);
    dut->eval();
}

static int hex_char_to_nibble(char c) {
    if (c >= '0' && c <= '9') return c - '0';
    if (c >= 'A' && c <= 'F') return c - 'A' + 10;
    if (c >= 'a' && c <= 'f') return c - 'a' + 10;
    return 0;
}

// Parse 3-char hex token to 12-bit value.
static uint16_t hex3_to_val(const std::string& tok) {
    uint16_t val = 0;
    for (size_t i = 0; i < tok.length() && i < 3; i++) {
        val = (val << 4) | hex_char_to_nibble(tok[i]);
    }
    return val & 0xFFF;
}

// Format 12-bit value as 3-char hex (lowercase for consistency).
static std::string val_to_hex3(uint16_t val) {
    char buf[4];
    snprintf(buf, sizeof(buf), "%03X", val & 0xFFF);
    return std::string(buf);
}

int main(int argc, char** argv) {
    Verilated::commandArgs(argc, argv);

    // Parse +VECTOR_FILE= plusarg
    const char* vector_file = NULL;
    for (int i = 1; i < argc; i++) {
        std::string arg(argv[i]);
        if (arg.rfind("+VECTOR_FILE=", 0) == 0) {
            vector_file = argv[i] + 13;
        }
    }

    if (!vector_file) {
        std::cerr << "ERROR: +VECTOR_FILE= not specified" << std::endl;
        return 1;
    }

    std::ifstream infile(vector_file);
    if (!infile.is_open()) {
        std::cerr << "ERROR: Cannot open vector file: " << vector_file << std::endl;
        return 1;
    }

    // Instantiate DUT
    Vntt_sync* dut = new Vntt_sync;

    // Initialize
    dut->clk     = 0;
    dut->rst_n   = 0;
    dut->mode    = 0;
    dut->coeff_in = 0;
    dut->valid_i = 0;
    dut->ready_i = 0;

    // Reset: 2 full cycles
    for (int i = 0; i < 2; i++) posedge(dut);
    dut->rst_n = 1;

    std::string line;
    vluint64_t cycle = 0;
    int vec_count = 0;

    while (std::getline(infile, line)) {
        if (line.empty() || line[0] == '#') continue;

        // Parse: MODE COEFF0 COEFF1 ... COEFF255
        std::istringstream iss(line);
        std::string mode_str;
        if (!(iss >> mode_str)) continue;

        int mode_val = 0;
        if (mode_str == "I") mode_val = 1;

        // Parse 256 coefficients into vector
        std::vector<uint16_t> input_coeffs(256);
        std::string tok;
        int coeff_idx = 0;
        while (iss >> tok && coeff_idx < 256) {
            input_coeffs[coeff_idx] = hex3_to_val(tok);
            coeff_idx++;
        }
        if (coeff_idx != 256) {
            std::cerr << "ERROR: Expected 256 coefficients, got " << coeff_idx
                      << " (vec " << vec_count << ")" << std::endl;
            continue;
        }

        // Set mode
        dut->mode = mode_val;

        // ---- Load 256 coefficients ----
        while (!dut->ready_o) {
            posedge(dut); cycle++;
            if (cycle > TIMEOUT_CYCLES) goto timeout_err;
        }
        for (int i = 0; i < 256; i++) {
            dut->coeff_in = input_coeffs[i];
            dut->valid_i = 1;
            posedge(dut); cycle++;
            dut->valid_i = 0;
            if (cycle > TIMEOUT_CYCLES) goto timeout_err;
        }

        // At this point, the DUT has captured all 256 coeffs and
        // transitioned to S_COMPUTE_RD (ready_o went low).

        // ---- Wait for valid_o (DUT computing) ----
        dut->ready_i = 1;
        while (!dut->valid_o) {
            posedge(dut);
            cycle++;
            if (cycle > TIMEOUT_CYCLES) goto timeout_err;
        }

        // ---- Read 256 output coefficients ----
        printf("RESULT: ");
        for (int i = 0; i < 256; i++) {
            // Wait for valid_o to be asserted (data is valid NOW)
            while (!dut->valid_o) {
                posedge(dut);
                cycle++;
                if (cycle > TIMEOUT_CYCLES) goto timeout_err;
            }

            // Capture coefficient BEFORE consuming posedge
            uint16_t coeff_val = (uint16_t)(dut->coeff_out & 0xFFF);
            printf("%s%s", val_to_hex3(coeff_val).c_str(),
                   (i < 255) ? " " : "");

            // Consume this coefficient: posedge with ready_i=1
            posedge(dut);
            cycle++;
        }
        printf("\n");

        vec_count++;
    }

    std::cout << "Processed " << vec_count << " vectors" << std::endl;
    infile.close();
    delete dut;
    return (vec_count > 0) ? 0 : 1;

timeout_err:
    std::cerr << "ERROR: Timeout at cycle " << cycle
              << " (vec " << vec_count << ")" << std::endl;
    infile.close();
    delete dut;
    return 1;
}