// sha3_top.v - SHA3/SHAKE top wrapper with valid/ready interface
//
// Implements SHA3-512 (G), SHA3-256 (H), SHAKE-256 (J) over a single
// Keccak-f[1600] core. Supports single-block absorption, plus a separate
// multi-block SHA3-256 streaming absorb path for long inputs (e.g. H(ek)).
//
// Single-block modes (mb_en=0):
//   00 = G  (SHA3-512):   rate=576,  suffix=01, msg_len=264, out=512
//   01 = H  (SHA3-256):   rate=1088, suffix=01, msg_len=256, out=256
//   10 = J  (SHAKE-256):  rate=1088, suffix=1111,msg_len=512,out=256
//
// Multi-block SHA3-256 (mb_en=1):
//   Streams pre-padded 1088-bit rate blocks. The CALLER applies SHA3-256
//   padding (suffix 0x06 ... 0x80) to the final block so this module only
//   does the absorb loop: state ^= block; state = Keccak-p(state); repeat.
//   The single-block paths above are bit-identical when mb_en=0.
//
// Interface (single-block):
//   mode[1:0]   - 00=G, 01=H, 10=J
//   data_i      - 512-bit message input
//   valid_i / ready_o / hash_o[511:0] / valid_o / ready_i
//
// Interface (multi-block SHA3-256, active when mb_en=1):
//   mb_en        - 1 selects the multi-block absorb path
//   mb_block_i   - one pre-padded 1088-bit (136-byte) rate block, byte 0 in [7:0]
//   mb_valid_i   - block valid
//   mb_ready_o   - module can accept a block
//   mb_last_i    - asserted with the final (already-padded) block
//   result reuses hash_o[255:0] / valid_o / ready_i

