`timescale 1ns / 1ps

module spi_controller (
    input wire clk,
    input wire rst_n,
    
    input wire [7:0] cmd_opcode,
    input wire [31:0] cmd_address,
    input wire [15:0] cmd_length,
    input wire cmd_valid,
    output wire cmd_ready,
    
    input wire [7:0] data_in,
    input wire data_in_valid,
    output wire data_in_ready,
    
    output wire [7:0] data_out,
    output wire data_out_valid,
    input wire data_out_ready,
    
    output wire spi0_sck,
    output wire spi0_mosi,
    input wire spi0_miso,
    output wire spi0_cs,
    
    output wire spi1_sck,
    output wire spi1_mosi,
    input wire spi1_miso,
    output wire spi1_cs,
    
    output reg [3:0] flash_select,
    output reg [3:0] spi_state,
    output reg [3:0] status_reg
);

    localparam STATE_IDLE        = 4'h0;
    localparam STATE_CMD        = 4'h1;
    localparam STATE_ADDR       = 4'h2;
    localparam STATE_DUMMY      = 4'h3;
    localparam STATE_DATA_RX    = 4'h4;
    localparam STATE_DATA_TX    = 4'h5;
    localparam STATE_WAIT       = 4'h6;
    localparam STATE_ERROR      = 4'h7;
    
    localparam MODE_STANDARD    = 2'b00;
    localparam MODE_FAST       = 2'b01;
    localparam MODE_DUAL       = 2'b10;
    localparam MODE_QUAD       = 2'b11;
    
    localparam CMD_WREN         = 8'h06;
    localparam CMD_WRDI        = 8'h04;
    localparam CMD_RDSR        = 8'h05;
    localparam CMD_WRSR        = 8'h01;
    localparam CMD_READ        = 8'h03;
    localparam CMD_FAST_READ   = 8'h0B;
    localparam CMD_DREAD       = 8'h3B;
    localparam CMD_QREAD       = 8'h6B;
    localparam CMD_PP          = 8'h02;
    localparam CMD_SE          = 8'h20;
    localparam CMD_BE32        = 8'h52;
    localparam CMD_BE64        = 8'hD8;
    localparam CMD_CE          = 8'hC7;
    localparam CMD_RDID        = 8'h9F;
    localparam CMD_RDFR        = 8'h48;
    localparam CMD_WRFR        = 8'h42;
    
    reg [7:0] spi_mode;
    reg [2:0] spi_divider;
    reg [31:0] current_addr;
    reg [15:0] bytes_remaining;
    reg [7:0] command_reg;
    reg [3:0] state;
    reg [3:0] next_state;
    
    reg [7:0] tx_shift_reg;
    reg [7:0] rx_shift_reg;
    reg [3:0] bit_counter;
    reg tx_active;
    reg rx_active;
    
    reg cs0_active;
    reg cs1_active;
    reg sck_active;
    
    reg [23:0] delay_counter;
    
    reg [7:0] clk_div_counter;
    wire sck_posedge;
    wire sck_negedge;
    
    assign sck_posedge = (clk_div_counter == spi_divider);
    assign sck_negedge = (clk_div_counter == (spi_divider >> 1));
    
    assign spi0_sck = sck_active & cs0_active & ~spi0_cs;
    assign spi1_sck = sck_active & cs1_active & ~spi1_cs;
    assign spi0_cs = ~cs0_active;
    assign spi1_cs = ~cs1_active;
    
    assign spi0_mosi = tx_active ? tx_shift_reg[7] : 1'bz;
    assign spi1_mosi = tx_active ? tx_shift_reg[7] : 1'bz;
    
    assign cmd_ready = (state == STATE_IDLE);
    
    assign data_in_ready = (state == STATE_DATA_TX) && tx_active && sck_negedge;
    assign data_out_valid = (state == STATE_DATA_RX) && rx_active && sck_posedge;
    
    always @(posedge clk or negedge rst_n) begin
        if (!rst_n) begin
            state <= STATE_IDLE;
            spi_state <= STATE_IDLE;
            cs0_active <= 1'b0;
            cs1_active <= 1'b0;
            sck_active <= 1'b0;
            tx_active <= 1'b0;
            rx_active <= 1'b0;
            spi_mode <= MODE_STANDARD;
            spi_divider <= 3'd1;
            flash_select <= 2'b00;
            status_reg <= 4'b0000;
        end else begin
            state <= next_state;
            spi_state <= state;
            
            case (state)
                STATE_IDLE: begin
                    cs0_active <= 1'b0;
                    cs1_active <= 1'b0;
                    sck_active <= 1'b0;
                    tx_active <= 1'b0;
                    rx_active <= 1'b0;
                    
                    if (cmd_valid && cmd_ready) begin
                        if (cmd_address[24] == 1'b0) begin
                            cs0_active <= 1'b1;
                            flash_select <= 2'b01;
                        end else begin
                            cs1_active <= 1'b1;
                            flash_select <= 2'b10;
                        end
                        
                        command_reg <= cmd_opcode;
                        current_addr <= cmd_address;
                        bytes_remaining <= cmd_length;
                        next_state <= STATE_CMD;
                    end
                end
                
                STATE_CMD: begin
                    if (sck_negedge) begin
                        if (bit_counter == 4'd7) begin
                            bit_counter <= 4'd0;
                            if (command_reg == CMD_READ || 
                                command_reg == CMD_FAST_READ ||
                                command_reg == CMD_PP) begin
                                next_state <= STATE_ADDR;
                            end else if (command_reg == CMD_RDSR ||
                                       command_reg == CMD_RDID) begin
                                next_state <= STATE_DATA_RX;
                            end else begin
                                next_state <= STATE_WAIT;
                            end
                        end else begin
                            tx_shift_reg <= tx_shift_reg << 1;
                            bit_counter <= bit_counter + 1;
                        end
                    end
                end
                
                STATE_ADDR: begin
                    if (sck_negedge) begin
                        if (bit_counter == 4'd7) begin
                            bit_counter <= 4'd0;
                            if (addr_byte_counter == 3'd2) begin
                                if (command_reg == CMD_FAST_READ ||
                                    command_reg == CMD_DREAD ||
                                    command_reg == CMD_QREAD) begin
                                    next_state <= STATE_DUMMY;
                                end else if (command_reg == CMD_READ) begin
                                    next_state <= STATE_DATA_RX;
                                end else if (command_reg == CMD_PP) begin
                                    next_state <= STATE_DATA_TX;
                                end else begin
                                    next_state <= STATE_WAIT;
                                end
                            end else begin
                                addr_byte_counter <= addr_byte_counter + 1;
                            end
                        end else begin
                            tx_shift_reg <= tx_shift_reg << 1;
                            bit_counter <= bit_counter + 1;
                        end
                    end
                end
                
                STATE_DATA_RX: begin
                    rx_active <= 1'b1;
                    if (sck_posedge) begin
                        rx_shift_reg <= {rx_shift_reg[6:0], (cs0_active ? spi0_miso : spi1_miso)};
                        if (bit_counter == 4'd7) begin
                            bit_counter <= 4'd0;
                            data_out <= rx_shift_reg;
                            if (bytes_remaining > 0) begin
                                bytes_remaining <= bytes_remaining - 1;
                            end else begin
                                next_state <= STATE_IDLE;
                                rx_active <= 1'b0;
                            end
                        end else begin
                            bit_counter <= bit_counter + 1;
                        end
                    end
                end
                
                STATE_DATA_TX: begin
                    tx_active <= 1'b1;
                    if (sck_negedge) begin
                        if (bit_counter == 4'd7) begin
                            bit_counter <= 4'd0;
                            if (data_in_valid && data_in_ready) begin
                                tx_shift_reg <= data_in;
                                if (bytes_remaining > 0) begin
                                    bytes_remaining <= bytes_remaining - 1;
                                end else begin
                                    next_state <= STATE_WAIT;
                                    tx_active <= 1'b0;
                                end
                            end
                        end else begin
                            tx_shift_reg <= tx_shift_reg << 1;
                            bit_counter <= bit_counter + 1;
                        end
                    end
                end
                
                STATE_WAIT: begin
                    if (delay_counter == 24'hFFFFFF) begin
                        next_state <= STATE_IDLE;
                    end else begin
                        delay_counter <= delay_counter + 1;
                    end
                end
                
                STATE_ERROR: begin
                    cs0_active <= 1'b0;
                    cs1_active <= 1'b0;
                    sck_active <= 1'b0;
                    next_state <= STATE_IDLE;
                end
            endcase
        end
    end
    
    always @(posedge clk or negedge rst_n) begin
        if (!rst_n) begin
            clk_div_counter <= 8'd0;
            sck_active <= 1'b0;
        end else begin
            if (cs0_active || cs1_active) begin
                if (clk_div_counter >= spi_divider) begin
                    clk_div_counter <= 8'd0;
                    sck_active <= ~sck_active;
                end else begin
                    clk_div_counter <= clk_div_counter + 1;
                end
            end else begin
                clk_div_counter <= 8'd0;
                sck_active <= 1'b0;
            end
        end
    end
    
    reg [2:0] addr_byte_counter;
    
    always @(posedge clk or negedge rst_n) begin
        if (!rst_n) begin
            addr_byte_counter <= 3'd0;
        end else if (state == STATE_ADDR) begin
        end else begin
            addr_byte_counter <= 3'd0;
        end
    end
    
    initial begin
        tx_shift_reg = 8'h00;
        rx_shift_reg = 8'h00;
        bit_counter = 4'd0;
        delay_counter = 24'h0;
    end

endmodule