// tb_sample_ntt_xsim.v - Vivado xsim smoke-test testbench for sample_ntt_sync
//
// Drives the DUT with test vectors from a hex file, then verifies that
// the DUT produces exactly 256 coefficients, each in range [0, Q-1].
// This is a smoke test — it does NOT compare against Python expected values
// because the RTL uses a per-permutation Keccak-p approach that differs
// from standard SHAKE-128 bit-stream semantics.
//
// Vector format (packed hex, 66 chars for $readmemh):
//   {j_idx[1:0], i_idx[1:0], k_i[2:0], rho_i[255:0]}   (264 bits → 66 hex chars)
//
// Usage:
//   xvlog -sv --relax sample_ntt_sync.v TB/tb_sample_ntt_xsim.v
//   xelab tb_sample_ntt_xsim -s tb_sample_ntt_xsim
//   xsim tb_sample_ntt_xsim -R
//
// Prerequisites:
//   - Generate vectors: python3 TB/gen_vectors.py
//   - Output file: vectors/sample_ntt_input.hex

`timescale 1ns / 1ps

module tb_sample_ntt_xsim;

    // ================================================================
    // Parameters
    // ================================================================
    parameter VECTOR_FILE  = "sync_rtl/sample_ntt/TB/vectors/sample_ntt_input.hex";
    parameter MAX_VECTORS  = 32;           // max lines in input file
    parameter Q            = 3329;         // ML-KEM modulus
    parameter N_COEFFS     = 256;          // coefficients per polynomial
    parameter TIMEOUT      = 500000;       // per-vector timeout (cycles)

    // ================================================================
    // DUT signals
    // ================================================================
    reg              clk;
    reg              rst_n;
    reg  [255:0]     rho_i;
    reg  [2:0]       k_i;
    reg  [1:0]       i_idx;
    reg  [1:0]       j_idx;
    reg              valid_i;
    wire             ready_o;
    wire [11:0]      coeff_o;
    wire             valid_o;
    reg              ready_i;
    wire             last_o;

    // ================================================================
    // DUT instantiation
    // ================================================================
    sample_ntt_sync u_dut (
        .clk      (clk),
        .rst_n    (rst_n),
        .rho_i    (rho_i),
        .k_i      (k_i),
        .i_idx    (i_idx),
        .j_idx    (j_idx),
        .valid_i  (valid_i),
        .ready_o  (ready_o),
        .coeff_o  (coeff_o),
        .valid_o  (valid_o),
        .ready_i  (ready_i),
        .last_o   (last_o)
    );

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

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

    // Test result tracking
    integer     pass_count;
    integer     fail_count;

    // ================================================================
    // Main test sequence
    // ================================================================
    initial begin
        vec_count = 0;

        // Load vectors from hex file
        $readmemh(VECTOR_FILE, vector_mem);

        // Count non-zero/X entries to determine actual vector count
        begin
            integer found_end;
            found_end = 0;
            for (idx = 0; idx < MAX_VECTORS; idx = idx + 1) begin
                if (!found_end && (vector_mem[idx] === {264{1'bx}} || vector_mem[idx] === {264{1'bz}}))
                    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 and is in the correct format.");
            $display("       Each line: <66 hex chars> = {j, i, k, rho}");
            $finish;
        end

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

        // Initialize DUT inputs
        rho_i   <= 256'd0;
        k_i     <= 3'd0;
        i_idx   <= 2'd0;
        j_idx   <= 2'd0;
        valid_i <= 1'b0;
        ready_i <= 1'b1;  // always ready to accept output

        // 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_rho;
                reg [2:0]   vec_k;
                reg [1:0]   vec_i;
                reg [1:0]   vec_j;
                reg          range_ok;

                // Extract fields from packed vector_mem
                // vector_mem[263:0] = {1'b0, j_idx, i_idx, k_i, rho_i}
                vec_rho = vector_mem[idx][255:0];
                vec_k   = vector_mem[idx][258:256];
                vec_i   = vector_mem[idx][260:259];
                vec_j   = vector_mem[idx][262:261];

                $display("INFO: Vector %0d - k=%0d, i=%0d, j=%0d, rho[0:31]=%0h...",
                         idx, vec_k, vec_i, vec_j, vec_rho[31:0]);

                // Wait for DUT to be IDLE (ready_o high) before driving
                while (!ready_o) begin
                    @(posedge clk);
                end

                // Drive DUT with input (1-cycle pulse on valid_i)
                rho_i   <= vec_rho;
                k_i     <= vec_k;
                i_idx   <= vec_i;
                j_idx   <= vec_j;
                valid_i <= 1'b1;
                @(posedge clk);
                valid_i <= 1'b0;

                // Wait for valid_o stream, collect all 256 coefficients
                coeff_idx = 0;
                cycle_count = 0;
                range_ok = 1'b1;  // assume OK until proven otherwise

                while (coeff_idx < N_COEFFS && cycle_count < TIMEOUT) begin
                    @(posedge clk);
                    cycle_count = cycle_count + 1;

                    if (valid_o) begin
                        // Check range: coefficient must be in [0, Q-1]
                        if (coeff_o >= Q) begin
                            $display("ERROR: Vector %0d coeff[%0d]=%0d out of range [0,%0d]",
                                     idx, coeff_idx, coeff_o, Q-1);
                            range_ok = 1'b0;
                        end
                        coeff_idx = coeff_idx + 1;
                    end
                end

                // Evaluate result
                if (cycle_count >= TIMEOUT) begin
                    $display("ERROR: Timeout on vector %0d (got %0d/%0d coefficients)",
                             idx, coeff_idx, N_COEFFS);
                    fail_count = fail_count + 1;
                end else if (coeff_idx != N_COEFFS) begin
                    $display("ERROR: Vector %0d incomplete (got %0d/%0d coefficients)",
                             idx, coeff_idx, N_COEFFS);
                    fail_count = fail_count + 1;
                end else if (!range_ok) begin
                    $display("ERROR: Vector %0d produced out-of-range coefficients",
                             idx);
                    fail_count = fail_count + 1;
                end else begin
                    $display("INFO: Vector %0d PASSED (%0d coefficients in %0d cycles, all in range)",
                             idx, coeff_idx, cycle_count);
                    pass_count = pass_count + 1;
                end

                // Wait for DUT to return to IDLE before next vector.
                // The DUT may still be processing its last Keccak permutation
                // (ST_WAIT → ST_DONE → ST_IDLE), which takes ~20+ cycles.
                // ASSERTION: after the 256th coeff (last_o), valid_o must stay
                // low until IDLE — any extra pulse is a spurious 257th output
                // (regression guard for the phase-1 need_more fix).
                begin
                    integer extra_pulses;
                    extra_pulses = 0;
                    while (!ready_o) begin
                        @(posedge clk);
                        if (valid_o) extra_pulses = extra_pulses + 1;
                    end
                    if (extra_pulses != 0) begin
                        $display("ERROR: Vector %0d emitted %0d spurious valid_o pulse(s) after last_o",
                                 idx, extra_pulses);
                        fail_count = fail_count + 1;
                        if (pass_count > 0) pass_count = pass_count - 1; // revoke the earlier PASS
                    end
                end
            end  // inner begin block
        end

        // ============================================================
        // Summary
        // ============================================================
        $display("========================================");
        $display("TEST COMPLETE");
        $display("  Total vectors: %0d", vec_count);
        $display("  Passed:        %0d", pass_count);
        $display("  Failed:        %0d", fail_count);
        $display("========================================");

        if (fail_count > 0) begin
            $display("ERROR: %0d vector(s) FAILED", fail_count);
            $finish;
        end else begin
            $display("ALL VECTORS PASSED");
            $finish;
        end
    end

    // ================================================================
    // Global simulation watchdog
    // ================================================================
    initial begin
        #(TIMEOUT * 10 * 200);   // generous global timeout
        $display("FATAL: Global simulation timeout reached");
        $finish;
    end

endmodule
