// sample_ntt_sync_shared.v - Synchronous SampleNTT for ML-KEM A matrix generation
//                            with external Keccak interface ports.
//
// Identical to sample_ntt_sync.v except the internal keccak_core instance
// is replaced by explicit kc_* ports. This allows multiple SampleNTT
// instances to share a single keccak_core (via a round-robin arbiter).
//
// Generates one k×k polynomial (256 coefficients) via SHAKE-128 XOF
// rejection sampling from seed rho || j || i.
//
// FIPS 202/203 conformant SHAKE-128 squeeze:
//   - Absorb: S = Keccak-p(padded(rho || j || i))
//   - Squeeze the full 1344-bit (168-byte) rate as 56 groups of 3 bytes,
//     each group g read from S[24*g +: 24]; extract d1[11:0], d2[23:12]
//   - Accept d if d < Q=3329
//   - Only after all 56 groups of the block are consumed: S = Keccak-p(S)
//     (re-permute once per rate block, NOT per 3-byte group)
//   - Repeat until 256 coefficients collected (~3 blocks)
//
// Parameters:
//   K = 4 (ML-KEM parameter)
//
// Interface:
//   clk, rst_n    - clock, active-low reset
//   rho_i[255:0]  - 256-bit seed rho (32 bytes)
//   k_i[2:0]      - actual k value (2/3/4)
//   i_idx[1:0]    - row index (0..k-1)
//   j_idx[1:0]    - column index (0..k-1)
//   valid_i        - start generation for this (i,j) pair
//   ready_o        - module can accept request
//   coeff_o[11:0]  - one coefficient output per cycle
//   valid_o        - coefficient output valid
//   ready_i        - consumer accepts coefficient
//   last_o         - high when 256th coefficient is output
//
//   --- Keccak interface (external, shared) ---
//   kc_state_o[1599:0] - keccak result (from keccak_core)
//   kc_valid_o         - keccak output valid (from keccak_core)
//   kc_ready_i         - sample_ntt ready to accept keccak output (always 1'b1)
//   kc_state_i[1599:0] - keccak input state (to keccak_core)
//   kc_valid_i         - request keccak permutation (to keccak_core)
//   kc_ready_o         - keccak ready to accept new request (from keccak_core)

