// 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..0x3fff
//   FFT:    0x8000..0xcfff
// Route the low ML-KEM block to ML-KEM and the remaining accelerator window to
// FFT. ML-KEM runs in the same clock domain as this AXI slave.
(* use_dsp = "no" *) 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
);

    localparam SEL_MLKEM = 2'd0;
    localparam SEL_FFT   = 2'd1;

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

    wire [1:0] aw_sel = aw_to_mlkem ? SEL_MLKEM : SEL_FFT;
    wire [1:0] ar_sel = ar_to_mlkem ? SEL_MLKEM : SEL_FFT;

    reg [1:0] write_sel;
    reg write_active;
    reg [1:0] read_sel;
    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        mlkem_m_awvalid;
    wire        mlkem_m_awready;
    wire [31:0] mlkem_m_awaddr;
    wire [4:0]  mlkem_m_awid;
    wire [7:0]  mlkem_m_awlen;
    wire [2:0]  mlkem_m_awsize;
    wire [1:0]  mlkem_m_awburst;
    wire [0:0]  mlkem_m_awlock;
    wire [3:0]  mlkem_m_awcache;
    wire [2:0]  mlkem_m_awprot;
    wire        mlkem_m_wvalid;
    wire        mlkem_m_wready;
    wire [31:0] mlkem_m_wdata;
    wire [3:0]  mlkem_m_wstrb;
    wire        mlkem_m_wlast;
    wire        mlkem_m_bvalid;
    wire        mlkem_m_bready;
    wire [4:0]  mlkem_m_bid;
    wire [1:0]  mlkem_m_bresp;
    wire        mlkem_m_arvalid;
    wire        mlkem_m_arready;
    wire [31:0] mlkem_m_araddr;
    wire [4:0]  mlkem_m_arid;
    wire [7:0]  mlkem_m_arlen;
    wire [2:0]  mlkem_m_arsize;
    wire [1:0]  mlkem_m_arburst;
    wire [0:0]  mlkem_m_arlock;
    wire [3:0]  mlkem_m_arcache;
    wire [2:0]  mlkem_m_arprot;
    wire        mlkem_m_rvalid;
    wire        mlkem_m_rready;
    wire [31:0] mlkem_m_rdata;
    wire [4:0]  mlkem_m_rid;
    wire [1:0]  mlkem_m_rresp;
    wire        mlkem_m_rlast;
    wire        mlkem_intr_m;

    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 [1:0] sel_w = write_active ? write_sel : aw_sel;
    wire [1:0] sel_r = read_active  ? read_sel  : ar_sel;

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

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

    always @(posedge clk or negedge resetn) begin
        if (!resetn) begin
            write_active    <= 1'b0;
            write_sel       <= SEL_MLKEM;
            read_active     <= 1'b0;
            read_sel        <= SEL_MLKEM;
        end else begin
            if (!write_active && s_axi_awvalid && s_axi_awready) begin
                write_active    <= 1'b1;
                write_sel       <= aw_sel;
            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       <= ar_sel;
            end else if (read_active && s_axi_rvalid && s_axi_rready && s_axi_rlast) begin
                read_active <= 1'b0;
            end
        end
    end

    assign mlkem_intr = mlkem_intr_m;

    assign mlkem_awready = mlkem_m_awready;
    assign mlkem_wready  = mlkem_m_wready;
    assign mlkem_bvalid  = mlkem_m_bvalid;
    assign mlkem_bid     = mlkem_m_bid;
    assign mlkem_bresp   = mlkem_m_bresp;
    assign mlkem_arready = mlkem_m_arready;
    assign mlkem_rvalid  = mlkem_m_rvalid;
    assign mlkem_rlast   = mlkem_m_rlast;
    assign mlkem_rid     = mlkem_m_rid;
    assign mlkem_rdata   = mlkem_m_rdata;
    assign mlkem_rresp   = mlkem_m_rresp;

    assign mlkem_m_awvalid = s_axi_awvalid && !write_active && (aw_sel == SEL_MLKEM);
    assign mlkem_m_awaddr  = s_axi_awaddr;
    assign mlkem_m_awid    = s_axi_awid;
    assign mlkem_m_awlen   = s_axi_awlen;
    assign mlkem_m_awsize  = s_axi_awsize;
    assign mlkem_m_awburst = s_axi_awburst;
    assign mlkem_m_awlock  = s_axi_awlock;
    assign mlkem_m_awcache = s_axi_awcache;
    assign mlkem_m_awprot  = s_axi_awprot;
    assign mlkem_m_wvalid  = s_axi_wvalid && (sel_w == SEL_MLKEM);
    assign mlkem_m_wdata   = s_axi_wdata;
    assign mlkem_m_wstrb   = s_axi_wstrb;
    assign mlkem_m_wlast   = s_axi_wlast;
    assign mlkem_m_bready  = s_axi_bready && (sel_w == SEL_MLKEM);
    assign mlkem_m_arvalid = s_axi_arvalid && !read_active && (ar_sel == SEL_MLKEM);
    assign mlkem_m_araddr  = s_axi_araddr;
    assign mlkem_m_arid    = s_axi_arid;
    assign mlkem_m_arlen   = s_axi_arlen;
    assign mlkem_m_arsize  = s_axi_arsize;
    assign mlkem_m_arburst = s_axi_arburst;
    assign mlkem_m_arlock  = s_axi_arlock;
    assign mlkem_m_arcache = s_axi_arcache;
    assign mlkem_m_arprot  = s_axi_arprot;
    assign mlkem_m_rready  = s_axi_rready && (sel_r == SEL_MLKEM);

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

        .s_axi_awid     (mlkem_m_awid),
        .s_axi_awaddr   (mlkem_m_awaddr),
        .s_axi_awlen    (mlkem_m_awlen),
        .s_axi_awsize   (mlkem_m_awsize),
        .s_axi_awburst  (mlkem_m_awburst),
        .s_axi_awlock   (mlkem_m_awlock[0]),
        .s_axi_awcache  (mlkem_m_awcache),
        .s_axi_awprot   (mlkem_m_awprot),
        .s_axi_awvalid  (mlkem_m_awvalid),
        .s_axi_awready  (mlkem_m_awready),

        .s_axi_wid      (5'd0),
        .s_axi_wdata    (mlkem_m_wdata),
        .s_axi_wstrb    (mlkem_m_wstrb),
        .s_axi_wlast    (mlkem_m_wlast),
        .s_axi_wvalid   (mlkem_m_wvalid),
        .s_axi_wready   (mlkem_m_wready),

        .s_axi_bid      (mlkem_m_bid),
        .s_axi_bresp    (mlkem_m_bresp),
        .s_axi_bvalid   (mlkem_m_bvalid),
        .s_axi_bready   (mlkem_m_bready),

        .s_axi_arid     (mlkem_m_arid),
        .s_axi_araddr   (mlkem_m_araddr),
        .s_axi_arlen    (mlkem_m_arlen),
        .s_axi_arsize   (mlkem_m_arsize),
        .s_axi_arburst  (mlkem_m_arburst),
        .s_axi_arlock   (mlkem_m_arlock[0]),
        .s_axi_arcache  (mlkem_m_arcache),
        .s_axi_arprot   (mlkem_m_arprot),
        .s_axi_arvalid  (mlkem_m_arvalid),
        .s_axi_arready  (mlkem_m_arready),

        .s_axi_rid      (mlkem_m_rid),
        .s_axi_rdata    (mlkem_m_rdata),
        .s_axi_rresp    (mlkem_m_rresp),
        .s_axi_rlast    (mlkem_m_rlast),
        .s_axi_rvalid   (mlkem_m_rvalid),
        .s_axi_rready   (mlkem_m_rready),

        .intr           (mlkem_intr_m)
    );

    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_sel == SEL_FFT)),
        .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 == SEL_FFT)),
        .s_wready       (fft_wready),

        .s_bid          (fft_bid),
        .s_bresp        (fft_bresp),
        .s_bvalid       (fft_bvalid),
        .s_bready       (s_axi_bready && (sel_w == SEL_FFT)),

        .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_sel == SEL_FFT)),
        .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 == SEL_FFT)),

        .fft_finish     (fft_finish)
    );

endmodule