// tb_sha3_mb_xsim.v - Multi-block SHA3-256 absorb self-check for sha3_top.
//
// Streams pre-padded 1088-bit rate blocks (mb_h_blocks.hex) into sha3_top
// with mb_en=1, then compares the 256-bit digest against mb_h_expected.hex
// (= hashlib.sha3_256(ek) for KAT count=0, ek = 800 bytes -> 6 blocks).
//
// Run (from project root):
//   xvlog -sv sync_rtl/sha3/keccak_round.v sync_rtl/sha3/keccak_core.v \
//             sync_rtl/sha3/sha3_top.v sync_rtl/sha3/TB/tb_sha3_mb_xsim.v
//   xelab tb_sha3_mb_xsim -s mb_sim --timescale 1ns/1ps
//   xsim mb_sim -R

`timescale 1ns/1ps
module tb_sha3_mb_xsim;

    parameter BLOCK_FILE    = "sync_rtl/sha3/TB/vectors/mb_h_blocks.hex";
    parameter EXPECTED_FILE = "sync_rtl/sha3/TB/vectors/mb_h_expected.hex";
    parameter MAX_BLOCKS    = 16;

    reg              clk = 0;
    reg              rst_n = 0;
    // single-block ports (unused here, tie off)
    reg  [1:0]       mode = 2'b01;
    reg  [511:0]     data_i = 0;
    reg              valid_i = 0;
    wire             ready_o;
    wire [511:0]     hash_o;
    wire             valid_o;
    reg              ready_i = 0;
    // multi-block ports
    reg              mb_en = 1'b1;
    reg  [1087:0]    mb_block_i = 0;
    reg              mb_valid_i = 0;
    reg              mb_last_i = 0;
    wire             mb_ready_o;

    sha3_top dut (
        .clk(clk), .rst_n(rst_n),
        .mode(mode), .data_i(data_i), .valid_i(valid_i),
        .ready_o(ready_o), .hash_o(hash_o), .valid_o(valid_o), .ready_i(ready_i),
        .mb_en(mb_en), .mb_block_i(mb_block_i),
        .mb_valid_i(mb_valid_i), .mb_last_i(mb_last_i), .mb_ready_o(mb_ready_o)
    );

    always #5 clk = ~clk;

    reg [1087:0] blocks [0:MAX_BLOCKS-1];
    reg [255:0]  expected_mem [0:0];
    reg [255:0]  expected;
    integer      nblocks;
    integer      i;

    // count valid (non-x) lines loaded
    function integer count_blocks;
        integer k;
        begin
            count_blocks = 0;
            for (k = 0; k < MAX_BLOCKS; k = k + 1)
                if (blocks[k] !== {1088{1'bx}}) count_blocks = k + 1;
        end
    endfunction

    initial begin
        for (i = 0; i < MAX_BLOCKS; i = i + 1) blocks[i] = {1088{1'bx}};
        $readmemh(BLOCK_FILE, blocks);
        $readmemh(EXPECTED_FILE, expected_mem);
        expected = expected_mem[0];
        nblocks = count_blocks();
        $display("=== Multi-block SHA3-256 absorb TB ===");
        $display("Loaded %0d blocks; expected digest = %064x", nblocks, expected);

        // reset
        rst_n = 0; ready_i = 0;
        repeat (4) @(posedge clk);
        rst_n = 1;
        @(posedge clk);

        // stream blocks — hold valid/last stable across the accept edge
        // (deassert only after mb_ready_o drops) to avoid sampling races.
        for (i = 0; i < nblocks; i = i + 1) begin
            mb_block_i = blocks[i];
            mb_last_i  = (i == nblocks - 1);
            mb_valid_i = 1'b1;
            // wait until module is ready, then one more edge is the accept
            while (!mb_ready_o) @(posedge clk);
            @(posedge clk);               // accept edge (valid & ready both high)
            while (mb_ready_o) @(posedge clk);  // ready drops => accepted, now safe
            mb_valid_i = 1'b0;
            mb_last_i  = 1'b0;
            // wait for this block's permutation to finish (ready again or DONE)
            while (!mb_ready_o && !valid_o) @(posedge clk);
            $display("  after block %0d: state[255:0] = %064x", i, dut.mb_state_r[255:0]);
        end

        // wait for digest
        i = 0;
        while (!valid_o && i < 1000) begin @(posedge clk); i = i + 1; end

        if (!valid_o) begin
            $display("FAIL: digest never became valid (timeout)");
            $finish;
        end

        if (hash_o[255:0] === expected) begin
            $display("PASS: H(ek) = %064x", hash_o[255:0]);
            $display("ALL TESTS PASSED");
        end else begin
            $display("FAIL: digest mismatch");
            $display("  got      = %064x", hash_o[255:0]);
            $display("  expected = %064x", expected);
        end
        ready_i = 1;          // ack the digest
        @(posedge clk);
        $finish;
    end

    // global safety timeout
    initial begin
        #100000;
        $display("FAIL: global timeout");
        $finish;
    end

endmodule