module axi_demux_1to2 (
    input  wire        clk,
    input  wire        rst_n,
    
    // External selection signals
    input  wire        aw_sel, // 0: m0, 1: m1
    input  wire        ar_sel, // 0: m0, 1: m1
    
    // Master interface (s_*)
    // Write Address Channel
    input  wire [31:0] s_awaddr,
    input  wire [7:0]  s_awlen,
    input  wire [2:0]  s_awsize,
    input  wire [1:0]  s_awburst,
    input  wire        s_awvalid,
    output wire        s_awready,
    
    // Write Data Channel
    input  wire [31:0] s_wdata,
    input  wire [3:0]  s_wstrb,
    input  wire        s_wlast,
    input  wire        s_wvalid,
    output wire        s_wready,
    
    // Write Response Channel
    output wire [1:0]  s_bresp,
    output wire        s_bvalid,
    input  wire        s_bready,
    
    // Read Address Channel
    input  wire [31:0] s_araddr,
    input  wire [7:0]  s_arlen,
    input  wire [2:0]  s_arsize,
    input  wire [1:0]  s_arburst,
    input  wire        s_arvalid,
    output wire        s_arready,
    
    // Read Data Channel
    output wire [31:0] s_rdata,
    output wire [1:0]  s_rresp,
    output wire        s_rlast,
    output wire        s_rvalid,
    input  wire        s_rready,
    
    // Slave 0 interface (m0_*)
    // Write Address Channel
    output wire [31:0] m0_awaddr,
    output wire [7:0]  m0_awlen,
    output wire [2:0]  m0_awsize,
    output wire [1:0]  m0_awburst,
    output wire        m0_awvalid,
    input  wire        m0_awready,
    
    // Write Data Channel
    output wire [31:0] m0_wdata,
    output wire [3:0]  m0_wstrb,
    output wire        m0_wlast,
    output wire        m0_wvalid,
    input  wire        m0_wready,
    
    // Write Response Channel
    input  wire [1:0]  m0_bresp,
    input  wire        m0_bvalid,
    output wire        m0_bready,
    
    // Read Address Channel
    output wire [31:0] m0_araddr,
    output wire [7:0]  m0_arlen,
    output wire [2:0]  m0_arsize,
    output wire [1:0]  m0_arburst,
    output wire        m0_arvalid,
    input  wire        m0_arready,
    
    // Read Data Channel
    input  wire [31:0] m0_rdata,
    input  wire [1:0]  m0_rresp,
    input  wire        m0_rlast,
    input  wire        m0_rvalid,
    output wire        m0_rready,
    
    // Slave 1 interface (m1_*)
    // Write Address Channel
    output wire [31:0] m1_awaddr,
    output wire [7:0]  m1_awlen,
    output wire [2:0]  m1_awsize,
    output wire [1:0]  m1_awburst,
    output wire        m1_awvalid,
    input  wire        m1_awready,
    
    // Write Data Channel
    output wire [31:0] m1_wdata,
    output wire [3:0]  m1_wstrb,
    output wire        m1_wlast,
    output wire        m1_wvalid,
    input  wire        m1_wready,
    
    // Write Response Channel
    input  wire [1:0]  m1_bresp,
    input  wire        m1_bvalid,
    output wire        m1_bready,
    
    // Read Address Channel
    output wire [31:0] m1_araddr,
    output wire [7:0]  m1_arlen,
    output wire [2:0]  m1_arsize,
    output wire [1:0]  m1_arburst,
    output wire        m1_arvalid,
    input  wire        m1_arready,
    
    // Read Data Channel
    input  wire [31:0] m1_rdata,
    input  wire [1:0]  m1_rresp,
    input  wire        m1_rlast,
    input  wire        m1_rvalid,
    output wire        m1_rready
);

//-------------------------------------------------------------------------
// Write Channel Control Logic
//-------------------------------------------------------------------------
reg        aw_active;     // Active write transaction
reg        aw_sel_reg;    // Registered selection for write data path
wire       aw_handshake = s_awvalid & s_awready;
wire       w_handshake  = s_wvalid & s_wready;

// AW channel
assign s_awready = ~aw_active & (aw_sel ? m1_awready : m0_awready);
assign m0_awvalid = ~aw_sel & s_awvalid & ~aw_active;
assign m1_awvalid =  aw_sel & s_awvalid & ~aw_active;

// AW selection and active flag
always @(posedge clk or negedge rst_n) begin
    if (!rst_n) begin
        aw_active <= 1'b0;
        aw_sel_reg <= 1'b0;
    end else begin
        if (aw_handshake) begin
            aw_active <= 1'b1;
            aw_sel_reg <= aw_sel;
        end else if (w_handshake && s_wlast) begin
            aw_active <= 1'b0;
        end
    end
