module matmul_axi_slave (
    input             clk,
    input             resetn,

    input      [4:0]  s_axi_awid,
    input      [31:0] s_axi_awaddr,
    input      [7:0]  s_axi_awlen,
    input      [2:0]  s_axi_awsize,
    input      [1:0]  s_axi_awburst,
    input             s_axi_awlock,
    input      [3:0]  s_axi_awcache,
    input      [2:0]  s_axi_awprot,
    input             s_axi_awvalid,
    output reg        s_axi_awready,

    input      [4:0]  s_axi_wid,
    input      [31:0] s_axi_wdata,
    input      [3:0]  s_axi_wstrb,
    input             s_axi_wlast,
    input             s_axi_wvalid,
    output reg        s_axi_wready,

    output     [4:0]  s_axi_bid,
    output     [1:0]  s_axi_bresp,
    output reg        s_axi_bvalid,
    input             s_axi_bready,

    input      [4:0]  s_axi_arid,
    input      [31:0] s_axi_araddr,
    input      [7:0]  s_axi_arlen,
    input      [2:0]  s_axi_arsize,
    input      [1:0]  s_axi_arburst,
    input             s_axi_arlock,
    input      [3:0]  s_axi_arcache,
    input      [2:0]  s_axi_arprot,
    input             s_axi_arvalid,
    output reg        s_axi_arready,

    output reg [4:0]  s_axi_rid,
    output reg [31:0] s_axi_rdata,
    output     [1:0]  s_axi_rresp,
    output reg        s_axi_rlast,
    output reg        s_axi_rvalid,
    input             s_axi_rready
);

localparam ADDR_CTRL      = 12'h000;
localparam ADDR_STATUS    = 12'h004;
localparam ADDR_SRC_BASE  = 12'h008;
localparam ADDR_DST_BASE  = 12'h00c;
localparam ADDR_GROUP_NUM = 12'h010;
localparam ADDR_A_BASE    = 12'h020;
localparam ADDR_B_BASE    = 12'h060;
localparam ADDR_C_BASE    = 12'h0a0;

reg [31:0] src_base;
reg [31:0] dst_base;
reg [31:0] group_num;
reg [31:0] a_data [0:15];
reg [31:0] b_data [0:15];
reg [65:0] c_data [0:15];

// ══ Write burst state machine (same pattern as fft_controller) ══
localparam W_IDLE = 2'd0;
localparam W_DATA = 2'd1;
localparam W_RESP = 2'd2;
reg [1:0]  w_state;
reg [31:0] w_addr;
reg [4:0]  w_id;

wire [11:0] w_offset = w_addr[11:0];
wire w_active = (w_state == W_DATA) && s_axi_wvalid && s_axi_wready;
wire w_start  = (w_offset == ADDR_CTRL) && s_axi_wdata[0] && !busy;

assign s_axi_bid   = w_id;
assign s_axi_bresp = 2'b00;

// ══ Read burst state machine ══
localparam R_IDLE = 2'd0;
localparam R_DATA = 2'd1;
reg [1:0]  r_state;
reg [31:0] r_addr;
reg [4:0]  r_id;
reg [7:0]  r_cnt;

wire [11:0] r_offset = r_addr[11:0];
assign s_axi_rresp = 2'b00;

// ══ Compute engine ══
reg        busy;
reg        done;
reg        error;
reg [3:0]  calc_index;

function [31:0] apply_wstrb;
    input [31:0] old_value;
    input [31:0] new_value;
    input [3:0]  strobe;
    begin
        apply_wstrb = old_value;
        if (strobe[0]) apply_wstrb[7:0]   = new_value[7:0];
        if (strobe[1]) apply_wstrb[15:8]  = new_value[15:8];
        if (strobe[2]) apply_wstrb[23:16] = new_value[23:16];
        if (strobe[3]) apply_wstrb[31:24] = new_value[31:24];
    end
endfunction

