// Shared AXI slave for the legacy FFT accelerator and the ML-KEM accelerator.
//
// Both software stacks use the same base address, 0xbf400000. Their internal
// maps do not overlap:
//   ML-KEM: 0x0000..0x0114
//   FFT:    0x1000, 0x2000, 0x3000, 0x4000, 0xf000 windows
// Route 0x0xxx to ML-KEM and all other offsets to FFT.
module fft_mlkem_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            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            s_axi_wready,

    output     [4:0]  s_axi_bid,
    output     [1:0]  s_axi_bresp,
    output            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            s_axi_arready,

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

    output            fft_finish,
    output            mlkem_intr
);

    wire aw_to_mlkem = (s_axi_awaddr[15:12] == 4'h0);
    wire ar_to_mlkem = (s_axi_araddr[15:12] == 4'h0);

    reg write_sel_mlkem;
    reg write_active;
    reg read_sel_mlkem;
    reg read_active;

    wire mlkem_awready, mlkem_wready, mlkem_bvalid;
    wire [4:0] mlkem_bid;
    wire [1:0] mlkem_bresp;
    wire mlkem_arready, mlkem_rvalid, mlkem_rlast;
    wire [4:0] mlkem_rid;
    wire [31:0] mlkem_rdata;
    wire [1:0] mlkem_rresp;

    wire fft_awready, fft_wready, fft_bvalid;
    wire [4:0] fft_bid;
    wire [1:0] fft_bresp;
    wire fft_arready, fft_rvalid, fft_rlast;
    wire [4:0] fft_rid;
    wire [31:0] fft_rdata;
    wire [1:0] fft_rresp;

    wire sel_w_mlkem = write_active ? write_sel_mlkem : aw_to_mlkem;
    wire sel_r_mlkem = read_active ? read_sel_mlkem : ar_to_mlkem;

    assign s_axi_awready = !write_active && (aw_to_mlkem ? mlkem_awready : fft_awready);
    assign s_axi_wready  = sel_w_mlkem ? mlkem_wready  : fft_wready;
    assign s_axi_bvalid  = sel_w_mlkem ? mlkem_bvalid  : fft_bvalid;
    assign s_axi_bid     = sel_w_mlkem ? mlkem_bid     : fft_bid;
    assign s_axi_bresp   = sel_w_mlkem ? mlkem_bresp   : fft_bresp;

    assign s_axi_arready = !read_active && (ar_to_mlkem ? mlkem_arready : fft_arready);
    assign s_axi_rvalid  = sel_r_mlkem ? mlkem_rvalid  : fft_rvalid;
    assign s_axi_rlast   = sel_r_mlkem ? mlkem_rlast   : fft_rlast;
    assign s_axi_rid     = sel_r_mlkem ? mlkem_rid     : fft_rid;
    assign s_axi_rdata   = sel_r_mlkem ? mlkem_rdata   : fft_rdata;
    assign s_axi_rresp   = sel_r_mlkem ? mlkem_rresp   : fft_rresp;

    always @(posedge clk or negedge resetn) begin
        if (!resetn) begin
            write_active    <= 1'b0;
            write_sel_mlkem <= 1'b0;
            read_active     <= 1'b0;
            read_sel_mlkem  <= 1'b0;
        end else begin
            if (!write_active && s_axi_awvalid && s_axi_awready) begin
                write_active    <= 1'b1;
                write_sel_mlkem <= aw_to_mlkem;
            end else if (write_active && s_axi_bvalid && s_axi_bready) begin
                write_active <= 1'b0;
            end

            if (!read_active && s_axi_arvalid && s_axi_arready) begin
                read_active    <= 1'b1;
                read_sel_mlkem <= ar_to_mlkem;
            end else if (read_active && s_axi_rvalid && s_axi_rready && s_axi_rlast) begin
                read_active <= 1'b0;
            end
        end
    end

    mlkem_axi_slave u_mlkem_axi_slave (
        .clk            (clk),
        .resetn         (resetn),

        .s_axi_awid     (s_axi_awid),
        .s_axi_awaddr   (s_axi_awaddr),
        .s_axi_awlen    (s_axi_awlen),
        .s_axi_awsize   (s_axi_awsize),
        .s_axi_awburst  (s_axi_awburst),
        .s_axi_awlock   (s_axi_awlock),
        .s_axi_awcache  (s_axi_awcache),
        .s_axi_awprot   (s_axi_awprot),
        .s_axi_awvalid  (s_axi_awvalid && !write_active && aw_to_mlkem),
        .s_axi_awready  (mlkem_awready),

        .s_axi_wid      (s_axi_wid),
        .s_axi_wdata    (s_axi_wdata),
        .s_axi_wstrb    (s_axi_wstrb),
        .s_axi_wlast    (s_axi_wlast),
        .s_axi_wvalid   (s_axi_wvalid && sel_w_mlkem),
        .s_axi_wready   (mlkem_wready),

        .s_axi_bid      (mlkem_bid),
        .s_axi_bresp    (mlkem_bresp),
        .s_axi_bvalid   (mlkem_bvalid),
        .s_axi_bready   (s_axi_bready && sel_w_mlkem),

        .s_axi_arid     (s_axi_arid),
        .s_axi_araddr   (s_axi_araddr),
        .s_axi_arlen    (s_axi_arlen),
        .s_axi_arsize   (s_axi_arsize),
        .s_axi_arburst  (s_axi_arburst),
        .s_axi_arlock   (s_axi_arlock),
        .s_axi_arcache  (s_axi_arcache),
        .s_axi_arprot   (s_axi_arprot),
        .s_axi_arvalid  (s_axi_arvalid && !read_active && ar_to_mlkem),
        .s_axi_arready  (mlkem_arready),

        .s_axi_rid      (mlkem_rid),
        .s_axi_rdata    (mlkem_rdata),
        .s_axi_rresp    (mlkem_rresp),
        .s_axi_rlast    (mlkem_rlast),
        .s_axi_rvalid   (mlkem_rvalid),
        .s_axi_rready   (s_axi_rready && sel_r_mlkem),

        .intr           (mlkem_intr)
    );

    fft_controller u_fft_controller (
        .aclk           (clk),
        .aresetn        (resetn),

        .s_awid         (s_axi_awid),
        .s_awaddr       (s_axi_awaddr),
        .s_awlen        (s_axi_awlen),
        .s_awsize       (s_axi_awsize),
        .s_awburst      (s_axi_awburst),
        .s_awlock       (s_axi_awlock),
        .s_awcache      (s_axi_awcache),
        .s_awprot       (s_axi_awprot),
        .s_awvalid      (s_axi_awvalid && !write_active && !aw_to_mlkem),
        .s_awready      (fft_awready),

        .s_wid          (s_axi_wid),
        .s_wdata        (s_axi_wdata),
        .s_wstrb        (s_axi_wstrb),
        .s_wlast        (s_axi_wlast),
        .s_wvalid       (s_axi_wvalid && !sel_w_mlkem),
        .s_wready       (fft_wready),

        .s_bid          (fft_bid),
        .s_bresp        (fft_bresp),
        .s_bvalid       (fft_bvalid),
        .s_bready       (s_axi_bready && !sel_w_mlkem),

        .s_arid         (s_axi_arid),
        .s_araddr       (s_axi_araddr),
        .s_arlen        (s_axi_arlen),
        .s_arsize       (s_axi_arsize),
        .s_arburst      (s_axi_arburst),
        .s_arlock       (s_axi_arlock),
        .s_arcache      (s_axi_arcache),
        .s_arprot       (s_axi_arprot),
        .s_arvalid      (s_axi_arvalid && !read_active && !ar_to_mlkem),
        .s_arready      (fft_arready),

        .s_rid          (fft_rid),
        .s_rdata        (fft_rdata),
        .s_rresp        (fft_rresp),
        .s_rlast        (fft_rlast),
        .s_rvalid       (fft_rvalid),
        .s_rready       (s_axi_rready && !sel_r_mlkem),

        .fft_finish     (fft_finish)
    );

endmodule