// sample_cbd_sync_shared.v - Centered Binomial Distribution sampling via SHAKE-256 PRF
// (shared keccak_core variant — external instance)
//
// Generates 256 polynomial coefficients from a 256-bit seed and 8-bit nonce
// using SHAKE-256 PRF followed by Centered Binomial Distribution (CBD).
//
// Algorithm (FIPS 203 / ML-KEM):
//   PRF(sigma, N) = SHAKE-256(sigma || N) → squeeze eta*64 bytes
//   For each of 256 coefficients:
//     eta=2: read 4 bits, coeff = (b0+b1) - (b2+b3)
//     eta=3: read 6 bits, coeff = (b0+b1+b2) - (b3+b4+b5)
//   Each coefficient in range [-eta, eta], stored as 12-bit signed.
//
// SHAKE-256 parameters:
//   rate = 1088 bits, capacity = 512 bits
//   suffix = 4'b1111
//   pad10*1 padding
//
// Multi-squeeze (eta=3):
//   SHAKE-256 squeezes 1088-bit blocks. For eta=3 we need 1536 bits.
//   First squeeze provides bits [0:1087], second provides [1088:1535].
//   The 1536-bit squeeze_buf accumulates both blocks contiguously:
//     After 1st: buf[1087:0] = squeeze1, buf[1535:1088] = 0
//     After 2nd: buf[1535:1088] = squeeze2[447:0] (remapped)
//   Bit ordering matches Python reference PRF output.
//
// Interface:
//   clk, rst_n     - clock, active-low reset
//   seed_i [255:0]  - sigma (256-bit seed)
//   nonce_i [7:0]   - N counter byte
//   eta_i [1:0]     - 2'd2 or 2'd3
//   valid_i         - input valid (start sampling)
//   ready_o         - module can accept new input
//   coeff_o [11:0]  - one coefficient per cycle, 12-bit signed
//   valid_o         - output valid
//   ready_i         - downstream accepts output
//   last_o          - high when last coefficient (255th, 0-indexed)
//   kc_state_o      - keccak result (from external keccak_core.state_o)
//   kc_valid_o      - keccak done (from external keccak_core.valid_o)
//   kc_ready_i      - always ready to accept keccak output (to keccak_core.ready_i)
//   kc_state_i      - keccak input state (to external keccak_core.state_i)
//   kc_valid_i      - request keccak permutation (to external keccak_core.valid_i)
//   kc_ready_o      - keccak ready to accept (from external keccak_core.ready_o)

