// tb_mlkem_dec_katK_xsim.v - ML-KEM Decaps vs NIST KAT, parametric K (KP) + CASE.
// D0 stage: stream dk (=sk) into the design via dk_in_* (routed to
// dk_pke/ek_pke/h/z by region) and ct via c_in_*, pulse start with op=2, and
// verify the dk PARSE: H(ek) (dbg_dech_o), z (dbg_decz_o), and round-trip a few
// ek_pke bytes (dbg_byte sel=0) and dk_pke bytes (sel=1) back out of BRAM.
//
// xelab -generic_top KP=2|3|4 ; xsim -testplusarg CASE=n
//   dk/ct/ss vectors: sync_rtl/top/TB/vectors/dec_k{K}_c{N}_{dk,ct,ss,ctn,ssn}.hex
`timescale 1ns/1ps
module tb_mlkem_dec_katK_xsim;
    parameter KP = 2;
    localparam DKB = 768*KP + 96;                  // dk (=sk) bytes: 1632/2400/3168
    localparam EKB = 384*KP + 32;                  // ek_pke bytes within dk
    localparam DKPB = 384*KP;                      // dk_pke bytes
    localparam CTB = (KP==4) ? 1568 : (32*(10*KP+4)); // ct bytes: 768/1088/1568

    reg clk=0, rst_n=0, start_i=0;
    reg [2:0]   k_i;
    reg [255:0] d_i=0, z_i=0, m_i=0;
    wire busy_o, done_o;
    // ek preload port (unused in Decaps; ek_pke comes from dk)
    reg         ek_in_we=0; reg [10:0] ek_in_addr=0; reg [7:0] ek_in_byte=0;
    // dk / c input ports
    reg         dk_in_we=0; reg [11:0] dk_in_addr=0; reg [7:0] dk_in_byte=0;
    reg         c_in_we=0;  reg [10:0] c_in_addr=0;  reg [7:0] c_in_byte=0;
    wire [255:0] ss_o;
    reg  [10:0] dbg_ct_idx_i=0; wire [7:0] dbg_ct_o;
    reg [5: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;
    wire [255:0] dbg_mprime_o, dbg_kbar_o, dbg_decz_o, dbg_dech_o;

    mlkem_top dut (
        .clk(clk), .rst_n(rst_n), .k_i(k_i), .op_i(2'd2),
        .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),
        .dk_in_we(dk_in_we), .dk_in_addr(dk_in_addr), .dk_in_byte(dk_in_byte),
        .c_in_we(c_in_we), .c_in_addr(c_in_addr), .c_in_byte(c_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),
        .dbg_mprime_o(dbg_mprime_o), .dbg_kbar_o(dbg_kbar_o),
        .dbg_decz_o(dbg_decz_o), .dbg_dech_o(dbg_dech_o)
    );
    always #5 clk = ~clk;

    reg [7:0] dk_b [0:DKB-1];
    reg [7:0] ct_b [0:CTB-1];
    reg [7:0] ctn_b [0:CTB-1];
    reg [7:0] ss_b [0:31];
    reg [7:0] ssn_b [0:31];
    integer c, i, j, errors, casenum;
    reg [8*80-1:0] tag, dkfile, ctfile, ssfile, ctnfile, ssnfile;

    initial begin
        if (!$value$plusargs("CASE=%d", casenum)) casenum = 0;
        $sformat(tag, "k%0d", KP);
        $sformat(dkfile, "sync_rtl/top/TB/vectors/dec_%0s_c%0d_dk.hex", tag, casenum);
        $sformat(ctfile, "sync_rtl/top/TB/vectors/dec_%0s_c%0d_ct.hex", tag, casenum);
        $sformat(ssfile, "sync_rtl/top/TB/vectors/dec_%0s_c%0d_ss.hex", tag, casenum);
        $sformat(ctnfile, "sync_rtl/top/TB/vectors/dec_%0s_c%0d_ctn.hex", tag, casenum);
        $sformat(ssnfile, "sync_rtl/top/TB/vectors/dec_%0s_c%0d_ssn.hex", tag, casenum);
        $readmemh(dkfile, dk_b);
        $readmemh(ctfile, ct_b);
        $readmemh(ssfile, ss_b);
        $readmemh(ctnfile, ctn_b);
        $readmemh(ssnfile, ssn_b);
        k_i = KP[2:0];

        $display("=== ML-KEM K=%0d Decaps KAT case %0d (D0: load+parse) ===", KP, casenum);

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

        // ---- stream dk into the design (1 byte/cycle); persists for both runs ----
        for (i = 0; i < DKB; i = i + 1) begin
            dk_in_we = 1'b1; dk_in_addr = i[11:0]; dk_in_byte = dk_b[i];
            @(posedge clk);
        end
        dk_in_we = 1'b0;

        errors = 0;

        // ================= Run 1: valid ciphertext (accept path) =================
        // ct == c, so c' == c, no implicit reject: ss == K' == KAT ss.
        run_decaps(0);   // 0 = use ct_b
        verify_d0;        // dk parse (errors NOT reset here; see verify_d0 note)
        verify_d4;        // m' (in m_r, survives the re-encrypt)
        verify_d5;        // K'/r'/K-bar
        verify_d6;        // c' == golden
        verify_d7(0);     // ss_o == KAT ss, dec_reject == 0

        // ================= Run 2: corrupted ciphertext (reject path) =============
        // ct_n != c, so c' != c -> implicit reject: ss == K-bar == J(z||ct_n) == ss_n.
        run_decaps(1);   // 1 = use ctn_b
        verify_d7(1);     // ss_o == KAT ss_n, dec_reject == 1

        if (errors == 0) $display("K=%0d CASE %0d PASS (D7): end-to-end Decaps (accept + reject) OK", KP, casenum);
        else             $display("K=%0d CASE %0d FAIL (D7): %0d total errors", KP, casenum, errors);
        $finish;
    end

    initial begin #120000000; $display("FAIL: global timeout"); $finish; end

    // read one stored coefficient at (slot, idx): present addr, wait for the
    // bank read (sd_bram, 1 cyc) + dbg_coeff_r register (1 cyc) + settle.
    task rdcoeff;
        input  [5:0] slot; input [7:0] idx; output [11:0] val;
        begin
            dbg_slot_i = slot; dbg_idx_i = idx;
            @(posedge clk); @(posedge clk); @(posedge clk);
            val = dbg_coeff_o;
        end
    endtask

    // Stream a ciphertext into c_in_bram (sel 0 = valid ct_b, 1 = corrupted ctn_b),
    // pulse start_i, and wait for done_o. dk is already loaded and persists.
    task run_decaps;
        input sel;
        integer ii;
        begin
            for (ii = 0; ii < CTB; ii = ii + 1) begin
                c_in_we = 1'b1; c_in_addr = ii[10:0];
                c_in_byte = sel ? ctn_b[ii] : ct_b[ii];
                @(posedge clk);
            end
            c_in_we = 1'b0; @(posedge clk);
            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 (sel=%0d)", KP, casenum, sel); $finish; end
            $display("=== Decaps run (sel=%0d) done in %0d cyc ===", sel, c);
            // settle so dbg taps reflect the finished run
            repeat(2) @(posedge clk);
        end
    endtask

    // D7: end-to-end shared secret. sel=0 (accept): ss_o == KAT ss, dec_reject=0.
    // sel=1 (reject): ss_o == KAT ss_n (= K-bar = J(z||ct_n)), dec_reject=1.
    task verify_d7;
        input sel;
        integer be;
        reg [255:0] exp;
        begin
            be = 0;
            for (j = 0; j < 32; j = j + 1)
                exp[8*j +: 8] = sel ? ssn_b[j] : ss_b[j];
            for (j = 0; j < 32; j = j + 1)
                if (ss_o[8*j +: 8] !== exp[8*j +: 8]) begin
                    if (be < 4) $display("  ss[%0d] got=%02x exp=%02x", j, ss_o[8*j +: 8], exp[8*j +: 8]);
                    be = be + 1;
                end
            if (be == 0)
                $display("  PASS: ss == KAT %0s (32 bytes), reject=%b",
                    sel ? "ss_n (reject)" : "ss (accept)", dut.dec_reject);
            else
                $display("  FAIL: ss %0d byte mismatches (sel=%0d, reject=%b)", be, sel, dut.dec_reject);
            errors = errors + be;
            // also check the reject flag matches the expectation
            if (dut.dec_reject !== sel) begin
                $display("  FAIL: dec_reject=%b expected %b", dut.dec_reject, sel);
                errors = errors + 1;
            end
        end
    endtask

    // D0: verify dk parse. H(ek)=dk[768K+32:+32], z=dk[768K+64:+32] captured into
    // hek_r/z_r (dbg_dech_o/dbg_decz_o). ek_pke=dk[384K:768K+32] in ek_bram
    // (dbg_byte sel=0), dk_pke=dk[0:384K] in dkp_bram (sel=1).
    task verify_d0;
        integer be;
        reg [7:0] got;
        begin
            // errors is initialized by the caller (main flow), not here.
            // H(ek)
            be = 0;
            for (j = 0; j < 32; j = j + 1)
                if (dbg_dech_o[8*j +: 8] !== dk_b[DKPB + EKB + j]) be = be + 1;
            if (be == 0) $display("  PASS: H(ek) parsed == dk[768K+32 ..]");
            else         $display("  FAIL: H(ek) %0d byte mismatches", be);
            errors = errors + be;

            // z
            be = 0;
            for (j = 0; j < 32; j = j + 1)
                if (dbg_decz_o[8*j +: 8] !== dk_b[DKPB + EKB + 32 + j]) be = be + 1;
            if (be == 0) $display("  PASS: z parsed == dk[768K+64 ..]");
            else         $display("  FAIL: z %0d byte mismatches", be);
            errors = errors + be;

            // ek_pke round-trip (every 97th byte to keep it quick)
            be = 0;
            for (i = 0; i < EKB; i = i + 97) begin
                dbg_byte_sel_i = 1'b0; dbg_byte_idx_i = i[10:0];
                @(posedge clk); @(posedge clk);
                if (dbg_byte_o !== dk_b[DKPB + i]) begin
                    if (be < 6) $display("  ekpke[%0d] got=%02x exp=%02x", i, dbg_byte_o, dk_b[DKPB+i]);
                    be = be + 1;
                end
            end
            if (be == 0) $display("  PASS: ek_pke round-trip (BRAM) == dk[384K ..]");
            else         $display("  FAIL: ek_pke %0d byte mismatches", be);
            errors = errors + be;

            // dk_pke round-trip
            be = 0;
            for (i = 0; i < DKPB; i = i + 97) begin
                dbg_byte_sel_i = 1'b1; dbg_byte_idx_i = i[10:0];
                @(posedge clk); @(posedge clk);
                if (dbg_byte_o !== dk_b[i]) begin
                    if (be < 6) $display("  dkpke[%0d] got=%02x exp=%02x", i, dbg_byte_o, dk_b[i]);
                    be = be + 1;
                end
            end
            if (be == 0) $display("  PASS: dk_pke round-trip (BRAM) == dk[0 ..]");
            else         $display("  FAIL: dk_pke %0d byte mismatches", be);
            errors = errors + be;
            if (errors == 0) $display("  D0 subcheck PASS: dk parse OK");
            else             $display("  D0 subcheck FAIL: %0d errors", errors);
        end
    endtask

    // D1: verify u'[i] (bank_se rel i, abs slot K*K+i) and v' (bank_t rel
    // DEC_VSLOT=2, abs slot K*K+2*K+2) against the decode-decompress golden.
    reg [11:0] up_g [0:255];
    reg [11:0] vp_g [0:255];
    task verify_d1;
        integer i, j, be, ndiff;
        reg [8*100-1:0] fn;
        reg [11:0] got;
        begin
            ndiff = 0;
            // u'[i] is NOT checked here: D2's forward NTT transforms u' in place
            // in bank_se rel 0..K-1, so by the time the run finishes those slots
            // hold u_hat. u' correctness is proven transitively in verify_d2
            // (u_hat == NTT(u') golden). Only v' (bank_a slot 1, untouched) checked.
            // v' lives in bank_a DEC_VASLOT=1 (abs slot 1).
            $sformat(fn, "sync_rtl/top/TB/vectors/decgold/dc_k%0d_c%0d_vp.hex", KP, casenum);
            $readmemh(fn, vp_g);
            be = 0;
            for (j = 0; j < 256; j = j + 1) begin
                rdcoeff(1, j[7:0], got);
                if (got !== vp_g[j]) begin
                    if (be < 4) $display("  v'[%0d] got=%03x exp=%03x", j, got, vp_g[j]);
                    be = be + 1;
                end
            end
            if (be == 0) $display("  PASS: v' == golden (256 coeffs)");
            else         $display("  FAIL: v' %0d coeff mismatches", be);
            ndiff = ndiff + be;
            errors = errors + ndiff;
        end
    endtask

    // D2: verify s_hat[i] (bank_a slot i*K, byteDecode12 dk_pke) and
    // u_hat[i] (bank_se rel slot i, = NTT(u'[i])) against golden.
    reg [11:0] sh_g [0:255];
    reg [11:0] uh_g [0:255];
    task verify_d2;
        integer i, j, be, ndiff;
        reg [8*100-1:0] fn;
        reg [11:0] got;
        begin
            ndiff = 0;
            for (i = 0; i < KP; i = i + 1) begin
                // s_hat[i] at bank_a slot i*KP
                $sformat(fn, "sync_rtl/top/TB/vectors/decgold/dc_k%0d_c%0d_shat_%0d.hex", KP, casenum, i);
                $readmemh(fn, sh_g);
                be = 0;
                for (j = 0; j < 256; j = j + 1) begin
                    rdcoeff(i*KP, j[7:0], got);
                    if (got !== sh_g[j]) begin
                        if (be < 4) $display("  s_hat[%0d][%0d] got=%03x exp=%03x", i, j, got, sh_g[j]);
                        be = be + 1;
                    end
                end
                if (be == 0) $display("  PASS: s_hat[%0d] == golden", i);
                else         $display("  FAIL: s_hat[%0d] %0d mismatches", i, be);
                ndiff = ndiff + be;
                // u_hat[i] at bank_se rel slot i (abs KP*KP + i)
                $sformat(fn, "sync_rtl/top/TB/vectors/decgold/dc_k%0d_c%0d_uhat_%0d.hex", KP, casenum, i);
                $readmemh(fn, uh_g);
                be = 0;
                for (j = 0; j < 256; j = j + 1) begin
                    rdcoeff(KP*KP + i, j[7:0], got);
                    if (got !== uh_g[j]) begin
                        if (be < 4) $display("  u_hat[%0d][%0d] got=%03x exp=%03x", i, j, got, uh_g[j]);
                        be = be + 1;
                    end
                end
                if (be == 0) $display("  PASS: u_hat[%0d] == golden", i);
                else         $display("  FAIL: u_hat[%0d] %0d mismatches", i, be);
                ndiff = ndiff + be;
            end
            errors = errors + ndiff;
        end
    endtask

    // D3: verify w (bank_t rel UPSUM=1, abs slot K*K+2*K+1) == golden.
    reg [11:0] w_g [0:255];
    task verify_d3;
        integer j, be;
        reg [8*100-1:0] fn;
        reg [11:0] got;
        begin
            $sformat(fn, "sync_rtl/top/TB/vectors/decgold/dc_k%0d_c%0d_w.hex", KP, casenum);
            $readmemh(fn, w_g);
            be = 0;
            for (j = 0; j < 256; j = j + 1) begin
                rdcoeff(KP*KP + 2*KP + 1, j[7:0], got);
                if (got !== w_g[j]) begin
                    if (be < 4) $display("  w[%0d] got=%03x exp=%03x", j, got, w_g[j]);
                    be = be + 1;
                end
            end
            if (be == 0) $display("  PASS: w == golden (256 coeffs)");
            else         $display("  FAIL: w %0d coeff mismatches", be);
            errors = errors + be;
        end
    endtask

    // D4: verify m' = byteEncode_1(Compress_1(w)) == golden (32 bytes via dbg).
    reg [7:0] mp_g [0:31];
    task verify_d4;
        integer j, be;
        reg [8*100-1:0] fn;
        begin
            $sformat(fn, "sync_rtl/top/TB/vectors/decgold/dc_k%0d_c%0d_mprime.hex", KP, casenum);
            $readmemh(fn, mp_g);
            be = 0;
            for (j = 0; j < 32; j = j + 1) begin
                if (dbg_mprime_o[8*j +: 8] !== mp_g[j]) begin
                    if (be < 4) $display("  m'[%0d] got=%02x exp=%02x", j, dbg_mprime_o[8*j +: 8], mp_g[j]);
                    be = be + 1;
                end
            end
            if (be == 0) $display("  PASS: m' == golden (32 bytes)");
            else         $display("  FAIL: m' %0d byte mismatches", be);
            errors = errors + be;
        end
    endtask

    // D5: verify K' (ss_o), r' (dbg_r_o), K-bar (dbg_kbar_o) against golden.
    reg [7:0] kp_b [0:31];
    reg [7:0] rp_b [0:31];
    reg [7:0] kb_b [0:31];
    task verify_d5;
        integer j, be;
        reg [8*100-1:0] fn;
        begin
            // K' = G(m'||h) low half -> ss_o
            $sformat(fn, "sync_rtl/top/TB/vectors/decgold/dc_k%0d_c%0d_kprime.hex", KP, casenum);
            $readmemh(fn, kp_b);
            be = 0;
            for (j = 0; j < 32; j = j + 1)
                if (ss_o[8*j +: 8] !== kp_b[j]) begin
                    if (be < 4) $display("  K'[%0d] got=%02x exp=%02x", j, ss_o[8*j +: 8], kp_b[j]);
                    be = be + 1;
                end
            if (be == 0) $display("  PASS: K' == golden (32 bytes)");
            else         $display("  FAIL: K' %0d byte mismatches", be);
            errors = errors + be;
            // r' = G(m'||h) high half -> dbg_r_o
            $sformat(fn, "sync_rtl/top/TB/vectors/decgold/dc_k%0d_c%0d_rprime.hex", KP, casenum);
            $readmemh(fn, rp_b);
            be = 0;
            for (j = 0; j < 32; j = j + 1)
                if (dbg_r_o[8*j +: 8] !== rp_b[j]) begin
                    if (be < 4) $display("  r'[%0d] got=%02x exp=%02x", j, dbg_r_o[8*j +: 8], rp_b[j]);
                    be = be + 1;
                end
            if (be == 0) $display("  PASS: r' == golden (32 bytes)");
            else         $display("  FAIL: r' %0d byte mismatches", be);
            errors = errors + be;
            // K-bar = J(z||c) -> dbg_kbar_o
            $sformat(fn, "sync_rtl/top/TB/vectors/decgold/dc_k%0d_c%0d_kbar.hex", KP, casenum);
            $readmemh(fn, kb_b);
            be = 0;
            for (j = 0; j < 32; j = j + 1)
                if (dbg_kbar_o[8*j +: 8] !== kb_b[j]) begin
                    if (be < 4) $display("  K-bar[%0d] got=%02x exp=%02x", j, dbg_kbar_o[8*j +: 8], kb_b[j]);
                    be = be + 1;
                end
            if (be == 0) $display("  PASS: K-bar == golden (32 bytes)");
            else         $display("  FAIL: K-bar %0d byte mismatches", be);
            errors = errors + be;
        end
    endtask

    // D6: verify c' = K-PKE.Encrypt(ek_pke, m', r') in ct_bram == golden.
    // ct length = 32*(du*K + dv): K2=768, K3=1088, K4=1568.
    reg [7:0] cp_b [0:1567];
    task verify_d6;
        integer i, be, ctlen;
        reg [8*100-1:0] fn;
        reg [7:0] got;
        begin
            ctlen = (KP == 2) ? 768 : (KP == 3) ? 1088 : 1568;
            $sformat(fn, "sync_rtl/top/TB/vectors/decgold/dc_k%0d_c%0d_cprime.hex", KP, casenum);
            $readmemh(fn, cp_b);
            be = 0;
            for (i = 0; i < ctlen; i = i + 1) begin
                dbg_ct_idx_i = i[10:0];
                @(posedge clk); @(posedge clk); @(posedge clk);
                got = dbg_ct_o;
                if (got !== cp_b[i]) begin
                    if (be < 6) $display("  c'[%0d] got=%02x exp=%02x", i, got, cp_b[i]);
                    be = be + 1;
                end
            end
            if (be == 0) $display("  PASS: c' == golden (%0d bytes)", ctlen);
            else         $display("  FAIL: c' %0d byte mismatches", be);
            errors = errors + be;
        end
    endtask
endmodule
