// tb_mlkem_kg_katK_xsim.v - ML-KEM KeyGen vs NIST KAT, parametric K (KP) + CASE.
// xelab -generic_top KP=2|3|4 ; xsim -testplusarg CASE=n
//   KP=2 -> k2 vectors (ML-KEM-512), KP=3 -> k3 (768), KP=4 -> k4 (1024).
//   ek = 384*KP+32 bytes (==KAT pk), dk = 768*KP+96 bytes (==KAT sk).
`timescale 1ns/1ps
module tb_mlkem_kg_katK_xsim;
    parameter KP = 2;
    localparam EKB = 384*KP + 32;
    localparam DKB = 768*KP + 96;
    localparam integer CLK_PERIOD_NS = 100;  // 10 MHz
    localparam integer CLK_HALF_NS   = CLK_PERIOD_NS / 2;

    reg clk=0, rst_n=0, start_i=0;
    reg [2:0] k_i;
    reg [255:0] d_i, z_i;
    wire busy_o, done_o;
    reg [3:0] dbg_slot_i=0; reg [7:0] dbg_idx_i=0; wire [11:0] dbg_coeff_o;
    reg dbg_byte_sel_i=0; reg [10:0] dbg_byte_idx_i=0; wire [7:0] dbg_byte_o;
    reg [11:0] dbg_dk_idx_i=0; wire [7:0] dbg_dk_o;
    wire [255:0] dbg_rho_o, dbg_sigma_o;

    // KMAX defaults to 4 (worst-case sizing); KP selects the runtime k value.
    mlkem_top dut (
        .clk(clk), .rst_n(rst_n), .k_i(k_i), .op_i(2'd0),
        .d_i(d_i), .z_i(z_i), .msg_i(256'd0), .start_i(start_i),
        .busy_o(busy_o), .done_o(done_o),
        .ek_in_we(1'b0), .ek_in_addr(11'd0), .ek_in_byte(8'd0),
        .dk_in_we(1'b0), .dk_in_addr(12'd0), .dk_in_byte(8'd0),
        .c_in_we(1'b0), .c_in_addr(11'd0), .c_in_byte(8'd0),
        .ss_o(), .dbg_ct_idx_i(11'd0), .dbg_ct_o(),
        .dbg_slot_i(dbg_slot_i), .dbg_idx_i(dbg_idx_i), .dbg_coeff_o(dbg_coeff_o),
        .dbg_byte_sel_i(dbg_byte_sel_i), .dbg_byte_idx_i(dbg_byte_idx_i), .dbg_byte_o(dbg_byte_o),
        .dbg_dk_idx_i(dbg_dk_idx_i), .dbg_dk_o(dbg_dk_o),
        .dbg_rho_o(dbg_rho_o), .dbg_sigma_o(dbg_sigma_o),
        .dbg_r_o(), .dbg_hek_o(),
        .dbg_mprime_o(), .dbg_kbar_o(), .dbg_decz_o(), .dbg_dech_o()
    );
    always #(CLK_HALF_NS) clk = ~clk;

    reg [255:0] dmem [0:0];
    reg [255:0] zmem [0:0];
    reg [7:0]   ek_gold [0:EKB-1];
    reg [7:0]   dk_gold [0:DKB-1];
    reg [7:0]   ek_got  [0:EKB-1];   // ek bytes read back from DUT
    reg [7:0]   dk_got  [0:DKB-1];   // dk bytes read back from DUT
    integer c, i, errors, casenum, j;
    integer state_cyc [0:31];
    integer si;
    reg [8*80-1:0] tag, dfile, zfile, ekfile, dkfile;

    function [8*16-1:0] state_name;
        input [4:0] s;
        begin
            case (s)
                5'd0:  state_name = "IDLE";
                5'd1:  state_name = "G";
                5'd2:  state_name = "A";
                5'd3:  state_name = "C";
                5'd4:  state_name = "N";
                5'd5:  state_name = "M";
                5'd6:  state_name = "E";
                5'd7:  state_name = "H";
                5'd8:  state_name = "ENC_LOAD";
                5'd9:  state_name = "ENC_H";
                5'd10: state_name = "ENC_G";
                5'd11: state_name = "ENC_A";
                5'd12: state_name = "ENC_TDEC";
                5'd13: state_name = "ENC_C";
                5'd14: state_name = "ENC_N";
                5'd15: state_name = "ENC_U";
                5'd16: state_name = "ENC_C1";
                5'd17: state_name = "ENC_V";
                5'd18: state_name = "ENC_C2";
                5'd19: state_name = "ENC_E2MV";
                5'd20: state_name = "DEC_LOAD";
                5'd21: state_name = "DEC_DECOMP";
                5'd22: state_name = "DEC_SDEC";
                5'd23: state_name = "DEC_NTT";
                5'd24: state_name = "DEC_W";
                5'd25: state_name = "DEC_MENC";
                5'd26: state_name = "DEC_G";
                5'd27: state_name = "DEC_J";
                5'd28: state_name = "DEC_CMP";
                5'd31: state_name = "DONE";
                default: state_name = "UNKNOWN";
            endcase
        end
    endfunction

    task reset_timing;
        begin
            for (si = 0; si < 32; si = si + 1) state_cyc[si] = 0;
        end
    endtask

    task print_timing;
        input [8*16-1:0] op_name;
        input integer cycles;
        integer runtime_ns;
        begin
            runtime_ns = cycles * CLK_PERIOD_NS;
            $display("TIME K=%0d CASE=%0d OP=%0s cycles=%0d clk=10MHz period=%0dns runtime=%0d ns (%0d.%03d us, %0d.%03d ms)",
                     KP, casenum, op_name, cycles, CLK_PERIOD_NS, runtime_ns,
                     runtime_ns/1000, runtime_ns%1000,
                     runtime_ns/1000000, (runtime_ns%1000000)/1000);
            $display("TIME_BREAKDOWN K=%0d CASE=%0d OP=%0s", KP, casenum, op_name);
            for (si = 0; si < 32; si = si + 1)
                if (state_cyc[si] != 0)
                    $display("  STATE %-12s cycles=%0d time_ns=%0d", state_name(si[4:0]), state_cyc[si], state_cyc[si]*CLK_PERIOD_NS);
        end
    endtask

    // Dump a byte array as offset-prefixed hex, 32 bytes/line.
    task dump_bytes(input [8*16-1:0] name, input integer n);
        integer a, b;
        begin
            for (a = 0; a < n; a = a + 32) begin
                $write("  %0s[%4d] ", name, a);
                for (b = a; b < a+32 && b < n; b = b + 1) begin
                    if (name == "ek") $write("%02x", ek_got[b]);
                    else               $write("%02x", dk_got[b]);
                end
                $write("\n");
            end
        end
    endtask

    initial begin
        if (!$value$plusargs("CASE=%d", casenum)) casenum = 0;
        $sformat(tag, "k%0d", KP);
        $sformat(dfile,  "sync_rtl/top/TB/vectors/kat_%0s_c%0d_d.hex",  tag, casenum);
        $sformat(zfile,  "sync_rtl/top/TB/vectors/kat_%0s_c%0d_z.hex",  tag, casenum);
        $sformat(ekfile, "sync_rtl/top/TB/vectors/kat_%0s_c%0d_ek.hex", tag, casenum);
        $sformat(dkfile, "sync_rtl/top/TB/vectors/kat_%0s_c%0d_dk.hex", tag, casenum);
        $readmemh(dfile,  dmem);
        $readmemh(zfile,  zmem);
        $readmemh(ekfile, ek_gold);
        $readmemh(dkfile, dk_gold);
        d_i = dmem[0]; z_i = zmem[0];
        k_i = KP[2:0];

        // ---- show the KeyGen inputs (d, z seeds: 32 bytes each, MSB-first) ----
        $display("=== ML-KEM K=%0d KAT case %0d INPUTS ===", KP, casenum);
        $write("  d = "); for (j = 0; j < 32; j = j + 1) $write("%02x", d_i[8*(31-j) +: 8]); $write("\n");
        $write("  z = "); for (j = 0; j < 32; j = j + 1) $write("%02x", z_i[8*(31-j) +: 8]); $write("\n");

        rst_n=0; repeat(4) @(posedge clk); rst_n=1; @(posedge clk);
        start_i=1; @(posedge clk); start_i=0;
        #1;
        reset_timing;
        c=0; while(!done_o && c<2000000) begin state_cyc[dut.st] = state_cyc[dut.st] + 1; @(posedge clk); c=c+1; end
        if(!done_o) begin $display("FAIL K=%0d case %0d: timeout", KP, casenum); $finish; end
        $display("=== ML-KEM K=%0d KAT case %0d: KeyGen done in %0d cyc ===", KP, casenum, c);
        print_timing("KeyGen", c);

        errors = 0;
        dbg_byte_sel_i = 1'b0;
        for (i = 0; i < EKB; i = i + 1) begin
            dbg_byte_idx_i = i[10:0]; @(posedge clk); @(posedge clk);
            ek_got[i] = dbg_byte_o;
            if (dbg_byte_o !== ek_gold[i]) begin
                if (errors < 8) $display("  EK[%0d] got=%02x exp=%02x", i, dbg_byte_o, ek_gold[i]);
                errors = errors + 1;
            end
        end
        for (i = 0; i < DKB; i = i + 1) begin
            dbg_dk_idx_i = i[11:0]; @(posedge clk); @(posedge clk);
            dk_got[i] = dbg_dk_o;
            if (dbg_dk_o !== dk_gold[i]) begin
                if (errors < 8) $display("  DK[%0d] got=%02x exp=%02x", i, dbg_dk_o, dk_gold[i]);
                errors = errors + 1;
            end
        end

        // ---- show the KeyGen outputs (ek, dk byte strings read from DUT) ----
        $display("=== ML-KEM K=%0d KAT case %0d OUTPUTS ===", KP, casenum);
        $display("  ek (%0d bytes):", EKB); dump_bytes("ek", EKB);
        $display("  dk (%0d bytes):", DKB); dump_bytes("dk", DKB);

        if (errors == 0) $display("K=%0d CASE %0d PASS: ek (%0dB)==pk, dk (%0dB)==sk", KP, casenum, EKB, DKB);
        else             $display("K=%0d CASE %0d FAIL: %0d mismatches", KP, casenum, errors);
        $finish;
    end
    initial begin #120000000; $display("FAIL: global timeout"); $finish; end
endmodule