module sha3_top (
    input              clk,
    input              rst_n,
    input  [1:0]       mode,
    input  [511:0]     data_i,
    input              valid_i,
    output             ready_o,
    output [511:0]     hash_o,
    output             valid_o,
    input              ready_i,
    // --- multi-block SHA3-256 absorb (tie mb_en=0 to disable) ---
    input              mb_en,
    input  [1087:0]    mb_block_i,
    input              mb_valid_i,
    input              mb_last_i,
    output             mb_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;

    // ================================================================
    // Absorb state: message || suffix || pad10*1 into rate bits
    // Built combinationally from data_i and mode.
    //
    // G: padded_block = {1, 308'b0, 1, 2'b01, data_i[263:0]}
    //    absorb_state = {1024'b0, padded_block_576}
    //
    // H: padded_block = {1, 828'b0, 1, 2'b01, data_i[255:0]}
    //    absorb_state = {512'b0, padded_block_1088}
    //
    // J: padded_block = {1, 570'b0, 1, 4'b1111, data_i[511:0]}
    //    absorb_state = {512'b0, padded_block_1088}
    // ================================================================
    wire [575:0]  g_pad;
    wire [1087:0] h_pad;
    wire [1087:0] j_pad;

    assign g_pad = {1'b1, {308{1'b0}}, 1'b1, 2'b10, data_i[263:0]};
    assign h_pad = {1'b1, {828{1'b0}}, 1'b1, 2'b10, data_i[255:0]};
    // J: SHAKE suffix is "1111" — all ones, order irrelevant
    assign j_pad = {1'b1, {570{1'b0}}, 1'b1, 4'b1111, data_i[511:0]};

    wire [1599:0] absorb_state;
    assign absorb_state = (mode == 2'b00) ? {{(1600-576){1'b0}}, g_pad} :
                          (mode == 2'b01) ? {{(1600-1088){1'b0}}, h_pad} :
                          (mode == 2'b10) ? {{(1600-1088){1'b0}}, j_pad} :
                          1600'd0;

    // ================================================================
    // Multi-block SHA3-256 absorb FSM (active only when mb_en=1)
    //
    // Running sponge state mb_state_r (init 0). For each pre-padded
    // 1088-bit rate block: mb_state_r ^= block; mb_state_r = Keccak-p(...).
    // After the last block, squeeze 256 bits. The caller pads the final
    // block (SHA3-256 suffix 0x06 ... 0x80), so this FSM is pure absorb.
    // ================================================================
    localparam MB_IDLE    = 2'd0;   // ready for a block (or first block)
    localparam MB_PERMUTE = 2'd1;   // keccak running on xored state
    localparam MB_DONE    = 2'd2;   // squeeze: present 256-bit digest

    reg  [1:0]    mb_state, mb_state_next;
    reg  [1599:0] mb_state_r;        // running sponge state
    reg           mb_last_r;         // captured last-block flag
    reg  [255:0]  mb_digest_r;       // latched 256-bit digest (sticky in MB_DONE)

    // XOR the incoming block into the low 1088 bits (rate) of the state.
    wire [1599:0] mb_xored;
    assign mb_xored = mb_state_r ^ {{(1600-1088){1'b0}}, mb_block_i};

    // Accept a block only in MB_IDLE while enabled.
    assign mb_ready_o = mb_en && (mb_state == MB_IDLE);
    wire   mb_accept  = mb_en && (mb_state == MB_IDLE) && mb_valid_i;
    wire   mb_kc_valid = mb_accept;   // start keccak on the accept cycle

    always @(*) begin
        mb_state_next = mb_state;
        case (mb_state)
            MB_IDLE:    if (mb_accept)   mb_state_next = MB_PERMUTE;
            MB_PERMUTE: if (kc_valid_o)  mb_state_next = mb_last_r ? MB_DONE : MB_IDLE;
            MB_DONE:    if (ready_i)      mb_state_next = MB_IDLE;
            default:    mb_state_next = MB_IDLE;
        endcase
    end

    // ================================================================
    // Keccak core
    // ================================================================
    wire              kc_valid_i;
    /* verilator lint_off UNUSEDSIGNAL */
    wire [1599:0]     kc_state_o;
    wire              kc_ready_o;
    /* verilator lint_on UNUSEDSIGNAL */
    wire              kc_valid_o;

    // Keccak input: multi-block xored state when mb_en, else single-block absorb.
    wire [1599:0] kc_state_i_mux;
    assign kc_state_i_mux = mb_en ? mb_xored : absorb_state;

    keccak_core #(.ROUNDS(24)) u_keccak (
        .clk      (clk),
        .rst_n    (rst_n),
        .state_i  (kc_state_i_mux),
        .valid_i  (kc_valid_i),
        .ready_o  (kc_ready_o),      // unused but must connect
        .state_o  (kc_state_o),
        .valid_o  (kc_valid_o),
        .ready_i  (1'b1)   // always accept output
    );

    // ================================================================
    // FSM combinational logic
    // ================================================================
    assign ready_o = !mb_en && (state_r == ST_IDLE);

    // kc_valid_i: single-block start (state_next==PERMUTE) OR multi-block accept.
    assign kc_valid_i = mb_en ? mb_kc_valid : (state_next == ST_PERMUTE);

    always @(*) begin
        state_next = state_r;
        case (state_r)
            ST_IDLE:    if (valid_i && ready_o)  state_next = ST_PERMUTE;
            ST_PERMUTE: if (kc_valid_o)           state_next = ST_SQUEEZE;
            ST_SQUEEZE: if (valid_o && ready_i)   state_next = ST_IDLE;
            default:    state_next = ST_IDLE;
        endcase
    end

    // ================================================================
    // Output
    // ================================================================
    // Register for squeezed output (only 512 bits needed)
    reg [511:0] squeezed_state_r;

    // valid_o / hash_o serve both paths, selected by mb_en.
    assign valid_o = mb_en ? (mb_state == MB_DONE) : (state_r == ST_SQUEEZE);
    assign hash_o  = mb_en ? {256'b0, mb_digest_r} : squeezed_state_r;

    // ================================================================
    // Sequential logic
    // ================================================================

    always @(posedge clk or negedge rst_n) begin
        if (!rst_n) begin
            state_r          <= ST_IDLE;
            squeezed_state_r <= 512'd0;
            mb_state         <= MB_IDLE;
            mb_state_r       <= 1600'd0;
            mb_last_r        <= 1'b0;
            mb_digest_r      <= 256'd0;
        end else begin
            state_r  <= state_next;
            mb_state <= mb_state_next;

            // --- single-block: latch squeezed output ---
            if (state_r == ST_PERMUTE && kc_valid_o) begin
                squeezed_state_r <= kc_state_o[511:0];
            end

            // --- multi-block: capture last flag on accept ---
            if (mb_accept) begin
                mb_last_r <= mb_last_i;
            end

            // --- multi-block: latch permuted state when keccak finishes ---
            if (mb_state == MB_PERMUTE && kc_valid_o) begin
                mb_state_r <= kc_state_o;
                // On the final block, latch the 256-bit digest (sticky for MB_DONE).
                if (mb_last_r) mb_digest_r <= kc_state_o[255:0];
            end

            // --- multi-block: clear running state after digest consumed ---
            if (mb_state == MB_DONE && ready_i) begin
                mb_state_r <= 1600'd0;
            end
        end
    end

endmodule
