// mlkem_top.v - ML-KEM KeyGen top-level integration. Runtime-selectable
// parameter set via k_i: k=2 (ML-KEM-512, eta1=3), k=3 (768), k=4 (1024).
// Storage is sized for KMAX (worst case = ML-KEM-1024); k_i picks the
// active sub-range at start_i.
//
// Streaming valid/ready interface. Given seeds d and z, computes the
// ML-KEM key pair per FIPS 203 Algorithm 16 (KeyGen_internal):
//   (rho,sigma) = G(d || K)
//   A_hat[i][j] = SampleNTT(rho || j || i)        i,j in 0..K-1
//   s[i] = CBD3(PRF(sigma, i)),  e[i] = CBD3(PRF(sigma, K+i))
//   s_hat[i] = NTT(s[i]);  e_hat[i] = NTT(e[i])
//   t_hat[i] = e_hat[i] + sum_j A_hat[i][j] o s_hat[j]
//   ek = byteEncode12(t_hat[0..K-1]) || rho
//   dk = byteEncode12(s_hat[0..K-1]) || ek || H(ek) || z
//
// Built incrementally and verified stage-by-stage against ml-kem-r golden
// vectors (test_framework/modules/mlkem_keygen/golden) and NIST KAT.
//
// Uses independent (verified) leaf modules, each with its own keccak_core:
// sha3_top, sample_ntt_sync, sample_cbd_sync, ntt_core, poly_mul_sync,
// mod_add_sync. No shared-keccak arbiter.