function [63:0] mul_u32_shift_add;
    input [31:0] lhs;
    input [31:0] rhs;
    integer bit_index;
    begin
        mul_u32_shift_add = 64'b0;
        for (bit_index = 0; bit_index < 32; bit_index = bit_index + 1) begin
            if (rhs[bit_index]) begin
                mul_u32_shift_add = mul_u32_shift_add + ({32'b0, lhs} << bit_index);
            end
        end
    end
endfunction

function [65:0] calc_c_element;
    input [3:0] index;
    reg [1:0] row;
    reg [1:0] col;
    begin
        row = index[3:2];
        col = index[1:0];
        calc_c_element = {2'b0, mul_u32_shift_add(a_data[{row, 2'd0}], b_data[{2'd0, col}])}
                       + {2'b0, mul_u32_shift_add(a_data[{row, 2'd1}], b_data[{2'd1, col}])}
                       + {2'b0, mul_u32_shift_add(a_data[{row, 2'd2}], b_data[{2'd2, col}])}
                       + {2'b0, mul_u32_shift_add(a_data[{row, 2'd3}], b_data[{2'd3, col}])};
    end
endfunction

function [31:0] c_read_word;
    input [5:0] word_index;
    begin
        case (word_index)
            6'd0:  c_read_word = c_data[0][31:0];
            6'd1:  c_read_word = c_data[0][63:32];
            6'd2:  c_read_word = {30'b0, c_data[0][65:64]};
            6'd3:  c_read_word = c_data[1][31:0];
            6'd4:  c_read_word = c_data[1][63:32];
            6'd5:  c_read_word = {30'b0, c_data[1][65:64]};
            6'd6:  c_read_word = c_data[2][31:0];
            6'd7:  c_read_word = c_data[2][63:32];
            6'd8:  c_read_word = {30'b0, c_data[2][65:64]};
            6'd9:  c_read_word = c_data[3][31:0];
            6'd10: c_read_word = c_data[3][63:32];
            6'd11: c_read_word = {30'b0, c_data[3][65:64]};
            6'd12: c_read_word = c_data[4][31:0];
            6'd13: c_read_word = c_data[4][63:32];
            6'd14: c_read_word = {30'b0, c_data[4][65:64]};
            6'd15: c_read_word = c_data[5][31:0];
            6'd16: c_read_word = c_data[5][63:32];
            6'd17: c_read_word = {30'b0, c_data[5][65:64]};
            6'd18: c_read_word = c_data[6][31:0];
            6'd19: c_read_word = c_data[6][63:32];
            6'd20: c_read_word = {30'b0, c_data[6][65:64]};
            6'd21: c_read_word = c_data[7][31:0];
            6'd22: c_read_word = c_data[7][63:32];
            6'd23: c_read_word = {30'b0, c_data[7][65:64]};
            6'd24: c_read_word = c_data[8][31:0];
            6'd25: c_read_word = c_data[8][63:32];
            6'd26: c_read_word = {30'b0, c_data[8][65:64]};
            6'd27: c_read_word = c_data[9][31:0];
            6'd28: c_read_word = c_data[9][63:32];
            6'd29: c_read_word = {30'b0, c_data[9][65:64]};
            6'd30: c_read_word = c_data[10][31:0];
            6'd31: c_read_word = c_data[10][63:32];
            6'd32: c_read_word = {30'b0, c_data[10][65:64]};
            6'd33: c_read_word = c_data[11][31:0];
            6'd34: c_read_word = c_data[11][63:32];
            6'd35: c_read_word = {30'b0, c_data[11][65:64]};
            6'd36: c_read_word = c_data[12][31:0];
            6'd37: c_read_word = c_data[12][63:32];
            6'd38: c_read_word = {30'b0, c_data[12][65:64]};
            6'd39: c_read_word = c_data[13][31:0];
            6'd40: c_read_word = c_data[13][63:32];
            6'd41: c_read_word = {30'b0, c_data[13][65:64]};
            6'd42: c_read_word = c_data[14][31:0];
            6'd43: c_read_word = c_data[14][63:32];
            6'd44: c_read_word = {30'b0, c_data[14][65:64]};
            6'd45: c_read_word = c_data[15][31:0];
            6'd46: c_read_word = c_data[15][63:32];
            6'd47: c_read_word = {30'b0, c_data[15][65:64]};
            default: c_read_word = 32'b0;
        endcase
    end
endfunction

