// tb_sha3.cpp - Verilator C++ testbench for sha3_top
//
// Reads test vectors from +VECTOR_FILE=plusarg.
// Format: "MM DDDD..." (mode hex, then 512-bit message hex)
//   MM: "00"=G, "01"=H, "10"=J
//   DDDD...: 128 hex chars representing 512-bit data_i
//
// Drives DUT with mode and input, waits for ready_o/valid_o,
// prints "RESULT: OUTPUT_HEX\n" to stdout.
//
// Clock: 10ns period. Reset: 2 cycles.

#include <iostream>
#include <fstream>
#include <string>
#include <sstream>
#include <cstdlib>
#include <cstring>
#include <cstdint>
#include "Vsha3_top.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(Vsha3_top* dut) {
    dut->clk = !dut->clk;
    main_time += (vluint64_t)(CLK_PERIOD_NS / 2.0);
    dut->eval();
    dut->clk = !dut->clk;
    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 a hex string into Verilator WData words (16 x 32-bit = 512 bits).
// Hex string is MSB-first (leftmost hex char = bits 511:508).
// data[0] = bits[31:0], data[15] = bits[511:480].
static void hex_to_512(const std::string& hex, uint32_t data_words[16]) {
    for (int w = 0; w < 16; w++) data_words[w] = 0;

    int len = (int)hex.length();
    // nibble_idx increments for each hex char from RIGHT to LEFT
    int nibble_idx = 0;
    for (int i = len - 1; i >= 0; i--) {
        char c = hex[i];
        if (c == ' ' || c == '\t') continue;
        int nib = hex_char_to_nibble(c);
        int word_idx = nibble_idx / 8;       // 8 nibbles per 32-bit word
        int shift    = (nibble_idx % 8) * 4; // shift within the word
        if (word_idx < 16) {
            data_words[word_idx] |= ((uint32_t)nib << shift);
        }
        nibble_idx++;
    }
}

// Print hash_o as hex (MSB-first).
// data[0] = bits[31:0], data[15] = bits[511:480].
static void print_hex(uint32_t data_words[16], int bits) {
    int full_words = bits / 32;
    // Build hex from MSB nibble to LSB nibble
    for (int w = full_words - 1; w >= 0; w--) {
        uint32_t val = data_words[w];
        for (int j = 28; j >= 0; j -= 4) {
            int nib = (int)((val >> j) & 0xF);
            printf("%01X", nib);
        }
    }
    printf("\n");
}

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
    Vsha3_top* dut = new Vsha3_top;

    // Initialize
    dut->clk     = 0;
    dut->rst_n   = 0;
    dut->mode    = 0;
    for (int w = 0; w < 16; w++) dut->data_i[w] = 0;
    dut->valid_i = 0;
    dut->ready_i = 0;
    // multi-block absorb path disabled for single-block G/H/J tests
    dut->mb_en     = 0;
    dut->mb_valid_i = 0;
    dut->mb_last_i  = 0;

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

    // Consumer always ready
    dut->ready_i = 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_HEX MESSAGE_HEX
        std::istringstream iss(line);
        std::string mode_str, data_hex;
        if (!(iss >> mode_str >> data_hex)) continue;
        if (mode_str.length() < 2) continue;

        int mode_val = hex_char_to_nibble(mode_str[1]);

        // Set mode
        dut->mode = mode_val & 0x3;

        // Set data_i from hex string
        uint32_t data_words[16];
        hex_to_512(data_hex, data_words);
        for (int w = 0; w < 16; w++) dut->data_i[w] = data_words[w];

        // Assert valid_i for one cycle
        dut->valid_i = 1;

        // posedge: DUT samples valid_i, starts keccak_core
        posedge(dut);
        cycle++;
        dut->valid_i = 0;

        // Wait for valid_o (keccak_core takes ~25 cycles)
        do {
            posedge(dut);
            cycle++;
            if (cycle > TIMEOUT_CYCLES) {
                std::cerr << "ERROR: Timeout waiting for valid_o (vec " << vec_count << ")" << std::endl;
                goto done;
            }
        } while (!dut->valid_o);

        // Read hash_o and print
        uint32_t hash_words[16];
        for (int w = 0; w < 16; w++) hash_words[w] = dut->hash_o[w];

        int out_bits = (mode_val == 0) ? 512 : 256;
        printf("RESULT: ");
        print_hex(hash_words, out_bits);

        // One more cycle for valid_o handshake
        posedge(dut);
        cycle++;
        vec_count++;
    }

done:
    infile.close();
    delete dut;

    if (vec_count == 0) {
        std::cerr << "ERROR: No vectors processed" << std::endl;
        return 1;
    }

    return 0;
}
