// tb_sha3_chain_xsim.v - Vivado xsim testbench for sha3_chain_top
//
// Tests the ML-KEM G function: SHA3-512(d_in || 0x02) → {rho, sigma}.
// Uses $readmemh file-based test vectors for d_in values.
//
// Vector format (256 bits = 64 hex chars per line):
//   Each line is a 256-bit d_in value (MSB-first hex).
//
// Expected rho/sigma values are hardcoded (computed by gen_vectors.py
// using hashlib.sha3_512).
//
// Protocol: start_i / done_o handshake
//   1. Drive d_in, assert start_i=1
//   2. Wait for done_o=1 (with timeout watchdog)
//   3. Capture rho_out, sigma_out
//   4. Compare with expected values
//   5. Deassert start_i (FSM returns to IDLE)

`timescale 1ns / 1ps

module tb_sha3_chain_xsim;

    // ================================================================
    // Parameters
    // ================================================================
    parameter VECTOR_FILE   = "sync_rtl/sha3_chain/TB/vectors/sha3_chain_input.hex";
    parameter RESULT_FILE   = "sync_rtl/sha3_chain/TB/vectors/sha3_chain_result.hex";
    parameter MAX_VECTORS   = 256;
    parameter TIMEOUT_CYCLES = 2000;

    // ================================================================
    // DUT signals
    // ================================================================
    reg              clk;
    reg              rst_n;
    reg  [255:0]     d_in;
    reg              start_i;
    wire             done_o;
    wire [255:0]     rho_out;
    wire [255:0]     sigma_out;

    // ================================================================
    // Expected values (computed by gen_vectors.py)
    // 8 test vectors total
    // ================================================================
    reg [255:0] expected_rho   [0:7];
    reg [255:0] expected_sigma [0:7];

    initial begin
        // all-zeros
        expected_rho[0]   = 256'h6a0af64a85e909df8e2816605d20b4e382b30bbb61bf3a5f821a0b5dba9ad3e7;
        expected_sigma[0] = 256'he367d3e9ab3b86b64230aa8bf8815d408ef819f9e9d956ce22bbc2f68eaa9e3a;
        // all-ones
        expected_rho[1]   = 256'hf393670510fe33b5df9efb515bf1515c84bdf625e3c53154c10c7d2c92b7f234;
        expected_sigma[1] = 256'h9fac853f1e61e8389a04fc710845963c18e48fbec11c66d7d6567f8e168160cc;
        // lsb-one
        expected_rho[2]   = 256'h90ad3aafcc6bde61dc5014cf6cdee8065504733fc0caa8bddf3e1689a7b7b302;
        expected_sigma[2] = 256'h3c8686e00e96619b07c7d56fa1effce7ec5597f94a9109ec40c2a6314f6e4ada;
        // msb-one
        expected_rho[3]   = 256'h6c914803c353f4f3b4d8cc541b67019dd2d04cd0fb2ba51b8ebaba973d8f5cbe;
        expected_sigma[3] = 256'h33d64b904e1afa541dbe72162021b54c7aef182125da38d88dd7d076636af6cf;
        // pattern-aa
        expected_rho[4]   = 256'h561a98ff76434dc201ca8152cf2b97342090157a522354e299db35fec29690aa;
        expected_sigma[4] = 256'hed57f5826ed6e77c756cf821244b350e6b34eb05f9639e227f9a5e130eca649f;
        // pattern-55
        expected_rho[5]   = 256'hff31fe4fc62849226d9197992b1a2d0429c2c325011d2c7bb48860920c04636c;
        expected_sigma[5] = 256'hdb5f9599ae43d39fa182896d0ab1d1af663f19355dfeb951fed7b1c8c94eb0d3;
        // random-1
        expected_rho[6]   = 256'h4538916fc357ba9fe24033fa8c1e18cbfb14faf7e1771e0d7521c80cb0d4379f;
        expected_sigma[6] = 256'h78565d212000aabb8086e3373ad82d45c07589bcca35409a3f5d9fd30623340e;
        // random-2
        expected_rho[7]   = 256'h176c0b7e91b90ae50757b14bbec130a9e328d09d020466565829de47f6b387cd;
        expected_sigma[7] = 256'h413f41c838828e079d5e06b06de61eed6b5bd5005b334e516be6e3b900da7938;
    end

    // ================================================================
    // DUT instantiation
    // ================================================================
    sha3_chain_top u_dut (
        .clk      (clk),
        .rst_n    (rst_n),
        .d_in     (d_in),
        .start_i  (start_i),
        .done_o   (done_o),
        .rho_out  (rho_out),
        .sigma_out(sigma_out)
    );

    // ================================================================
    // Clock generation: 100 MHz (10 ns period)
    // ================================================================
    initial clk = 1'b0;
    always #5 clk = ~clk;

    // ================================================================
    // Vector memory (loaded by $readmemh)
    // ================================================================
    reg [255:0] vector_mem [0:MAX_VECTORS-1];
    integer     vec_count;
    integer     idx;
    integer     cycle_count;
    integer     result_fd;
    integer     pass_count;
    integer     fail_count;

    // ================================================================
    // Hex-to-ASCII conversion helper
    // ================================================================
    function [7:0] nibble_to_ascii;
        input [3:0] nibble;
        begin
            if (nibble < 4'd10)
                nibble_to_ascii = 8'h30 + {4'd0, nibble};
            else
                nibble_to_ascii = 8'h41 + ({4'd0, nibble} - 4'd10);
        end
    endfunction

    // ================================================================
    // Main test sequence
    // ================================================================
    initial begin
        vec_count = 0;
        $readmemh(VECTOR_FILE, vector_mem);

        // Count non-X entries
        begin
            integer found_end;
            found_end = 0;
            for (idx = 0; idx < MAX_VECTORS; idx = idx + 1) begin
                if (!found_end && (vector_mem[idx] === 256'hx || vector_mem[idx] === 256'hz))
                    found_end = 1;
                else if (!found_end)
                    vec_count = vec_count + 1;
            end
        end

        if (vec_count == 0) begin
            $display("ERROR: No vectors loaded from %s", VECTOR_FILE);
            $display("       Check that the file exists.");
            $display("       Format: 64 hex chars per line (256-bit d_in)");
            $finish;
        end

        $display("INFO: Loaded %0d test vectors from %s", vec_count, VECTOR_FILE);

        // Open result file
        result_fd = $fopen(RESULT_FILE, "w");
        if (result_fd == 0) begin
            $display("ERROR: Cannot open result file: %s", RESULT_FILE);
            $finish;
        end

        // Initialize DUT inputs
        d_in    <= 256'd0;
        start_i <= 1'b0;

        // Reset sequence: rst_n low for 3 cycles, then high
        rst_n <= 1'b0;
        repeat (3) @(posedge clk);
        rst_n <= 1'b1;
        @(posedge clk);

        pass_count = 0;
        fail_count = 0;

        // ============================================================
        // Process each vector
        // ============================================================
        for (idx = 0; idx < vec_count; idx = idx + 1) begin
            begin
                reg [255:0] vec_d_in;
                reg [255:0] captured_rho;
                reg [255:0] captured_sigma;
                integer     bit_idx;
                reg [3:0]   nib;

                vec_d_in = vector_mem[idx];

                $display("INFO: Vector %0d - d_in=0x%064h", idx, vec_d_in);

                // Drive d_in and assert start_i
                d_in    <= vec_d_in;
                start_i <= 1'b1;
                @(posedge clk);

                // Wait for done_o (with timeout)
                cycle_count = 0;
                while (!done_o && cycle_count < TIMEOUT_CYCLES) begin
                    @(posedge clk);
                    cycle_count = cycle_count + 1;
                end

                if (cycle_count >= TIMEOUT_CYCLES) begin
                    $display("ERROR: Timeout waiting for done_o on vector %0d", idx);
                    $fwrite(result_fd, "TIMEOUT: vector %0d d_in=0x%064h\n", idx, vec_d_in);
                    fail_count = fail_count + 1;
                    // Force reset the FSM
                    start_i <= 1'b0;
                    @(posedge clk);
                end else begin
                    // Capture outputs (valid on cycle done_o=1)
                    captured_rho   = rho_out;
                    captured_sigma = sigma_out;

                    // Deassert start_i (FSM returns to IDLE on next cycle)
                    start_i <= 1'b0;

                    // Check rho
                    if (captured_rho !== expected_rho[idx]) begin
                        $display("FAIL: Vector %0d RHO mismatch", idx);
                        $display("       expected = 0x%064h", expected_rho[idx]);
                        $display("       got      = 0x%064h", captured_rho);
                        $fwrite(result_fd, "FAIL: vector %0d RHO expected=0x%064h got=0x%064h\n",
                                idx, expected_rho[idx], captured_rho);
                        fail_count = fail_count + 1;
                    end else begin
                        pass_count = pass_count + 1;
                    end

                    // Check sigma
                    if (captured_sigma !== expected_sigma[idx]) begin
                        $display("FAIL: Vector %0d SIGMA mismatch", idx);
                        $display("       expected = 0x%064h", expected_sigma[idx]);
                        $display("       got      = 0x%064h", captured_sigma);
                        $fwrite(result_fd, "FAIL: vector %0d SIGMA expected=0x%064h got=0x%064h\n",
                                idx, expected_sigma[idx], captured_sigma);
                        fail_count = fail_count + 1;
                    end else begin
                        pass_count = pass_count + 1;
                    end

                    // Write results to output file
                    $fwrite(result_fd, "RESULT: %0d ", idx);
                    $fwrite(result_fd, "d_in=0x");
                    for (bit_idx = 63; bit_idx >= 0; bit_idx = bit_idx - 1) begin
                        nib = vec_d_in[(bit_idx*4)+:4];
                        $fwrite(result_fd, "%c", nibble_to_ascii(nib));
                    end
                    $fwrite(result_fd, " rho=0x");
                    for (bit_idx = 63; bit_idx >= 0; bit_idx = bit_idx - 1) begin
                        nib = captured_rho[(bit_idx*4)+:4];
                        $fwrite(result_fd, "%c", nibble_to_ascii(nib));
                    end
                    $fwrite(result_fd, " sigma=0x");
                    for (bit_idx = 63; bit_idx >= 0; bit_idx = bit_idx - 1) begin
                        nib = captured_sigma[(bit_idx*4)+:4];
                        $fwrite(result_fd, "%c", nibble_to_ascii(nib));
                    end
                    $fwrite(result_fd, "\n");

                    @(posedge clk);
                end
            end  // inner begin block
        end

        // ============================================================
        // Summary
        // ============================================================
        $fclose(result_fd);

        $display("========================================");
        $display("SHA3 CHAIN TEST COMPLETE");
        $display("  Total vectors:   %0d", vec_count);
        $display("  Checks (rho+sigma): %0d", vec_count * 2);
        $display("  Passed:          %0d", pass_count);
        $display("  Failed:          %0d", fail_count);
        $display("  Results written to: %s", RESULT_FILE);
        if (fail_count == 0)
            $display("  RESULT: ALL TESTS PASSED");
        else
            $display("  RESULT: SOME TESTS FAILED");
        $display("========================================");

        $finish;
    end

    // ================================================================
    // Timeout watchdog
    // ================================================================
    initial begin
        #(TIMEOUT_CYCLES * 10 * 100);
        $display("FATAL: Global simulation timeout reached");
        $finish;
    end

endmodule