integer i;
always @(posedge clk) begin
    if (!resetn) begin
        src_base      <= 32'b0;
        dst_base      <= 32'b0;
        group_num     <= 32'd1;
        w_state       <= W_IDLE;
        s_axi_awready <= 1'b1;
        s_axi_wready  <= 1'b0;
        s_axi_bvalid  <= 1'b0;
        w_addr        <= 32'b0;
        w_id          <= 5'b0;
        r_state       <= R_IDLE;
        s_axi_arready <= 1'b1;
        s_axi_rvalid  <= 1'b0;
        s_axi_rlast   <= 1'b0;
        r_addr        <= 32'b0;
        r_id          <= 5'b0;
        r_cnt         <= 8'b0;
        s_axi_rid     <= 5'b0;
        s_axi_rdata   <= 32'b0;
        busy          <= 1'b0;
        done          <= 1'b0;
        error         <= 1'b0;
        calc_index    <= 4'b0;

        for (i = 0; i < 16; i = i + 1) begin
            a_data[i] <= 32'b0;
            b_data[i] <= 32'b0;
            c_data[i] <= 66'b0;
        end
    end else begin
        // ═══════════════════════════════════════════════════════
        //  Write burst FSM (FFT pattern)
        // ═══════════════════════════════════════════════════════
        case (w_state)
            W_IDLE: begin
                s_axi_awready <= 1'b1;
                if (s_axi_awvalid && s_axi_awready) begin
                    s_axi_awready <= 1'b0;
                    s_axi_wready  <= 1'b1;
                    w_addr        <= s_axi_awaddr;
                    w_id          <= s_axi_awid;
                    w_state       <= W_DATA;
                end
            end
            W_DATA: begin
                if (s_axi_wvalid && s_axi_wready) begin
                    // Register write
                    if (w_offset == ADDR_SRC_BASE) begin
                        src_base <= apply_wstrb(src_base, s_axi_wdata, s_axi_wstrb);
                    end else if (w_offset == ADDR_DST_BASE) begin
                        dst_base <= apply_wstrb(dst_base, s_axi_wdata, s_axi_wstrb);
                    end else if (w_offset == ADDR_GROUP_NUM) begin
                        group_num <= apply_wstrb(group_num, s_axi_wdata, s_axi_wstrb);
                    end else if ((w_offset >= ADDR_A_BASE) && (w_offset < ADDR_B_BASE)) begin
                        a_data[(w_offset - ADDR_A_BASE) >> 2] <=
                            apply_wstrb(a_data[(w_offset - ADDR_A_BASE) >> 2], s_axi_wdata, s_axi_wstrb);
                    end else if ((w_offset >= ADDR_B_BASE) && (w_offset < ADDR_C_BASE)) begin
                        b_data[(w_offset - ADDR_B_BASE) >> 2] <=
                            apply_wstrb(b_data[(w_offset - ADDR_B_BASE) >> 2], s_axi_wdata, s_axi_wstrb);
                    end else if (w_offset == ADDR_STATUS) begin
                        if (s_axi_wdata[1]) done <= 1'b0;
                        if (s_axi_wdata[2]) error <= 1'b0;
                    end else if ((w_offset != ADDR_CTRL) && (w_offset < ADDR_C_BASE)) begin
                        error <= 1'b1;
                    end

                    // Auto-increment address for next beat
                    w_addr <= w_addr + 32'd4;

                    if (s_axi_wlast) begin
                        s_axi_wready <= 1'b0;
                        s_axi_bvalid <= 1'b1;
                        w_state      <= W_RESP;
                    end
                end
            end
            W_RESP: begin
                if (s_axi_bvalid && s_axi_bready) begin
                    s_axi_bvalid  <= 1'b0;
                    s_axi_awready <= 1'b1;
                    w_state       <= W_IDLE;
                end
            end
            default: w_state <= W_IDLE;
        endcase

        // ═══════════════════════════════════════════════════════
        //  Read burst FSM (FFT pattern, no BRAM latency)
        // ═══════════════════════════════════════════════════════
        case (r_state)
            R_IDLE: begin
                s_axi_arready <= 1'b1;
                if (s_axi_arvalid && s_axi_arready) begin
                    s_axi_arready <= 1'b0;
                    r_addr        <= s_axi_araddr;
                    r_id          <= s_axi_arid;
                    r_cnt         <= 8'd0;
                    s_axi_rvalid  <= 1'b1;
                    s_axi_rlast   <= (s_axi_arlen == 8'd0);

                    // Output data for current address
                    if (r_offset == ADDR_CTRL)
                        s_axi_rdata <= 32'b0;
                    else if (r_offset == ADDR_STATUS)
                        s_axi_rdata <= {29'b0, error, done, busy};
                    else if (r_offset == ADDR_SRC_BASE)
                        s_axi_rdata <= src_base;
                    else if (r_offset == ADDR_DST_BASE)
                        s_axi_rdata <= dst_base;
                    else if (r_offset == ADDR_GROUP_NUM)
                        s_axi_rdata <= group_num;
                    else if ((r_offset >= ADDR_A_BASE) && (r_offset < ADDR_B_BASE))
                        s_axi_rdata <= a_data[(r_offset - ADDR_A_BASE) >> 2];
                    else if ((r_offset >= ADDR_B_BASE) && (r_offset < ADDR_C_BASE))
                        s_axi_rdata <= b_data[(r_offset - ADDR_B_BASE) >> 2];
                    else if ((r_offset >= ADDR_C_BASE) && (r_offset < (ADDR_C_BASE + 12'h0c0)))
                        s_axi_rdata <= c_read_word((r_offset - ADDR_C_BASE) >> 2);
                    else
                        s_axi_rdata <= 32'b0;

                    s_axi_rid <= s_axi_arid;
                    r_state   <= R_DATA;
                end
            end
            R_DATA: begin
                if (s_axi_rvalid && s_axi_rready) begin
                    if (r_cnt == s_axi_arlen) begin
                        s_axi_rvalid  <= 1'b0;
                        s_axi_rlast   <= 1'b0;
                        s_axi_arready <= 1'b1;
                        r_state       <= R_IDLE;
                    end else begin
                        r_cnt    <= r_cnt + 8'd1;
                        r_addr   <= r_addr + 32'd4;
                        s_axi_rlast <= (r_cnt + 8'd1 == s_axi_arlen);

                        // Output next beat data
                        if ((r_offset + 12'd4) == ADDR_CTRL)
                            s_axi_rdata <= 32'b0;
                        else if ((r_offset + 12'd4) == ADDR_STATUS)
                            s_axi_rdata <= {29'b0, error, done, busy};
                        else if ((r_offset + 12'd4) == ADDR_SRC_BASE)
                            s_axi_rdata <= src_base;
                        else if ((r_offset + 12'd4) == ADDR_DST_BASE)
                            s_axi_rdata <= dst_base;
                        else if ((r_offset + 12'd4) == ADDR_GROUP_NUM)
                            s_axi_rdata <= group_num;
                        else if (((r_offset + 12'd4) >= ADDR_A_BASE) && ((r_offset + 12'd4) < ADDR_B_BASE))
                            s_axi_rdata <= a_data[((r_offset + 12'd4) - ADDR_A_BASE) >> 2];
                        else if (((r_offset + 12'd4) >= ADDR_B_BASE) && ((r_offset + 12'd4) < ADDR_C_BASE))
                            s_axi_rdata <= b_data[((r_offset + 12'd4) - ADDR_B_BASE) >> 2];
                        else if (((r_offset + 12'd4) >= ADDR_C_BASE) && ((r_offset + 12'd4) < (ADDR_C_BASE + 12'h0c0)))
                            s_axi_rdata <= c_read_word(((r_offset + 12'd4) - ADDR_C_BASE) >> 2);
                        else
                            s_axi_rdata <= 32'b0;
                    end
                end
            end
            default: r_state <= R_IDLE;
        endcase

        // ═══════════════════════════════════════════════════════
        //  Compute engine (unchanged)
        // ═══════════════════════════════════════════════════════
        if (w_start) begin
            busy       <= 1'b1;
            done       <= 1'b0;
            error      <= 1'b0;
            calc_index <= 4'd0;
        end else if (busy) begin
            c_data[calc_index] <= calc_c_element(calc_index);
            if (calc_index == 4'd15) begin
                busy       <= 1'b0;
                done       <= 1'b1;
                calc_index <= 4'd0;
            end else begin
                calc_index <= calc_index + 4'd1;
            end
        end
    end
end

endmodule