end

// W channel (routed based on registered selection)
assign m0_wdata  = s_wdata;
assign m0_wstrb  = s_wstrb;
assign m0_wlast  = s_wlast;
assign m0_wvalid = ~aw_sel_reg & aw_active & s_wvalid;

assign m1_wdata  = s_wdata;
assign m1_wstrb  = s_wstrb;
assign m1_wlast  = s_wlast;
assign m1_wvalid = aw_sel_reg & aw_active & s_wvalid;

assign s_wready = aw_sel_reg ? (m1_wready & aw_active) : (m0_wready & aw_active);

//-------------------------------------------------------------------------
// Read Address Channel (direct routing)
//-------------------------------------------------------------------------
assign s_arready = ar_sel ? m1_arready : m0_arready;
assign m0_arvalid = ~ar_sel & s_arvalid;
assign m1_arvalid =  ar_sel & s_arvalid;

assign m0_araddr  = s_araddr;
assign m0_arlen   = s_arlen;
assign m0_arsize  = s_arsize;
assign m0_arburst = s_arburst;

assign m1_araddr  = s_araddr;
assign m1_arlen   = s_arlen;
assign m1_arsize  = s_arsize;
assign m1_arburst = s_arburst;

//-------------------------------------------------------------------------
// Read Data Channel Arbitration (Round Robin)
//-------------------------------------------------------------------------
reg        r_active;       // Active read transaction
reg        r_sel_reg;      // Current selected slave (0:m0, 1:m1)
reg        r_arb_state;    // RR arbitration state
wire       r_handshake = s_rvalid & s_rready;
wire       r_finish    = r_handshake & s_rlast;

// Arbitration and selection logic
wire [1:0] r_request = {m1_rvalid, m0_rvalid};
wire       r_grant;

round_robin_arb #(.WIDTH(2)) r_arb (
    .clk       (clk),
    .rst_n     (rst_n),
    .enable    (~r_active),  // Only arbitrate when not active
    .request   (r_request),
    .grant     (r_grant)
);

always @(posedge clk or negedge rst_n) begin
    if (!rst_n) begin
        r_active <= 1'b0;
        r_sel_reg <= 1'b0;
    end else begin
        if (~r_active && |r_request) begin
            r_active <= 1'b1;
            r_sel_reg <= r_grant; // 0 for m0, 1 for m1
        end else if (r_finish) begin
            r_active <= 1'b0;
        end
    end
end

// Mux for read data channel
assign s_rdata = r_sel_reg ? m1_rdata : m0_rdata;
assign s_rresp = r_sel_reg ? m1_rresp : m0_rresp;
assign s_rlast = r_sel_reg ? m1_rlast : m0_rlast;
assign s_rvalid = r_sel_reg ? m1_rvalid : m0_rvalid;

assign m0_rready = (~r_sel_reg) & s_rready;
assign m1_rready = r_sel_reg & s_rready;

//-------------------------------------------------------------------------
// Write Response Arbitration (Round Robin)
//-------------------------------------------------------------------------
reg        b_arb_state;    // RR arbitration state
wire [1:0] b_request = {m1_bvalid, m0_bvalid};
wire       b_grant;

round_robin_arb #(.WIDTH(2)) b_arb (
    .clk       (clk),
    .rst_n     (rst_n),
    .enable    (1'b1),      // Always arbitrate
    .request   (b_request),
    .grant     (b_grant)
);

// Mux for write response channel
assign s_bresp = b_grant ? m1_bresp : m0_bresp;
assign s_bvalid = b_grant ? m1_bvalid : m0_bvalid;

assign m0_bready = (~b_grant) & s_bready;
assign m1_bready = b_grant & s_bready;

endmodule

//-------------------------------------------------------------------------
// Round Robin Arbiter Module
//-------------------------------------------------------------------------
module round_robin_arb #(
    parameter WIDTH = 2
)(
    input  wire             clk,
    input  wire             rst_n,
    input  wire             enable,    // Arbitration enable
    input  wire [WIDTH-1:0] request,   // Request vector
    output reg  [WIDTH-1:0] grant      // One-hot grant
);

reg [WIDTH-1:0] pointer;

always @(posedge clk or negedge rst_n) begin
    if (!rst_n) begin
        pointer <= 1;
    end else if (enable && |request) begin
        pointer <= {pointer[0], pointer[WIDTH-1:1]}; // Rotate right
    end
end

// Grant generation
always @(*) begin
    grant = {WIDTH{1'b0}};
    if (|request) begin
        for (int i = 0; i < WIDTH; i++) begin
            int idx = (i + pointer) % WIDTH;
            if (request[idx]) begin
                grant = (1 << idx);
                break;
            end
        end
    end
end

endmodule