module sample_cbd_sync_shared (
    input              clk,
    input              rst_n,
    input  [255:0]     seed_i,
    input  [7:0]       nonce_i,
    input  [1:0]       eta_i,          // 2'd2 or 2'd3
    input              valid_i,
    output             ready_o,
    output [11:0]      coeff_o,        // 12-bit signed
    output             valid_o,
    input              ready_i,
    output             last_o,

    // External keccak_core interface
    input  [1599:0]    kc_state_o,     // keccak result (from keccak_core.state_o)
    input              kc_valid_o,     // keccak done (from keccak_core.valid_o)
    output             kc_ready_i,     // always ready to accept (to keccak_core.ready_i)
    output [1599:0]    kc_state_i,     // keccak input state (to keccak_core.state_i)
    output             kc_valid_i,     // request keccak permutation (to keccak_core.valid_i)
    input              kc_ready_o      // keccak ready to accept (from keccak_core.ready_o)
);

    // ================================================================
    // FSM state encoding
    // ================================================================
    localparam ST_IDLE    = 2'd0;
    localparam ST_PERMUTE = 2'd1;
    localparam ST_SQUEEZE = 2'd2;

    reg [1:0]  state_r, state_next;

    // ================================================================
    // SHAKE-256 pad10*1 construction (combinational)
    //
    // Message: {nonce_i, seed_i} = 264 bits
    // In FIPS 202 bit ordering: message[0] = seed_i[0], message[263] = nonce_i[7]
    // Padded block (1088 rate bits):
    //   {1'b1, 818'b0, 1'b1, 4'b1111, nonce_i, seed_i}
    //
    // Matches Python: PRF(sigma, N) = shake256(sigma||N, d)
    // where sigma bits come first (LSB), then N bits, then suffix+padding.
    // ================================================================
    wire [263:0] message_264;
    wire [1087:0] padded_block;

    assign message_264 = {nonce_i, seed_i};
    assign padded_block = {1'b1, {818{1'b0}}, 1'b1, 4'b1111, message_264};

    // Absorb state: capacity (512 bits of zero) || padded rate block
    wire [1599:0] absorb_state;
    assign absorb_state = {{(1600-1088){1'b0}}, padded_block};

    // ================================================================
    // Registered inputs (captured on valid_i && ready_o)
    // ================================================================
    reg [2:0]   eta_r;            // 2 or 3
    reg [7:0]   coeff_cnt;        // 0..255, number of coeffs output so far
    reg [1599:0] keccak_state_r;  // current 1600-bit keccak state (for re-permutation)
    reg         perm_done;        // 1 after first keccak permutation completes

    // Combinational mux select: absorb_state BEFORE first perm finishes,
    // keccak_state_r AFTER. Must be combinational to avoid NBA race
    // with kc_valid_i on the IDLE→PERMUTE transition edge.
    wire        first_perm_sel;
    assign      first_perm_sel = !perm_done;

    // ================================================================
    // Squeeze buffer (accumulated across multiple squeezes)
    // ================================================================
    reg [1535:0] squeeze_buf;   // accumulated squeeze data
    reg [10:0]   buf_fill;      // valid bits in squeeze_buf (0, 1088, or 1536)
    reg [10:0]   buf_ptr;       // read position within squeeze_buf
    reg          kc_done_d1;    // kc_valid_o delayed by 1 cycle (gate for valid_o)

    // ================================================================
    // External keccak_core interface connections
    //
    // The external keccak_core is instantiated at the top level.
    // This module drives kc_valid_i and kc_state_i, receives
    // kc_valid_o and kc_state_o, and provides kc_ready_i = 1'b1
    // (always ready to accept keccak output).  kc_ready_o is exposed
    // for the top-level arbiter to check keccak availability.
    // ================================================================

    // Mux: absorb_state for first perm, keccak_state_r for re-permutation
    assign kc_state_i = first_perm_sel ? absorb_state : keccak_state_r;

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

    // ================================================================
    // Bits per coefficient
    // ================================================================
    wire [3:0] bits_per_coeff;
    assign bits_per_coeff = (eta_r == 3'd2) ? 4'd4 : 4'd6;

    // ================================================================
    // CBD computation (combinational)
    //
    // Extract bits_per_coeff bits from squeeze_buf starting at buf_ptr.
    // For eta=2: b0,b1,b2,b3 → coeff = (b0+b1) - (b2+b3)
    // For eta=3: b0,b1,b2,b3,b4,b5 → coeff = (b0+b1+b2) - (b3+b4+b5)
    // ================================================================
    wire [5:0] cbd_bits;
    assign cbd_bits = squeeze_buf[buf_ptr +: 6];

    wire [2:0] sum_pos, sum_neg;

    // eta=2: sum_pos = b0+b1, sum_neg = b2+b3
    wire [2:0] sp2, sn2;
    assign sp2 = {2'b00, cbd_bits[0]} + {2'b00, cbd_bits[1]};
    assign sn2 = {2'b00, cbd_bits[2]} + {2'b00, cbd_bits[3]};

    // eta=3: sum_pos = b0+b1+b2, sum_neg = b3+b4+b5
    wire [2:0] sp3, sn3;
    assign sp3 = {2'b00, cbd_bits[0]} + {2'b00, cbd_bits[1]} + {2'b00, cbd_bits[2]};
    assign sn3 = {2'b00, cbd_bits[3]} + {2'b00, cbd_bits[4]} + {2'b00, cbd_bits[5]};

    assign sum_pos = (eta_r == 3'd2) ? sp2 : sp3;
    assign sum_neg = (eta_r == 3'd2) ? sn2 : sn3;

    wire signed [3:0] coeff_raw;
    assign coeff_raw = $signed({1'b0, sum_pos}) - $signed({1'b0, sum_neg});

    wire [11:0] coeff_signed;
    assign coeff_signed = {{8{coeff_raw[3]}}, coeff_raw[3:0]};  // sign-extend to 12 bits

    assign coeff_o = coeff_signed;
    // valid_o: suppress for 1 cycle after keccak finishes to avoid
    // buf_ptr race between squeeze capture and SQUEEZE advancement.
    assign valid_o = (state_r == ST_SQUEEZE) && (buf_fill >= {7'b0, bits_per_coeff}) && !kc_done_d1;
    assign last_o  = valid_o && (coeff_cnt == 8'd255);

    // ================================================================
    // FSM: ready_o
    // ================================================================
    assign ready_o = (state_r == ST_IDLE);

    // ================================================================
    // kc_valid_i: start external keccak_core when transitioning to PERMUTE
    // ================================================================
    assign kc_valid_i = (state_next == ST_PERMUTE) && (state_r != ST_PERMUTE);

    // ================================================================
    // Buffer exhaustion detection
    // ================================================================
    wire [11:0] next_buf_ptr;
    assign next_buf_ptr = {1'b0, buf_ptr} + {8'b0, bits_per_coeff};

    // buffer will be exhausted after outputting NEXT coefficient
    wire buffer_exhaust_next;
    assign buffer_exhaust_next = (next_buf_ptr + {8'b0, bits_per_coeff}) > {1'b0, buf_fill};

    // Coefficients remaining AFTER the current one
    wire [8:0] coeffs_remaining;
    assign coeffs_remaining = 9'd256 - {1'b0, coeff_cnt} - 9'd1;

    // ================================================================
    // FSM combinational next-state logic
    // ================================================================
    always @(*) begin
        state_next = state_r;
        case (state_r)
            ST_IDLE: begin
                if (valid_i && ready_o)
                    state_next = ST_PERMUTE;
            end

            ST_PERMUTE: begin
                // Wait for external keccak_core to finish
                if (kc_valid_o)
                    state_next = ST_SQUEEZE;
            end

            ST_SQUEEZE: begin
                // Output one coeff per cycle when ready
                if (valid_o && ready_i) begin
                    if (coeff_cnt == 8'd255) begin
                        // Last coefficient was output
                        state_next = ST_IDLE;
                    end else if (buffer_exhaust_next && (coeffs_remaining > 9'd0)) begin
                        // Need more squeeze data: start another keccak permutation
                        state_next = ST_PERMUTE;
                    end
                    // else: stay in SQUEEZE for next coefficient
                end
            end

            default: state_next = ST_IDLE;
        endcase
    end

    // ================================================================
    // Sequential logic
    // ================================================================
    always @(posedge clk or negedge rst_n) begin
        if (!rst_n) begin
            state_r           <= ST_IDLE;
            eta_r             <= 3'd0;
            coeff_cnt         <= 8'd0;
            keccak_state_r    <= 1600'd0;
            perm_done         <= 1'b0;
            squeeze_buf       <= 1536'd0;
            buf_fill          <= 11'd0;
            buf_ptr           <= 11'd0;
            kc_done_d1        <= 1'b0;
        end else begin
            state_r <= state_next;

            // Delay kc_valid_o by 1 cycle to gate valid_o
            kc_done_d1 <= kc_valid_o;

            // ---- Capture inputs on IDLE → PERMUTE transition ----
            if (state_r == ST_IDLE && valid_i && ready_o) begin
                // Determine eta: 2 or 3
                if (eta_i == 2'd2)
                    eta_r <= 3'd2;
                else if (eta_i == 2'd3)
                    eta_r <= 3'd3;
                else
                    eta_r <= 3'd2;  // default

                coeff_cnt      <= 8'd0;
                buf_fill       <= 11'd0;
                buf_ptr        <= 11'd0;
            end

            // ---- On keccak_core valid_o: latch squeeze data ----
            if (kc_valid_o) begin
                // Save full 1600-bit state for potential re-permutation
                keccak_state_r <= kc_state_o;

                if (!perm_done) begin
                    // First squeeze: fill lower 1088 bits of squeeze_buf
                    // buf[1087:0] = squeeze data, buf[1535:1088] = 0
                    squeeze_buf <= {{(1536 - 1088){1'b0}}, kc_state_o[1087:0]};
                    buf_fill    <= 11'd1088;
                    buf_ptr     <= 11'd0;
                    perm_done   <= 1'b1;
                end else begin
                    // Second squeeze: fill upper 448 bits of squeeze_buf
                    // buf[1535:1088] = squeeze2[447:0], buf[1087:0] preserved
                    // This remaps: squeeze2[0] → buf[1088], matching Python's contiguous output
                    squeeze_buf <= {kc_state_o[447:0], squeeze_buf[1087:0]};
                    buf_fill    <= 11'd1536;
                    // buf_ptr stays at current position (kept from SQUEEZE)
                end
            end

            // ---- SQUEEZE: advance on output ----
            if (state_r == ST_SQUEEZE && valid_o && ready_i) begin
                buf_ptr    <= buf_ptr + {7'b0, bits_per_coeff};
                coeff_cnt  <= coeff_cnt + 8'd1;

                // Clear state on last coeff
                if (coeff_cnt == 8'd255) begin
                    buf_fill  <= 11'd0;
                    buf_ptr   <= 11'd0;
                    // Reset perm_done here (before next IDLE→PERMUTE)
                    // so the mux selects absorb_state for the next vector.
                    perm_done <= 1'b0;
                end
            end
        end
    end

endmodule
