// tb_mlkem_enc_katK_xsim.v - ML-KEM Encaps vs NIST KAT, parametric K (KP) + CASE.
// E0 stage: verify H(ek), G(m||H(ek)) -> (ss=K, r). Preloads ek into ek_bram,
// pulses start with op=1, waits for done, checks ss == KAT.ss and dumps H(ek)/r.
//
// xelab -generic_top KP=2|3|4 ; xsim -testplusarg CASE=n
//   ek/m/ct/ss vectors: sync_rtl/top/TB/vectors/enc_k{K}_c{N}_{ek,m,ct,ss}.hex
//   (per-byte hex, byte 0 first).
`timescale 1ns/1ps
module tb_mlkem_enc_katK_xsim;
    parameter KP = 2;
    localparam EKB = 384*KP + 32;                  // ek (=pk) bytes
    localparam CTB = (KP==4) ? 1568 : (32*(10*KP+4)); // ct bytes: K2 768,K3 1088,K4 1568

    reg clk=0, rst_n=0, start_i=0;
    reg [2:0]   k_i;
    reg [255:0] d_i, z_i, m_i;
    wire busy_o, done_o;
    // ek preload port
    reg         ek_in_we=0; reg [10:0] ek_in_addr=0; reg [7:0] ek_in_byte=0;
    wire [255:0] ss_o;
    reg  [10:0] dbg_ct_idx_i=0; wire [7:0] dbg_ct_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, dbg_r_o, dbg_hek_o;

    mlkem_top dut (
        .clk(clk), .rst_n(rst_n), .k_i(k_i), .op_i(1'b1),
        .d_i(d_i), .z_i(z_i), .msg_i(m_i), .start_i(start_i),
        .busy_o(busy_o), .done_o(done_o),
        .ek_in_we(ek_in_we), .ek_in_addr(ek_in_addr), .ek_in_byte(ek_in_byte),
        .ss_o(ss_o), .dbg_ct_idx_i(dbg_ct_idx_i), .dbg_ct_o(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_r_o), .dbg_hek_o(dbg_hek_o)
    );
    always #5 clk = ~clk;

    reg [7:0] ek_b  [0:EKB-1];
    reg [7:0] m_b   [0:31];
    reg [7:0] ss_b  [0:31];
    integer c, i, errors, casenum, j;
    reg [8*80-1:0] tag, ekfile, mfile, ssfile;

    initial begin
        if (!$value$plusargs("CASE=%d", casenum)) casenum = 0;
        $sformat(tag, "k%0d", KP);
        $sformat(ekfile, "sync_rtl/top/TB/vectors/enc_%0s_c%0d_ek.hex", tag, casenum);
        $sformat(mfile,  "sync_rtl/top/TB/vectors/enc_%0s_c%0d_m.hex",  tag, casenum);
        $sformat(ssfile, "sync_rtl/top/TB/vectors/enc_%0s_c%0d_ss.hex", tag, casenum);
        $readmemh(ekfile, ek_b);
        $readmemh(mfile,  m_b);
        $readmemh(ssfile, ss_b);

        // build m_i: byte i in m_i[8*i +: 8]
        m_i = 256'd0;
        for (j = 0; j < 32; j = j + 1) m_i[8*j +: 8] = m_b[j];
        k_i = KP[2:0];

        $display("=== ML-KEM K=%0d Encaps KAT case %0d (E0) ===", KP, casenum);
        $write("  m  = "); for (j=0;j<32;j=j+1) $write("%02x", m_b[j]); $write("\n");

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

        // ---- preload ek into ek_bram (1 byte/cycle) ----
        for (i = 0; i < EKB; i = i + 1) begin
            ek_in_we = 1'b1; ek_in_addr = i[10:0]; ek_in_byte = ek_b[i];
            @(posedge clk);
        end
        ek_in_we = 1'b0; @(posedge clk);

        // ---- run Encaps ----
        start_i=1; @(posedge clk); start_i=0;
        c=0; while(!done_o && c<2000000) begin @(posedge clk); c=c+1; end
        if(!done_o) begin $display("FAIL K=%0d case %0d: timeout", KP, casenum); $finish; end
        $display("=== Encaps E0 done in %0d cyc ===", c);

        $write("  H(ek) = "); for (j=0;j<32;j=j+1) $write("%02x", dbg_hek_o[8*j +: 8]); $write("\n");
        $write("  r     = "); for (j=0;j<32;j=j+1) $write("%02x", dbg_r_o[8*j +: 8]); $write("\n");
        $write("  ss    = "); for (j=0;j<32;j=j+1) $write("%02x", ss_o[8*j +: 8]); $write("\n");

        // ---- check ss == KAT.ss ----
        errors = 0;
        for (j = 0; j < 32; j = j + 1)
            if (ss_o[8*j +: 8] !== ss_b[j]) begin
                if (errors < 8) $display("  SS[%0d] got=%02x exp=%02x", j, ss_o[8*j +: 8], ss_b[j]);
                errors = errors + 1;
            end

        if (errors == 0) $display("K=%0d CASE %0d PASS (E0): ss == KAT.ss", KP, casenum);
        else             $display("K=%0d CASE %0d FAIL (E0): %0d ss mismatches", KP, casenum, errors);

        // ---- E1: verify A_hat (slots 0..K^2-1) and t_hat (slots slot_t..) ----
        // The KAT case-0 Encaps reuses ek=pk from KeyGen case-0, so regenerated
        // A_hat and decoded t_hat equal KeyGen's golden vectors (K=2 c0 only).
        if (KP == 2 && casenum == 0) verify_e1;
        $finish;
    end

    // golden coeff arrays (K=2 c0): A_hat[i][j] and t_hat[i], 256 coeffs each
    reg [11:0] ga [0:4*256-1];   // A00,A01,A10,A11 concatenated (slot*256+idx)
    reg [11:0] gt [0:2*256-1];   // t_hat0, t_hat1
    integer ce, slot, idx;
    task verify_e1;
        begin
            $readmemh("test_framework/modules/mlkem_keygen/golden/c000_Ahat_0_0.hex", ga, 0,    255);
            $readmemh("test_framework/modules/mlkem_keygen/golden/c000_Ahat_0_1.hex", ga, 256,  511);
            $readmemh("test_framework/modules/mlkem_keygen/golden/c000_Ahat_1_0.hex", ga, 512,  767);
            $readmemh("test_framework/modules/mlkem_keygen/golden/c000_Ahat_1_1.hex", ga, 768, 1023);
            $readmemh("test_framework/modules/mlkem_keygen/golden/c000_that_0.hex",   gt, 0,    255);
            $readmemh("test_framework/modules/mlkem_keygen/golden/c000_that_1.hex",   gt, 256,  511);
            ce = 0;
            // A_hat: slots 0..3 (i*k+j)
            for (slot = 0; slot < 4; slot = slot + 1)
                for (idx = 0; idx < 256; idx = idx + 1) begin
                    dbg_slot_i = slot[3:0]; dbg_idx_i = idx[7:0];
                    @(posedge clk); @(posedge clk); @(posedge clk);
                    if (dbg_coeff_o !== ga[slot*256+idx]) begin
                        if (ce < 8) $display("  A[s%0d,%0d] got=%03x exp=%03x", slot, idx, dbg_coeff_o, ga[slot*256+idx]);
                        ce = ce + 1;
                    end
                end
            // t_hat: slots slot_t_rt(=8) .. +1
            for (slot = 0; slot < 2; slot = slot + 1)
                for (idx = 0; idx < 256; idx = idx + 1) begin
                    dbg_slot_i = (8+slot); dbg_idx_i = idx[7:0];
                    @(posedge clk); @(posedge clk); @(posedge clk);
                    if (dbg_coeff_o !== gt[slot*256+idx]) begin
                        if (ce < 12) $display("  T[%0d,%0d] got=%03x exp=%03x", slot, idx, dbg_coeff_o, gt[slot*256+idx]);
                        ce = ce + 1;
                    end
                end
            if (ce == 0) $display("K=2 CASE 0 PASS (E1): A_hat (1024) + t_hat (512) == KeyGen golden");
            else         $display("K=2 CASE 0 FAIL (E1): %0d coeff mismatches", ce);
        end
    endtask
    initial begin #120000000; $display("FAIL: global timeout"); $finish; end
endmodule