// sha3_top_shared.v - SHA3/SHAKE wrapper with EXTERNAL (shared) keccak_core.
//
// Identical to sha3_top.v except the internal keccak_core instance is
// replaced by explicit kc_* ports, so several Keccak consumers can share a
// single keccak_core via a phase mux (see mlkem_top). The sponge state
// (mb_state_r) and all FSM/squeeze logic stay here; only the 24-round
// permutation is external. Bit-identical to sha3_top when wired to a
// dedicated keccak_core.
//
// 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 to the final block.
//
// Keccak interface (port names mirror the keccak_core port each connects to):
//   kc_state_o[1599:0] - keccak result          (input,  from keccak_core.state_o)
//   kc_valid_o         - keccak output valid     (input,  from keccak_core.valid_o)
//   kc_ready_o         - keccak ready for input  (input,  from keccak_core.ready_o)
//   kc_state_i[1599:0] - keccak input state      (output, to   keccak_core.state_i)
//   kc_valid_i         - request permutation     (output, to   keccak_core.valid_i)
//   kc_ready_i         - accept keccak output    (output, to   keccak_core.ready_i, always 1)

module sha3_top_shared (
    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,
    // --- external (shared) keccak_core interface ---
    input  [1599:0]    kc_state_o,
    input              kc_valid_o,
    /* verilator lint_off UNUSEDSIGNAL */
    input              kc_ready_o,
    /* verilator lint_on UNUSEDSIGNAL */
    output [1599:0]    kc_state_i,
    output             kc_valid_i,
    output             kc_ready_i
);

    // ================================================================
    // 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
    // ================================================================
    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)
    // ================================================================
    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

    // ================================================================
    // External keccak_core interface (was internal keccak_core)
    // ================================================================
    // Keccak input: multi-block xored state when mb_en, else single-block absorb.
    assign kc_state_i = mb_en ? mb_xored : absorb_state;
    // Always accept keccak output (matches the dedicated-core ready_i=1'b1).
    assign kc_ready_i = 1'b1;
    // 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);

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

    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
    // ================================================================
    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;
                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