`include "sync_rtl/common/defines.vh"

/* verilator lint_off UNUSEDPARAM */
module sample_ntt_sync_shared #(parameter K = 4) (
/* verilator lint_on UNUSEDPARAM */
    input              clk,
    input              rst_n,
    input  [255:0]     rho_i,
    /* verilator lint_off UNUSEDSIGNAL */
    input  [2:0]       k_i,
    /* verilator lint_on UNUSEDSIGNAL */
    input  [1:0]       i_idx,
    input  [1:0]       j_idx,
    input              valid_i,
    output             ready_o,
    output [11:0]      coeff_o,
    output             valid_o,
    input              ready_i,
    output             last_o,

    // --- Keccak interface ports (external) ---
    input  [1599:0]    kc_state_o,
    input              kc_valid_o,
    output             kc_ready_i,
    output [1599:0]    kc_state_i,
    output             kc_valid_i,
    /* verilator lint_off UNUSEDSIGNAL */
    input              kc_ready_o
    /* verilator lint_on UNUSEDSIGNAL */
);

    // ================================================================
    // Local parameters
    // ================================================================
    localparam Q = `Q;   // 3329

    // SHAKE-128 rate = 1344 bits = 168 bytes = 56 groups of 3 bytes.
    // After consuming all 56 groups of a block, re-permute the state.
    localparam GRP_MAX = 6'd55;

    // ================================================================
    // FSM state encoding
    // ================================================================
    localparam ST_IDLE       = 3'd0;
    localparam ST_ABSORB     = 3'd1;   // keccak running on absorb_state
    localparam ST_SQUEEZE    = 3'd2;   // extract d1/d2, output coefficients
    localparam ST_WAIT       = 3'd3;   // wait for next keccak to finish
    localparam ST_DONE       = 3'd4;   // all 256 coefficients output

    reg [2:0]  state_r, state_next;

    // ================================================================
    // Registered inputs (captured on valid_i)
    // ================================================================
    /* verilator lint_off UNUSEDSIGNAL */
    reg [255:0] rho_r;
    reg [2:0]   k_r;
    reg [1:0]   i_r;
    reg [1:0]   j_r;
    /* verilator lint_on UNUSEDSIGNAL */

    // ================================================================
    // Coefficient counter (0..256). 9 bits to avoid overflow when
    // reaching 256 after the 256th output.
    // ================================================================
    reg [8:0] coeff_cnt_r;

    // ================================================================
    // Squeeze state register (1600-bit result of keccak_p)
    // ================================================================
    reg [1599:0] squeeze_state_r;

    // ================================================================
    // Group pointer: which 3-byte group within the 1344-bit rate
    // block is currently being consumed (0..GRP_MAX).
    // ================================================================
    reg [5:0] grp_ptr_r;

    // ================================================================
    // Registered d1, d2 and acceptance flags
    // ================================================================
    reg [11:0] d1_r, d2_r;
    reg        d1_acc_r, d2_acc_r;   // d1/d2 accepted?

    // ================================================================
    // Squeeze sub-phase (for multi-cycle coefficient output)
    //   0 → latch d1/d2, start keccak → 1
    //   1 → output d1 if accepted → 2
    //   2 → output d2 if accepted → WAIT
    // ================================================================
    reg [1:0] sq_phase_r;

    // ================================================================
    // Comb: build absorb state from rho, i, j INPUT PORTS directly
    // (not registered copies — avoids NBA race on IDLE→ABSORB edge)
    // ================================================================
    wire [7:0] abs_i_byte, abs_j_byte;
    assign abs_i_byte = {6'b0, i_idx};
    assign abs_j_byte = {6'b0, j_idx};

    wire [271:0] msg_bytes;
    assign msg_bytes = {abs_i_byte, abs_j_byte, rho_i};

    // SHAKE-128 absorb block: capacity(256b) | pad10*1 | suffix(1111) | msg(272b)
    wire [1599:0] absorb_state;
    assign absorb_state = {
        256'b0,                // capacity          [1599:1344]
        1'b1,                  // pad10*1 final 1   [1343]
        {1066{1'b0}},         // pad10*1 zeros     [1342:277]
        1'b1,                  // pad10*1 first 1   [276]
        4'b1111,               // SHAKE suffix      [275:272]
        msg_bytes              // message           [271:0]
    };

    // ================================================================
    // Comb: extract d1,d2 from the current 3-byte group of squeeze state
    // ================================================================
    // group g occupies bits [24*g +: 24]: c0=byte0, c1=byte1, c2=byte2
    // d1 = {c1[3:0], c0}
    // d2 = {c2, c1[7:4]}
    wire [10:0] grp_bit_off;
    assign grp_bit_off = grp_ptr_r * 11'd24;   // 0..1320, +24 ≤ 1344 (rate)

    wire [23:0] grp_bits;
    assign grp_bits = squeeze_state_r[ grp_bit_off +: 24 ];

    wire [7:0] c0, c1, c2;
    assign c0 = grp_bits[7:0];
    assign c1 = grp_bits[15:8];
    assign c2 = grp_bits[23:16];

    wire [11:0] d1_comb, d2_comb;
    assign d1_comb = {c1[3:0], c0};
    assign d2_comb = {c2, c1[7:4]};

    wire d1_ok, d2_ok;
    assign d1_ok = (d1_comb < Q);
    assign d2_ok = (d2_comb < Q);

    // ================================================================
    // Keccak interface: external ports instead of internal keccak_core.
    //
    // kc_valid_i: asserted during ABSORB and first phase of SQUEEZE.
    // kc_state_i: absorb_state in ABSORB, squeeze_state_r otherwise.
    // kc_ready_i: always 1'b1 — this module always accepts keccak output.
    //
    // These signals are output ports driven by the combinational logic
    // below. The upstream arbiter connects them to a shared keccak_core.
    // ================================================================

    // kc_ready_i: always ready to accept keccak output
    assign kc_ready_i = 1'b1;

    // kc_valid_i: asserted on the ABSORB load, and for one cycle when the
    // squeeze block is exhausted (SQUEEZE → WAIT) to re-permute the state.
    assign kc_valid_i = (state_next == ST_ABSORB) ||
                        (state_r == ST_SQUEEZE && state_next == ST_WAIT);

    // kc_state_i: absorb_state in ABSORB, squeeze_state_r otherwise
    assign kc_state_i = (state_next == ST_ABSORB) ? absorb_state : squeeze_state_r;

    // ================================================================
    // Output signals
    // ================================================================
    reg [11:0] coeff_o_r;
    reg        valid_o_r;
    reg        last_o_r;

    assign coeff_o = coeff_o_r;
    assign valid_o = valid_o_r;
    assign last_o  = last_o_r;

    // ================================================================
    // ready_o: accept new request in IDLE
    // ================================================================
    assign ready_o  = (state_r == ST_IDLE);
    wire need_more  = (coeff_cnt_r < 9'd256);

    // grp_done: current 3-byte group fully consumed (phase 2 resolved):
    // d2 rejected, no longer need coefficients, or d2 was just accepted out.
    wire grp_done = (state_r == ST_SQUEEZE) && (sq_phase_r == 2'd2) &&
                    (!d2_acc_r || !need_more || (valid_o_r && ready_i));

    // ================================================================
    // FSM: state_next (combinational)
    // ================================================================
    always @(*) begin
        state_next = state_r;

        case (state_r)
            ST_IDLE: begin
                if (valid_i && ready_o)
                    state_next = ST_ABSORB;
            end

            ST_ABSORB: begin
                // Wait for keccak to finish the absorb permutation
                if (kc_valid_o)
                    state_next = ST_SQUEEZE;
            end

            ST_SQUEEZE: begin
                // A group is fully consumed once phase 2 resolves (d2 output,
                // rejected, or no longer needed). Then either advance to the
                // next group in this block, re-permute (block exhausted), or
                // finish.
                if (grp_done) begin
                    if (!need_more)
                        state_next = ST_DONE;
                    else if (grp_ptr_r < GRP_MAX)
                        state_next = ST_SQUEEZE;  // next group, no re-permute
                    else
                        state_next = ST_WAIT;     // block exhausted: re-permute
                end
            end

            ST_WAIT: begin
                // Wait for keccak to finish squeeze permutation
                if (kc_valid_o) begin
                    if (!need_more)
                        state_next = ST_DONE;
                    else
                        state_next = ST_SQUEEZE;
                end
            end

            ST_DONE: begin
                state_next = ST_IDLE;
            end

            default: state_next = ST_IDLE;
        endcase
    end

    // ================================================================
    // FSM: sequential logic (registered)
    // ================================================================
    always @(posedge clk or negedge rst_n) begin
        if (!rst_n) begin
            state_r         <= ST_IDLE;
            sq_phase_r      <= 2'd0;
            coeff_cnt_r     <= 9'd0;
            squeeze_state_r <= 1600'd0;
            grp_ptr_r       <= 6'd0;
            d1_r            <= 12'd0;
            d2_r            <= 12'd0;
            d1_acc_r        <= 1'b0;
            d2_acc_r        <= 1'b0;
            coeff_o_r       <= 12'd0;
            valid_o_r       <= 1'b0;
            last_o_r        <= 1'b0;
            rho_r           <= 256'd0;
            k_r             <= 3'd0;
            i_r             <= 2'd0;
            j_r             <= 2'd0;
        end else begin
            state_r <= state_next;

            // ---------------------------------------------------------
            // Capture inputs on valid_i (IDLE → ABSORB transition)
            // ---------------------------------------------------------
            if (state_r == ST_IDLE && valid_i && ready_o) begin
                rho_r        <= rho_i;
                k_r          <= k_i;
                i_r          <= i_idx;
                j_r          <= j_idx;
                coeff_cnt_r  <= 9'd0;
            end

            // ---------------------------------------------------------
            // Latch keccak output when valid_o fires. A fresh block
            // starts at group 0 (ABSORB load or WAIT re-permute result).
            // ---------------------------------------------------------
            if (kc_valid_o) begin
                squeeze_state_r <= kc_state_o;
                grp_ptr_r       <= 6'd0;
            end

            // ---------------------------------------------------------
            // SQ_PHASE = 0: latch d1/d2 acceptance
            // ---------------------------------------------------------
            if (state_r == ST_SQUEEZE && sq_phase_r == 2'd0) begin
                d1_r     <= d1_comb;
                d2_r     <= d2_comb;
                d1_acc_r <= d1_ok;
                d2_acc_r <= d2_ok;
            end

            // ---------------------------------------------------------
            // SQ_PHASE management and coefficient output
            //
            // Handshake pattern (matching pipeline_reg convention):
            //   Cycle N:   valid_o_r ← 1, coeff_o_r ← value
            //   Cycle N+1: if ready_i: consume, valid_o_r ← 0, advance
            //
            // The testbench reads coeff_o after Cycle N and consumes
            // with the next posedge (Cycle N+1).
            // ---------------------------------------------------------
            if (state_r == ST_SQUEEZE) begin
                case (sq_phase_r)
                    // ----- Phase 0: latch, advance to phase 1 -----
                    2'd0: begin
                        sq_phase_r <= 2'd1;
                    end

                    // ----- Phase 1: output d1 -----
                    2'd1: begin
                        if (d1_acc_r && need_more) begin
                            if (!valid_o_r) begin
                                // First cycle: assert output
                                coeff_o_r <= d1_r;
                                valid_o_r <= 1'b1;
                                last_o_r  <= (coeff_cnt_r == 9'd255);
                            end else begin
                                // valid_o is high; wait for consumer
                                if (ready_i) begin
                                    coeff_cnt_r <= coeff_cnt_r + 9'd1;
                                    valid_o_r   <= 1'b0;
                                    sq_phase_r  <= 2'd2;
                                end
                            end
                        end else begin
                            // d1 rejected or no longer needed: skip to phase 2
                            valid_o_r  <= 1'b0;
                            sq_phase_r <= 2'd2;
                        end
                    end

                    // ----- Phase 2: output d2 -----
                    2'd2: begin
                        if (d2_acc_r && need_more) begin
                            if (!valid_o_r) begin
                                // First cycle: assert output
                                coeff_o_r <= d2_r;
                                valid_o_r <= 1'b1;
                                last_o_r  <= (coeff_cnt_r == 9'd255);
                            end else begin
                                // valid_o is high; wait for consumer
                                if (ready_i) begin
                                    coeff_cnt_r <= coeff_cnt_r + 9'd1;
                                    valid_o_r   <= 1'b0;
                                    // state_next → ST_WAIT (combinational)
                                end
                            end
                        end else begin
                            // d2 rejected or done
                            valid_o_r <= 1'b0;
                        end
                    end

                    default: begin
                        valid_o_r <= 1'b0;
                    end
                endcase

                // Group consumed but block not exhausted: advance to the
                // next 3-byte group within the same block (no re-permute).
                if (grp_done && need_more && grp_ptr_r < GRP_MAX) begin
                    grp_ptr_r  <= grp_ptr_r + 6'd1;
                    sq_phase_r <= 2'd0;
                end
            end else if (state_r != ST_SQUEEZE && state_next == ST_SQUEEZE) begin
                // About to enter SQUEEZE: reset phase and output
                sq_phase_r <= 2'd0;
                valid_o_r  <= 1'b0;
            end else begin
                // Not in SQUEEZE: clear
                valid_o_r  <= 1'b0;
                sq_phase_r <= 2'd0;
            end
        end
    end

endmodule