`include "sync_rtl/common/defines.vh"

module mlkem_top #(
    parameter KMAX = 4               // storage sizing (worst case = ML-KEM-1024)
) (
    input              clk,
    input              rst_n,
    input  [2:0]       k_i,          // RUNTIME ML-KEM param: 2=512, 3=768, 4=1024
    input  [255:0]     d_i,          // KeyGen seed d (byte 0 in d_i[7:0])
    input  [255:0]     z_i,          // implicit-rejection seed z
    input              start_i,      // pulse to begin KeyGen
    output             busy_o,       // high while running
    output             done_o,       // pulse when ek/dk ready

    // Debug readback tap: read one stored coefficient by (poly slot, index).
    // Lets stage TBs verify intermediates without wide buses.
    input  [3:0]       dbg_slot_i,   // poly slot (see localparams below)
    input  [7:0]       dbg_idx_i,    // coefficient index 0..255
    output [11:0]      dbg_coeff_o,
    // Debug byte readback: ek (sel=0, 0..799) / dk_pke (sel=1, 0..767)
    input              dbg_byte_sel_i,
    input  [10:0]      dbg_byte_idx_i,
    output [7:0]       dbg_byte_o,
    // Debug full-dk readback: dk = dk_pke(768) || ek(800) || H(ek)(32) || z(32)
    // = 1632 bytes. Index 0..1631.
    input  [11:0]      dbg_dk_idx_i,
    output [7:0]       dbg_dk_o,
    // Debug taps for hash outputs
    output [255:0]     dbg_rho_o,
    output [255:0]     dbg_sigma_o
);

    localparam Q = `Q;   // 3329

    // Runtime ML-KEM parameter, captured at start_i.
    reg  [2:0] k_r;
    // FIPS 203: eta1 = 3 for ML-KEM-512 (k=2), else 2 (k=3/4).
    wire [1:0] eta1_rt = (k_r == 3'd2) ? 2'd3 : 2'd2;

    // Runtime-derived sizes (k_r in {2,3,4}). Small multiplies are cheap.
    wire [5:0] kk_rt     = k_r * k_r;            // 4/9/16
    wire [5:0] slot_s_rt = kk_rt;                // s_hat base slot
    wire [5:0] slot_e_rt = kk_rt + k_r;          // e_hat base slot
    wire [5:0] slot_t_rt = kk_rt + {1'b0, k_r} + {1'b0, k_r};  // t_hat base = kk+2k
    wire [11:0] ek_bytes_rt = 12'd384 * {9'b0, k_r} + 12'd32;  // 800/1184/1568
    wire [11:0] dk_bytes_rt = 12'd384 * {9'b0, k_r};            // 768/1152/1536
    // H(ek) block count = ceil((ek_bytes+1)/136): 6/9/12 for k=2/3/4 (table)
    wire [3:0] h_nblk_rt = (k_r == 3'd2) ? 4'd6 : (k_r == 3'd3) ? 4'd9 : 4'd12;
    wire [11:0] h_last_rt = {6'b0, h_nblk_rt} * 12'd136 - 12'd1;  // final padded byte index

    // ================================================================
    // Polynomial storage, sized for KMAX (worst case). Runtime k uses a
    // sub-range. Slot layout (each slot = 256 coeffs):
    //   A_hat[i][j] : slots 0 .. k*k-1        at index i*k + j
    //   s_hat[i]    : slots slot_s_rt .. +k-1
    //   e_hat[i]    : slots slot_e_rt .. +k-1
    //   t_hat[i]    : slots slot_t_rt .. +k-1
    // NUM_SLOTS = KMAX*KMAX + 3*KMAX = 28 for KMAX=4.
    // ================================================================
    localparam NUM_SLOTS = KMAX*KMAX + 3*KMAX;
    localparam SAW       = 5;             // slot-address width (>=clog2(28))

    // ================================================================
    // Coefficient storage as 3 banks (Phase 2). Was one async-read reg
    // array polymem[0:28*256-1]; split by polynomial so each becomes a
    // registered-read sd_bram (infers BRAM / ASIC SRAM) in the final step.
    //   bank_a : A_hat[i][j]  slots 0..K^2-1   -> D=KMAX*KMAX*256=4096
    //   bank_se: s_hat||e_hat slots slot_s..   -> D=2*KMAX*256=2048 (rel 0..2K-1)
    //   bank_t : t_hat[i]     slots slot_t..   -> D=KMAX*256=1024  (rel 0..K-1)
    // Addresses are base-relative: bank index = abs_slot - base_slot.
    // Stage 2b (in progress): per-consumer registered read-ahead. Banks are
    // still async reg arrays here; converted consumers read via their own
    // 1-cycle pipeline reg (== sd_bram timing). Once every consumer of a
    // bank is registered, the array is replaced by an sd_bram with its read
    // port muxed across the (phase-disjoint) consumers.
    // ================================================================
    localparam PA_AW = 12;                // bank_a addr width (4096)
    localparam PSE_AW = 11;               // bank_se addr width (2048)
    localparam PT_AW = 10;                // bank_t addr width (1024)
    reg [11:0] bank_a  [0:(1<<PA_AW)-1];
    reg [11:0] bank_se [0:(1<<PSE_AW)-1];
    reg [11:0] bank_t  [0:(1<<PT_AW)-1];

    // Debug readback (registered for timing)
    reg [11:0] dbg_coeff_r;
    // Route debug read by absolute slot range (not on the KAT correctness
    // path — TB reads ek/dk BRAMs only). A: 0..K^2-1; s/e: slot_s..slot_t-1;
    // t: slot_t..  Index within bank is (abs_slot-base)*256 + idx.
    wire [13:0] dbg_a_addr  = dbg_slot_i*256 + dbg_idx_i;
    wire [13:0] dbg_se_addr = (dbg_slot_i - slot_s_rt)*256 + dbg_idx_i;
    wire [13:0] dbg_t_addr  = (dbg_slot_i - slot_t_rt)*256 + dbg_idx_i;
    always @(posedge clk) begin
        if (dbg_slot_i >= slot_t_rt)      dbg_coeff_r <= bank_t [dbg_t_addr[PT_AW-1:0]];
        else if (dbg_slot_i >= slot_s_rt) dbg_coeff_r <= bank_se[dbg_se_addr[PSE_AW-1:0]];
        else                              dbg_coeff_r <= bank_a [dbg_a_addr[PA_AW-1:0]];
    end
    assign dbg_coeff_o = dbg_coeff_r;

    // ek and dk_pke byte memories sized for KMAX.
    localparam EK_MAX = 384*KMAX + 32;   // 1568
    localparam DK_MAX = 384*KMAX;         // 1536

    // ek / dk_pke byte storage in BRAM (sd_bram: 1R/1W, 1-cycle read latency).
    // Each single read port is shared time-disjointly: ST_H assemble reads ek
    // (in ST_H), debug readback reads ek/dk (in ST_DONE). Write port driven by
    // ST_E (1 byte/cycle) and the rho-tail copy.
    wire [10:0] ek_rd_addr, dkp_rd_addr;
    wire [7:0]  ek_rd_data, dkp_rd_data;
    reg         ek_we, dkp_we;
    reg  [10:0] ek_wa, dkp_wa;
    reg  [7:0]  ek_wd, dkp_wd;

    sd_bram #(.W(8), .D(2048), .A(11)) u_ek_bram (
        .clk(clk),
        .rd_addr(ek_rd_addr), .rd_data(ek_rd_data),
        .wr_en(ek_we), .wr_addr(ek_wa), .wr_data(ek_wd)
    );
    sd_bram #(.W(8), .D(2048), .A(11)) u_dkp_bram (
        .clk(clk),
        .rd_addr(dkp_rd_addr), .rd_data(dkp_rd_data),
        .wr_en(dkp_we), .wr_addr(dkp_wa), .wr_data(dkp_wd)
    );

    // full dk = dk_pke(dk_bytes) || ek(ek_bytes) || H(ek)(32) || z(32)
    wire [11:0] dk_ek_end  = dk_bytes_rt + ek_bytes_rt;   // ek region end
    wire [11:0] dk_hek_end = dk_ek_end + 12'd32;           // H(ek) region end

    // Debug-region selects for dk readback (combinational region decode).
    wire dbgdk_in_dkp = (dbg_dk_idx_i < dk_bytes_rt);
    wire dbgdk_in_ek  = (dbg_dk_idx_i >= dk_bytes_rt) && (dbg_dk_idx_i < dk_ek_end);

    // Read-address muxes (declared here; depend on ST_H h_g defined below via
    // the function-free expression h_blk*136 + h_byte). See assigns after the
    // H-stage register declarations.
    // dbg_byte_o / dbg_dk_o are combinational taps over the (registered) BRAM
    // read data; net 1-cycle latency, well within the TB's 2-cycle read wait.
    assign dbg_byte_o = dbg_byte_sel_i ? dkp_rd_data : ek_rd_data;
    assign dbg_dk_o   = dbgdk_in_dkp ? dkp_rd_data :
                        dbgdk_in_ek  ? ek_rd_data  :
                        (dbg_dk_idx_i < dk_hek_end) ? hek_r[(dbg_dk_idx_i - dk_ek_end)*8 +: 8]
                                                    : z_i[(dbg_dk_idx_i - dk_hek_end)*8 +: 8];

    // ================================================================
    // Top-level FSM (built incrementally). Stage 2a: G only.
    // ================================================================
    localparam ST_IDLE   = 4'd0;
    localparam ST_G      = 4'd1;   // run G(d||K), capture rho/sigma
    localparam ST_A      = 4'd2;   // generate A_hat[i][j] via SampleNTT
    localparam ST_C      = 4'd3;   // generate s[i],e[i] via CBD
    localparam ST_N      = 4'd4;   // forward NTT of s[i],e[i] in place
    localparam ST_M      = 4'd5;   // matrix accumulate t_hat = e_hat + sum A o s_hat
    localparam ST_E      = 4'd6;   // byteEncode12 -> ek/dk BRAM
    localparam ST_H      = 4'd7;   // H(ek) via multi-block SHA3-256
    localparam ST_DONE   = 4'd15;

    reg [3:0]   st, st_next;
    reg [255:0] rho_r, sigma_r;

    // A-generation bookkeeping: explicit i/j counters (avoid runtime divide)
    reg  [2:0]   a_i;          // row 0..k-1
    reg  [2:0]   a_j;          // col 0..k-1
    reg  [4:0]   a_pair;       // 0..k*k pairs done (for done test)
    reg  [7:0]   a_widx;       // write index 0..255 within current poly
    reg          a_busy;       // 1 once current pair's request accepted (gates collect)
    wire [SAW-1:0] a_slot = a_i*k_r + a_j;   // A_hat[i][j] slot = i*k + j

    // C-generation bookkeeping: 2*k polys (s[0..k-1] then e[0..k-1])
    reg  [4:0]   c_poly;       // 0..2k
    reg  [7:0]   c_widx;
    reg          c_busy;
    wire [7:0]   c_nonce = {3'b0, c_poly};       // s:0..k-1  e:k..2k-1  == nonce
    // slot: c_poly < k -> s_hat[c_poly], else e_hat[c_poly-k]
    wire [SAW-1:0] c_slot = (c_poly < {2'b0, k_r}) ? (slot_s_rt + c_poly)
                                                   : (slot_e_rt + (c_poly - {2'b0, k_r}));

    assign busy_o = (st != ST_IDLE);
    assign done_o = (st == ST_DONE);
    assign dbg_rho_o   = rho_r;
    assign dbg_sigma_o = sigma_r;

    // ---- sha3_top in G mode: data_i = {K_byte, d} (d byte0 in [7:0]) ----
    reg          sha3_valid;
    wire         sha3_ready;
    wire [511:0] sha3_hash;
    wire         sha3_vo;
    reg          sha3_ack;       // consumer ready for hash
    wire [511:0] g_data = {248'b0, 5'b0, k_r, d_i};   // data_i[263:256]=k, [255:0]=d

    // ================================================================
    // Shared keccak_core + phase mux (3 consumers -> 1 core)
    //
    // G/H (u_sha3), SampleNTT (u_snt) and CBD (u_cbd) each need a Keccak
    // permutation, but they run in DISJOINT top-FSM phases (ST_G/ST_H, ST_A,
    // ST_C respectively), so at most one is ever active. One keccak_core
    // serves all three: kc_state_i/kc_valid_i are muxed by phase; kc_valid_o
    // is broadcast but GATED per consumer so an inactive consumer never
    // latches a permutation result meant for another (the samplers latch
    // squeeze state unconditionally on their kc_valid_o input).
    // ================================================================
    wire [1599:0] kc_state_o;        // shared core output (to all consumers)
    wire          kc_valid_o;        // shared core output-valid
    wire          kc_ready_o;        // shared core input-ready

    // per-consumer drives toward the core
    wire [1599:0] sha3_kc_state_i, snt_kc_state_i, cbd_kc_state_i;
    wire          sha3_kc_valid_i, snt_kc_valid_i, cbd_kc_valid_i;
    /* verilator lint_off UNUSEDSIGNAL */
    wire          sha3_kc_ready_i, snt_kc_ready_i, cbd_kc_ready_i; // all 1'b1
    /* verilator lint_on UNUSEDSIGNAL */

    // phase selects (mutually exclusive)
    wire sel_sha3 = (st == ST_G) || (st == ST_H);
    wire sel_snt  = (st == ST_A);
    wire sel_cbd  = (st == ST_C);

    // gated output-valid: only the active consumer sees kc_valid_o
    wire kc_valid_o_sha3 = kc_valid_o & sel_sha3;
    wire kc_valid_o_snt  = kc_valid_o & sel_snt;
    wire kc_valid_o_cbd  = kc_valid_o & sel_cbd;

    // input mux: route the active consumer's request to the core
    wire [1599:0] kc_state_i_mux = sel_snt ? snt_kc_state_i :
                                   sel_cbd ? cbd_kc_state_i :
                                             sha3_kc_state_i;
    wire kc_valid_i_mux = sel_snt ? snt_kc_valid_i :
                          sel_cbd ? cbd_kc_valid_i :
                          sel_sha3 ? sha3_kc_valid_i : 1'b0;

    keccak_core #(.ROUNDS(24)) u_keccak (
        .clk(clk), .rst_n(rst_n),
        .state_i(kc_state_i_mux),
        .valid_i(kc_valid_i_mux),
        .ready_o(kc_ready_o),
        .state_o(kc_state_o),
        .valid_o(kc_valid_o),
        .ready_i(1'b1)               // consumers always accept (kc_ready_i=1)
    );

    // ---- single shared sha3_top serving BOTH G and H ----
    // G (ST_G) uses single-block mode (mb_en=0); H(ek) (ST_H) uses the
    // multi-block absorb path (mb_en=1). These phases are disjoint in the
    // top FSM, so one sha3_top (one keccak_core) is sufficient. mb_en and
    // ready_i are muxed by phase; data_i/mode only matter while mb_en=0.
    wire sha3_mb_en = (st == ST_H);
    sha3_top_shared u_sha3 (
        .clk(clk), .rst_n(rst_n),
        .mode(2'b00),                 // G = SHA3-512 (only used when mb_en=0)
        .data_i(g_data),
        .valid_i(sha3_valid),
        .ready_o(sha3_ready),
        .hash_o(sha3_hash),
        .valid_o(sha3_vo),
        .ready_i(sha3_mb_en ? h_ack : sha3_ack),
        .mb_en(sha3_mb_en),
        .mb_block_i(h_block_r),
        .mb_valid_i(h_mbvalid),
        .mb_last_i(h_mblast),
        .mb_ready_o(h_mbready),
        // shared keccak_core interface (gated by phase in the mux below)
        .kc_state_o(kc_state_o),
        .kc_valid_o(kc_valid_o_sha3),
        .kc_ready_o(kc_ready_o),
        .kc_state_i(sha3_kc_state_i),
        .kc_valid_i(sha3_kc_valid_i),
        .kc_ready_i(sha3_kc_ready_i)
    );

    // ---- multi-block H(ek) state (SHA3-256, 6/9/12 blocks); fed to shared u_sha3 ----
    reg  [1087:0] h_block_r;       // current pre-padded rate block
    reg           h_mbvalid;
    reg           h_mblast;
    wire          h_mbready;
    wire [511:0]  h_hash;
    wire          h_vo;
    reg           h_ack;
    reg  [255:0]  hek_r;           // captured H(ek)
    reg  [3:0]    h_blk;           // 0..H_NBLK-1 block index (up to 11 for K=4)
    reg  [7:0]    h_byte;          // 0..136 byte address being presented (assemble)
    reg  [1:0]    h_phase;         // 0=assemble 1=feed 2=wait-perm 3=done
    // ek BRAM read is registered (1-cycle latency); assemble presents the
    // address for h_byte this cycle and writes the byte that arrived (for the
    // address presented last cycle) into h_block_r. Writeback pipeline regs:
    reg           h_wb_vld;        // a byte is arriving this cycle
    reg  [7:0]    h_wb_idx;        // its position within the 136-byte block
    reg  [11:0]   h_wb_g;          // its global ek byte index
    reg  [7:0]    h_wb_pad;        // pad constant to use if g is out of ek range
    reg           h_wb_inek;       // 1 if g is within ek range (use BRAM data)

    // h_hash / h_vo are now served by the shared u_sha3 above (mb_en=1 during
    // ST_H). They alias the single core's outputs; the H consumer logic below
    // already gates on st==ST_H, and u_sha3's valid_o/hash_o are mb-selected.
    assign h_hash = sha3_hash;
    assign h_vo   = sha3_vo;

    // SHA3-256 over ek (ek_bytes_rt bytes): rate=136. Padded length = h_nblk_rt*136.
    //   pad: byte ek_bytes_rt = 0x06 (domain + first pad bit), last byte |= 0x80.
    // byte b (0..135) of block blk: global g = blk*136 + b.
    // ek bytes come from the ek BRAM (registered read); only the pad/zero bytes
    // are constants, returned by h_padconst. Reads runtime ek_bytes_rt/h_last_rt.
    function [7:0] h_padconst(input [3:0] blk, input [7:0] b);
        integer g;
        begin
            g = blk*136 + b;
            if (g == h_last_rt && g == ek_bytes_rt) h_padconst = 8'h86; // 0x06|0x80
            else if (g == ek_bytes_rt)  h_padconst = 8'h06;
            else if (g == h_last_rt)    h_padconst = 8'h80;
            else                        h_padconst = 8'h00;
        end
    endfunction

    // Global ek byte index for the address currently presented in ST_H assemble.
    wire [11:0] h_g_addr = {8'd0, h_blk} * 12'd136 + {4'd0, h_byte};

    // ek BRAM read-address mux: ST_H assemble drives it; otherwise debug readback.
    // dbg_dk ek-region takes priority (its loop runs after the ek loop), so the
    // single read port serves dbg_byte_o (ek loop) then dbg_dk_o (dk loop).
    wire [11:0] dbgdk_ek_off = dbg_dk_idx_i - dk_bytes_rt;   // offset within ek region
    assign ek_rd_addr = (st == ST_H) ? h_g_addr[10:0] :
                        dbgdk_in_ek   ? dbgdk_ek_off[10:0]
                                      : dbg_byte_idx_i;
    // dkp BRAM read-address mux: dbg_byte (sel=1) or dbg_dk (dkp region).
    assign dkp_rd_addr = dbg_byte_sel_i ? dbg_byte_idx_i : dbg_dk_idx_i[10:0];

    // ---- sample_ntt_sync: Â[i][j] = SampleNTT(rho || j || i) ----
    reg          snt_valid;
    wire         snt_ready;
    wire [11:0]  snt_coeff;
    wire         snt_vo;
    wire         snt_last;
    reg          snt_ack;       // we accept coeffs

    sample_ntt_sync_shared #(.K(KMAX)) u_snt (
        .clk(clk), .rst_n(rst_n),
        .rho_i(rho_r),
        .k_i(k_r),
        .i_idx(a_i[1:0]),
        .j_idx(a_j[1:0]),
        .valid_i(snt_valid),
        .ready_o(snt_ready),
        .coeff_o(snt_coeff),
        .valid_o(snt_vo),
        .ready_i(snt_ack),
        .last_o(snt_last),
        // shared keccak_core interface
        .kc_state_o(kc_state_o),
        .kc_valid_o(kc_valid_o_snt),
        .kc_ready_o(kc_ready_o),
        .kc_state_i(snt_kc_state_i),
        .kc_valid_i(snt_kc_valid_i),
        .kc_ready_i(snt_kc_ready_i)
    );

    // ---- sample_cbd_sync: s[i]=CBD3(PRF(sigma,i)), e[i]=CBD3(PRF(sigma,K+i)) ----
    reg          cbd_valid;
    wire         cbd_ready;
    wire [11:0]  cbd_coeff;       // 12-bit signed (two's complement)
    wire         cbd_vo;
    wire         cbd_last;
    reg          cbd_ack;

    sample_cbd_sync_shared u_cbd (
        .clk(clk), .rst_n(rst_n),
        .seed_i(sigma_r),
        .nonce_i(c_nonce),
        .eta_i(eta1_rt),
        .valid_i(cbd_valid),
        .ready_o(cbd_ready),
        .coeff_o(cbd_coeff),
        .valid_o(cbd_vo),
        .ready_i(cbd_ack),
        .last_o(cbd_last),
        // shared keccak_core interface
        .kc_state_o(kc_state_o),
        .kc_valid_o(kc_valid_o_cbd),
        .kc_ready_o(kc_ready_o),
        .kc_state_i(cbd_kc_state_i),
        .kc_valid_i(cbd_kc_valid_i),
        .kc_ready_i(cbd_kc_ready_i)
    );

    // signed (two's complement) -> [0,Q): add Q when negative
    wire [11:0] cbd_modq = cbd_coeff[11] ? (cbd_coeff + 12'(Q)) : cbd_coeff;

    // ---- ntt_core: forward NTT (mode=0, no scaling) of s[i],e[i] in place ----
    // N-stage bookkeeping: process slots S0,S1,E0,E1 (= SLOT_S0 + n_slot).
    reg  [4:0]  n_slot;        // 0..2K (process s_hat[0..K-1] then e_hat[0..K-1])
    reg  [8:0]  n_ridx;        // load read-AHEAD pointer 0..256 (leads consume by 1)
    reg  [7:0]  n_widx;        // output write index 0..255
    reg         n_valid;       // feeding coeffs to ntt_core (delayed 1 cyc vs n_ridx)
    reg         n_loading;     // 1 while presenting load addresses to bank_se
    reg         n_pending;     // waiting for ntt_core IDLE to start next slot
    reg  [11:0] n_rd_data;     // registered bank_se read (== sd_bram timing)
    wire [SAW-1:0] n_slot_addr = slot_s_rt + n_slot;   // s_hat then e_hat contiguous

    wire        ntt_ready;
    wire [11:0] ntt_coeff;
    wire        ntt_vo;
    wire        ntt_done;
    // bank_se read addr for the NTT load (relative slot = n_slot); registered
    // into n_rd_data, which feeds ntt_core 1 cycle later.
    wire [13:0] ntt_rd_full = n_slot*256 + n_ridx[7:0];
    wire [11:0] ntt_in = n_rd_data;

    ntt_core u_ntt (
        .clk(clk), .rst_n(rst_n),
        .coeff_in(ntt_in),
        .valid_i(n_valid),
        .ready_o(ntt_ready),
        .mode(1'b0),                  // forward NTT, no scaling
        .coeff_out(ntt_coeff),
        .valid_o(ntt_vo),
        .ready_i(1'b1),               // always accept output
        .done_o(ntt_done)
    );

    // ---- poly_mul_sync: t_hat[i] = e_hat[i] + sum_j A_hat[i][j] o s_hat[j] ----
    // M-stage bookkeeping. For each (i,j): LOAD 256 (A,shat) pairs, then accumulate
    // 256 products into T_i (init from E_i when j==0, else from running T_i).
    reg  [2:0]  m_i;           // row 0..K (needs to reach K to exit)
    reg  [2:0]  m_j;           // col 0..K-1
    reg  [8:0]  m_ld;          // load index 0..256
    reg  [7:0]  m_oidx;        // output/accum index 0..255
    reg         m_loading;     // 1 while streaming pairs into poly_mul
    reg         m_pending;     // wait for poly_mul IDLE before next (i,j)

    // ---- Stage 2f: byteEncode12 serializer ----
    // Pack each poly (2 coeffs -> 3 bytes, LSB-first 12-bit). ek = t_hat[0..K-1]
    // bytes || rho; dk_pke = s_hat[0..K-1] bytes. Walk coeff pairs per poly.
    reg  [4:0]  e_poly;        // 0..2K-1: [0,K) = t_hat -> ek; [K,2K) = s_hat -> dk_pke
    reg  [7:0]  e_pair;        // 0..127 coeff-pair within poly
    reg  [1:0]  e_byte;        // 0..2 byte within the current pair (BRAM: 1 write/cycle)
    reg  [9:0]  e_rho;         // 0..31 rho byte copy index (ek tail)
    reg         e_done;        // serialization complete
    // source poly slot: t_hat[e_poly] for ek half, s_hat[e_poly-K] for dk half
    wire        e_is_dk = (e_poly >= {1'b0, k_r});
    wire [4:0]  e_pidx  = e_is_dk ? (e_poly - {1'b0, k_r}) : e_poly;   // index within target
    wire [SAW-1:0] e_slot = e_is_dk ? (slot_s_rt + e_pidx) : (slot_t_rt + e_pidx);
    // two coeffs of the current pair: ek half reads t_hat (bank_t), dk half
    // reads s_hat (bank_se). Relative index within the bank = e_pidx.
    wire [13:0] e_rd0_full = e_pidx*256 + {e_pair, 1'b0};
    wire [13:0] e_rd1_full = e_pidx*256 + {e_pair, 1'b1};
    wire [11:0] e_c0 = e_is_dk ? bank_se[e_rd0_full[PSE_AW-1:0]] : bank_t[e_rd0_full[PT_AW-1:0]];
    wire [11:0] e_c1 = e_is_dk ? bank_se[e_rd1_full[PSE_AW-1:0]] : bank_t[e_rd1_full[PT_AW-1:0]];
    // 3 packed bytes
    wire [7:0]  e_b0 = e_c0[7:0];
    wire [7:0]  e_b1 = {e_c1[3:0], e_c0[11:8]};
    wire [7:0]  e_b2 = e_c1[11:4];
    // byte base offset within target memory: poly index *384 (= 128 pairs *3)
    wire [11:0] e_base = e_pidx * 12'd384;
    wire [11:0] e_boff = e_base + {e_pair, 1'b0} + {2'b0, e_pair};  // pair*3
    // BRAM write is 1 byte/cycle: select packed byte by e_byte sub-counter
    wire [7:0]  e_byte_d = (e_byte == 2'd0) ? e_b0 : (e_byte == 2'd1) ? e_b1 : e_b2;

    wire [SAW-1:0]  m_aslot = m_i*k_r + m_j;     // A_hat[i][j] slot = i*k + j
    wire [SAW-1:0]  m_sslot = slot_s_rt + m_j;     // s_hat[j]
    wire [SAW-1:0]  m_eslot = slot_e_rt + m_i;     // e_hat[i]
    wire [SAW-1:0]  m_tslot = slot_t_rt + m_i;     // t_hat[i]

    reg          pm_valid;
    wire         pm_ready;
    wire [11:0]  pm_coeff;
    wire         pm_vo;
    // pm_a: A_hat[i][j] in bank_a (abs slot m_aslot). pm_b: s_hat[j] in
    // bank_se (relative slot = m_sslot - slot_s_rt = m_j). m_ld is a read-ahead
    // pointer; both bank reads are registered into pm_a_rd/pm_b_rd and fed to
    // poly_mul one cycle later (pm_valid delayed 1 cyc), like ST_N.
    wire [13:0] pm_a_full = m_aslot*256 + m_ld[7:0];
    wire [13:0] pm_b_full = m_j*256 + m_ld[7:0];
    reg  [11:0]  pm_a_rd, pm_b_rd;       // registered bank reads (sd_bram timing)
    wire [11:0]  pm_a_in = pm_a_rd;
    wire [11:0]  pm_b_in = pm_b_rd;

    poly_mul_sync u_pmul (
        .clk(clk), .rst_n(rst_n),
        .coeff_a_in(pm_a_in),
        .coeff_b_in(pm_b_in),
        .valid_i(pm_valid),
        .ready_o(pm_ready),
        .coeff_out(pm_coeff),
        .valid_o(pm_vo),
        .ready_i(1'b1)
    );

    // accumulator source: e_hat[i] for first term (j==0), else running t_hat[i].
    // e_hat[i] lives in bank_se at relative slot (slot_e_rt-slot_s_rt + m_i) = K+m_i.
    // t_hat[i] lives in bank_t at relative slot m_i.
    // Registered read-ahead: present the index that the NEXT pm_vo will consume
    //   m_acc_radr = pm_vo ? m_oidx+1 : m_oidx
    // and register the selected bank into m_acc_rd, consumed one cycle later when
    // pm_vo is high. The same-index read-modify-write on bank_t is read-old:
    // the running t_hat source was written by the previous (j-1) term, long
    // settled; the value consumed this cycle was registered before this cycle's
    // write to m_oidx. Cadence is CALC(vo=0)/C0(vo=1)/C1(vo=1).
    wire [7:0]  m_acc_radr = pm_vo ? (m_oidx + 8'd1) : m_oidx;
    wire [13:0] m_eacc_full = ({2'b0, k_r} + {2'b0, m_i})*256 + m_acc_radr;  // K+m_i
    wire [13:0] m_tacc_full = m_i*256 + m_acc_radr;
    wire [11:0] m_acc_rd_next = (m_j == 2'd0) ? bank_se[m_eacc_full[PSE_AW-1:0]]
                                              : bank_t [m_tacc_full[PT_AW-1:0]];
    reg  [11:0] m_acc_rd;
    wire [11:0]  m_acc_src = m_acc_rd;
    // (a + b) mod Q  (both < Q, sum < 2Q): one conditional subtract
    wire [12:0]  m_sum  = {1'b0, m_acc_src} + {1'b0, pm_coeff};
    wire [11:0]  m_accq = (m_sum >= 13'(Q)) ? (m_sum - 13'(Q)) : m_sum[11:0];

    always @(*) begin
        st_next = st;
        case (st)
            ST_IDLE: if (start_i) st_next = ST_G;
            ST_G:    if (sha3_vo) st_next = ST_A;
            ST_A:    if (a_pair >= kk_rt) st_next = ST_C;
            ST_C:    if (c_poly >= {1'b0, k_r, 1'b0}) st_next = ST_N;
            ST_N:    if (n_slot >= {1'b0, k_r, 1'b0}) st_next = ST_M;
            ST_M:    if (m_i >= k_r) st_next = ST_E;
            ST_E:    if (e_done) st_next = ST_H;
            ST_H:    if (h_phase == 2'd3) st_next = ST_DONE;
            ST_DONE: st_next = ST_IDLE;
            default: st_next = ST_IDLE;
        endcase
    end

    always @(posedge clk or negedge rst_n) begin
        if (!rst_n) begin
            st         <= ST_IDLE;
            k_r        <= 3'd0;
            rho_r      <= 256'd0;
            sigma_r    <= 256'd0;
            sha3_valid <= 1'b0;
            sha3_ack   <= 1'b0;
            snt_valid  <= 1'b0;
            snt_ack    <= 1'b0;
            a_pair     <= 5'd0;
            a_i        <= 3'd0;
            a_j        <= 3'd0;
            a_widx     <= 8'd0;
            a_busy     <= 1'b0;
            cbd_valid  <= 1'b0;
            cbd_ack    <= 1'b0;
            c_poly     <= 3'd0;
            c_widx     <= 8'd0;
            c_busy     <= 1'b0;
            n_slot     <= 3'd0;
            n_ridx     <= 9'd0;
            n_widx     <= 8'd0;
            n_valid    <= 1'b0;
            n_loading  <= 1'b0;
            n_rd_data  <= 12'd0;
            n_pending  <= 1'b0;
            m_i        <= 2'd0;
            m_j        <= 2'd0;
            m_ld       <= 9'd0;
            m_oidx     <= 8'd0;
            m_loading  <= 1'b0;
            m_pending  <= 1'b0;
            pm_valid   <= 1'b0;
            pm_a_rd    <= 12'd0;
            pm_b_rd    <= 12'd0;
            m_acc_rd   <= 12'd0;
            e_poly     <= 3'd0;
            e_pair     <= 8'd0;
            e_byte     <= 2'd0;
            e_rho      <= 10'd0;
            e_done     <= 1'b0;
            ek_we      <= 1'b0;
            dkp_we     <= 1'b0;
            ek_wa      <= 11'd0;
            dkp_wa     <= 11'd0;
            ek_wd      <= 8'd0;
            dkp_wd     <= 8'd0;
            h_block_r  <= 1088'd0;
            h_mbvalid  <= 1'b0;
            h_mblast   <= 1'b0;
            h_ack      <= 1'b0;
            hek_r      <= 256'd0;
            h_blk      <= 3'd0;
            h_byte     <= 8'd0;
            h_phase    <= 2'd0;
            h_wb_vld   <= 1'b0;
            h_wb_idx   <= 8'd0;
            h_wb_g     <= 12'd0;
            h_wb_pad   <= 8'd0;
            h_wb_inek  <= 1'b0;
        end else begin
            st <= st_next;

            // BRAM write-enables default low; pulsed where a byte is written.
            ek_we  <= 1'b0;
            dkp_we <= 1'b0;

            // Kick off G when entering ST_G
            if (st == ST_IDLE && start_i) begin
                k_r        <= k_i;          // capture runtime ML-KEM param
                sha3_valid <= 1'b1;
                sha3_ack   <= 1'b1;
            end
            // Drop valid once accepted
            if (sha3_valid && sha3_ready) sha3_valid <= 1'b0;

            // Capture rho/sigma when G completes; arm A stage
            if (st == ST_G && sha3_vo) begin
                rho_r   <= sha3_hash[255:0];     // rho = G output bytes 0..31
                sigma_r <= sha3_hash[511:256];   // sigma = bytes 32..63
                sha3_ack <= 1'b0;
                snt_valid <= 1'b1;               // start first SampleNTT
                snt_ack   <= 1'b1;
                a_pair    <= 5'd0;
                a_i       <= 3'd0;
                a_j       <= 3'd0;
                a_widx    <= 8'd0;
                a_busy    <= 1'b0;
            end

            // ---- ST_A: drive SampleNTT, store 256 coeffs per pair ----
            if (st == ST_A) begin
                // mark busy once this pair's request accepted
                if (snt_valid && snt_ready) begin
                    snt_valid <= 1'b0;
                    a_busy    <= 1'b1;
                end

                // store each output coefficient only while busy (ignore stale last coeff from prior poly)
                if (a_busy && snt_vo && snt_ack) begin
                    bank_a[(a_slot*256 + a_widx) & ((1<<PA_AW)-1)] <= snt_coeff;
                    if (snt_last) begin
                        // finished this poly; advance (i,j) in row-major order
                        a_pair <= a_pair + 5'd1;
                        a_widx <= 8'd0;
                        a_busy <= 1'b0;
                        if (a_j + 3'd1 < k_r) begin
                            a_j <= a_j + 3'd1;
                        end else begin
                            a_j <= 3'd0;
                            a_i <= a_i + 3'd1;
                        end
                        // start next SampleNTT if more pairs remain
                        if (a_pair + 5'd1 < kk_rt) snt_valid <= 1'b1;
                    end else begin
                        a_widx <= a_widx + 8'd1;
                    end
                end
            end

            // Arm C stage when A finishes
            if (st == ST_A && st_next == ST_C) begin
                cbd_valid <= 1'b1;
                cbd_ack   <= 1'b1;
                c_poly    <= 3'd0;
                c_widx    <= 8'd0;
                c_busy    <= 1'b0;
            end

            // ---- ST_C: drive CBD, store 256 mod-q coeffs per poly ----
            if (st == ST_C) begin
                if (cbd_valid && cbd_ready) begin
                    cbd_valid <= 1'b0;
                    c_busy    <= 1'b1;
                end

                if (c_busy && cbd_vo && cbd_ack) begin
                    // bank_se relative slot = c_slot - slot_s_rt = c_poly
                    bank_se[(c_poly*256 + c_widx) & ((1<<PSE_AW)-1)] <= cbd_modq;
                    if (cbd_last) begin
                        c_poly <= c_poly + 3'd1;
                        c_widx <= 8'd0;
                        c_busy <= 1'b0;
                        if (c_poly + 3'd1 < {1'b0, k_r, 1'b0}) cbd_valid <= 1'b1;
                    end else begin
                        c_widx <= c_widx + 8'd1;
                    end
                end
            end

            // Arm N stage when C finishes: prime load of slot S0. n_ridx is a
            // read-ahead pointer; bank_se read is registered into n_rd_data and
            // fed to ntt_core one cycle later, so valid starts low (priming).
            if (st == ST_C && st_next == ST_N) begin
                n_slot    <= 3'd0;
                n_ridx    <= 9'd0;
                n_widx    <= 8'd0;
                n_valid   <= 1'b0;
                n_loading <= 1'b1;       // begin presenting load addresses
                n_pending <= 1'b0;
            end

            // ---- ST_N: forward NTT each of S0,S1,E0,E1 in place ----
            if (st == ST_N) begin
                // LOAD phase: present read-ahead addr to bank_se; the value
                // registered last cycle (n_rd_data) is consumed by ntt_core
                // this cycle (n_valid). Cores hold ready high through LOAD, so
                // a fixed 1-cycle skew suffices (no backpressure gating).
                if (n_loading) begin
                    if (n_ridx == 9'd256) begin
                        // 256th coeff (bank_se[255]) consumed this cycle; stop
                        n_loading <= 1'b0;
                        n_valid   <= 1'b0;
                    end else begin
                        n_rd_data <= bank_se[ntt_rd_full[PSE_AW-1:0]];
                        n_ridx    <= n_ridx + 9'd1;
                        n_valid   <= 1'b1;     // data presented last cycle is valid
                    end
                end

                // OUTPUT phase: collect 256 results, write back to same slot
                if (ntt_vo) begin
                    // bank_se relative slot = n_slot_addr - slot_s_rt = n_slot
                    bank_se[(n_slot*256 + n_widx) & ((1<<PSE_AW)-1)] <= ntt_coeff;
                    n_widx <= n_widx + 8'd1;   // wraps 255->0 after last
                end

                // Slot complete when ntt_core returns to DONE
                if (ntt_done) begin
                    if (n_slot + 3'd1 < {1'b0, k_r, 1'b0}) begin
                        n_slot    <= n_slot + 3'd1;
                        n_widx    <= 8'd0;
                        n_pending <= 1'b1;     // wait one cycle for core IDLE
                    end else begin
                        n_slot <= n_slot + 3'd1;  // == 2K -> ST_DONE
                    end
                end

                // Kick next slot's load once core is back IDLE (re-prime)
                if (n_pending && ntt_ready && !ntt_done) begin
                    n_ridx    <= 9'd0;
                    n_valid   <= 1'b0;
                    n_loading <= 1'b1;
                    n_pending <= 1'b0;
                end
            end

            // Arm M stage when N finishes: prime first (i=0,j=0) poly_mul load.
            // m_ld read-ahead pointer; pm_a_rd/pm_b_rd registered, pm_valid
            // asserted one cycle after an address is presented (like ST_N).
            if (st == ST_N && st_next == ST_M) begin
                m_i       <= 2'd0;
                m_j       <= 2'd0;
                m_ld      <= 9'd0;
                m_oidx    <= 8'd0;
                m_loading <= 1'b1;
                m_pending <= 1'b0;
                pm_valid  <= 1'b0;
            end

            // ---- ST_M: t_hat[i] = e_hat[i] + sum_j A[i][j] o s_hat[j] ----
            if (st == ST_M) begin
                // registered accumulator read (sd_bram timing): captures the
                // index the next pm_vo will consume. Free-running within ST_M;
                // primed during LOAD (m_oidx held at 0).
                m_acc_rd <= m_acc_rd_next;

                // LOAD: present read-ahead addr to bank_a/bank_se; the pair
                // registered last cycle (pm_a_rd/pm_b_rd) is consumed by
                // poly_mul this cycle (pm_valid). poly_mul holds ready high
                // through LOAD, so a fixed 1-cycle skew suffices.
                if (m_loading) begin
                    if (m_ld == 9'd256) begin
                        pm_valid  <= 1'b0;     // 256th pair consumed this cycle
                        m_loading <= 1'b0;
                        m_ld      <= 9'd0;
                        m_oidx    <= 8'd0;
                    end else begin
                        pm_a_rd <= bank_a [pm_a_full[PA_AW-1:0]];
                        pm_b_rd <= bank_se[pm_b_full[PSE_AW-1:0]];
                        m_ld    <= m_ld + 9'd1;
                        pm_valid <= 1'b1;      // pair presented last cycle is valid
                    end
                end

                // ACCUMULATE: each product coeff += e_hat (j==0) or running t_hat
                if (pm_vo) begin
                    // bank_t relative slot = m_tslot - slot_t_rt = m_i
                    bank_t[(m_i*256 + m_oidx) & ((1<<PT_AW)-1)] <= m_accq;
                    if (m_oidx == 8'd255) begin
                        // finished this (i,j) term; advance
                        if (m_j + 2'd1 < k_r) begin
                            m_j       <= m_j + 2'd1;
                            m_pending <= 1'b1;     // next term, same row
                        end else begin
                            m_j <= 2'd0;
                            m_i <= m_i + 2'd1;     // next row (or == K -> DONE)
                            if (m_i + 2'd1 < k_r) m_pending <= 1'b1;
                        end
                    end else begin
                        m_oidx <= m_oidx + 8'd1;
                    end
                end

                // Start next (i,j) poly_mul load once core is IDLE again (re-prime)
                if (m_pending && pm_ready && !pm_vo) begin
                    m_loading <= 1'b1;
                    m_ld      <= 9'd0;
                    m_oidx    <= 8'd0;
                    pm_valid  <= 1'b0;
                    m_pending <= 1'b0;
                end
            end

            // Arm E stage when M finishes
            if (st == ST_M && st_next == ST_E) begin
                e_poly <= 3'd0;
                e_pair <= 8'd0;
                e_byte <= 2'd0;
                e_rho  <= 10'd0;
                e_done <= 1'b0;
            end

            // ---- ST_E: byteEncode12 t_hat -> ek, s_hat -> dk_pke, 1 byte/cycle ----
            if (st == ST_E && !e_done) begin
                if (e_poly < {1'b0, k_r, 1'b0}) begin
                    // write current byte (e_byte_d) to target BRAM
                    if (!e_is_dk) begin
                        ek_we <= 1'b1;
                        ek_wa <= e_boff[10:0] + {9'd0, e_byte};
                        ek_wd <= e_byte_d;
                    end else begin
                        dkp_we <= 1'b1;
                        dkp_wa <= e_boff[10:0] + {9'd0, e_byte};
                        dkp_wd <= e_byte_d;
                    end
                    // advance sub-counters
                    if (e_byte == 2'd2) begin
                        e_byte <= 2'd0;
                        if (e_pair == 8'd127) begin
                            e_pair <= 8'd0;
                            e_poly <= e_poly + 5'd1;   // next poly or -> rho phase
                        end else begin
                            e_pair <= e_pair + 8'd1;
                        end
                    end else begin
                        e_byte <= e_byte + 2'd1;
                    end
                end else begin
                    // rho copy: ek[384*K + r] = rho byte r  (r = 0..31), 1 byte/cycle
                    ek_we <= 1'b1;
                    ek_wa <= dk_bytes_rt[10:0] + {1'b0, e_rho};
                    ek_wd <= rho_r[e_rho*8 +: 8];
                    if (e_rho == 10'd31) e_done <= 1'b1;
                    else                 e_rho  <= e_rho + 10'd1;
                end
            end

            // Arm H stage when E finishes
            if (st == ST_E && st_next == ST_H) begin
                h_blk    <= 3'd0;
                h_byte   <= 8'd0;
                h_phase  <= 2'd0;     // assemble
                h_mbvalid<= 1'b0;
                h_mblast <= 1'b0;
                h_ack    <= 1'b1;     // ready to consume final digest
                h_wb_vld <= 1'b0;     // no pending writeback yet
            end

            // ---- ST_H: H(ek) via multi-block SHA3-256 (6/9/12 pre-padded blocks) ----
            if (st == ST_H) begin
                case (h_phase)
                    // assemble 136 bytes of block h_blk into h_block_r.
                    // ek BRAM read is registered: present addr for h_byte this
                    // cycle, write back the byte that arrived for the addr we
                    // presented last cycle (h_wb_*). h_byte runs 0..136 (one
                    // extra cycle to flush the final writeback).
                    2'd0: begin
                        // writeback the byte read for the previous address
                        if (h_wb_vld)
                            h_block_r[h_wb_idx*8 +: 8] <= h_wb_inek ? ek_rd_data : h_wb_pad;

                        if (h_byte <= 8'd135) begin
                            // set up writeback for the address presented this cycle
                            h_wb_vld  <= 1'b1;
                            h_wb_idx  <= h_byte;
                            h_wb_g    <= h_g_addr;
                            h_wb_inek <= (h_g_addr < ek_bytes_rt);
                            h_wb_pad  <= h_padconst(h_blk, h_byte);
                            h_byte    <= h_byte + 8'd1;
                        end else begin
                            // h_byte==136: final writeback (for byte 135) done above
                            h_wb_vld  <= 1'b0;
                            h_byte    <= 8'd0;
                            h_mbvalid <= 1'b1;
                            h_mblast  <= (h_blk == h_nblk_rt - 4'd1);
                            h_phase   <= 2'd1;        // feed
                        end
                    end
                    // feed: hold valid until accepted (mb_ready drops)
                    2'd1: begin
                        if (h_mbvalid && !h_mbready) begin
                            h_mbvalid <= 1'b0;
                            h_mblast  <= 1'b0;
                            h_phase   <= 2'd2;        // wait permute
                        end
                    end
                    // wait permute done: ready again (more blocks) or digest valid (last)
                    2'd2: begin
                        if (h_vo) begin
                            hek_r   <= h_hash[255:0];
                            h_phase <= 2'd3;          // done
                        end else if (h_mbready) begin
                            h_blk   <= h_blk + 3'd1;
                            h_phase <= 2'd0;          // assemble next block
                        end
                    end
                    default: ; // 2'd3 done: hold
                endcase
            end
        end
    end

endmodule
