`timescale 1ns/1ps

module i2c_master (
    input clk,
    input rst_n,
    input [6:0] addr,
    input rw,
    input [7:0] data_in,
    input start_tx,
    output reg [7:0] data_out,
    output reg busy,
    output reg ack_err,
    inout scl,
    inout sda
);
    localparam IDLE             = 4'd0;
    localparam START_A          = 4'd1;
    localparam START_B          = 4'd2;
    localparam ADDR_DRIVE       = 4'd3;
    localparam ADDR_SAMPLE      = 4'd4;
    localparam ACK1             = 4'd5;
    localparam WRITE_DRIVE      = 4'd6;
    localparam WRITE_SAMPLE     = 4'd7;
    localparam READ             = 4'd8;
    localparam ACK2             = 4'd9;
    localparam STOP_A           = 4'd10;
    localparam STOP_B           = 4'd11;
    localparam CLOCK_STRETCH    = 4'd12;
    localparam ARBITRATION_LOST = 4'd13;
    localparam ERROR            = 4'd14;

    reg [3:0] state;
    reg [2:0] bit_count;
    reg [7:0] tx_shift;
    reg [7:0] rx_shift;
    reg scl_drive_low;
    reg sda_drive_low;
    reg latched_rw;

    wire scl_in;
    wire sda_in;

    assign scl = scl_drive_low ? 1'b0 : 1'bz;
    assign sda = sda_drive_low ? 1'b0 : 1'bz;
    assign scl_in = scl;
    assign sda_in = sda;

    always @(posedge clk or negedge rst_n) begin
        if (!rst_n) begin
            state <= IDLE;
            bit_count <= 3'd0;
            tx_shift <= 8'd0;
            rx_shift <= 8'd0;
            scl_drive_low <= 1'b0;
            sda_drive_low <= 1'b0;
            latched_rw <= 1'b0;
            data_out <= 8'd0;
            busy <= 1'b0;
            ack_err <= 1'b0;
        end else begin
            case (state)
                IDLE: begin
                    busy <= 1'b0;
                    scl_drive_low <= 1'b0;
                    sda_drive_low <= 1'b0;
                    bit_count <= 3'd7;
                    rx_shift <= 8'd0;
                    if (start_tx) begin
                        busy <= 1'b1;
                        ack_err <= 1'b0;
                        tx_shift <= {addr, rw};
                        latched_rw <= rw;
                        state <= START_A;
                    end
                end

                START_A: begin
                    busy <= 1'b1;
                    if (scl_in && sda_in) begin
                        sda_drive_low <= 1'b1;
                        scl_drive_low <= 1'b0;
                        state <= START_B;
                    end else begin
                        ack_err <= 1'b1;
                        state <= ERROR;
                    end
                end

                START_B: begin
                    busy <= 1'b1;
                    scl_drive_low <= 1'b1;
                    bit_count <= 3'd7;
                    sda_drive_low <= ~tx_shift[7];
                    state <= ADDR_DRIVE;
                end

                ADDR_DRIVE: begin
                    busy <= 1'b1;
                    scl_drive_low <= 1'b0;
                    sda_drive_low <= ~tx_shift[7];
                    state <= ADDR_SAMPLE;
                end

                ADDR_SAMPLE: begin
                    busy <= 1'b1;
                    scl_drive_low <= 1'b0;
                    if (tx_shift[7] && !sda_in) begin
                        ack_err <= 1'b1;
                        state <= ARBITRATION_LOST;
                    end else if (bit_count == 3'd0) begin
                        sda_drive_low <= 1'b0;
                        state <= ACK1;
                    end else begin
                        tx_shift <= {tx_shift[6:0], 1'b0};
                        bit_count <= bit_count - 1'b1;
                        scl_drive_low <= 1'b1;
                        sda_drive_low <= ~tx_shift[6];
                        state <= ADDR_DRIVE;
                    end
                end

                ACK1: begin
                    busy <= 1'b1;
                    scl_drive_low <= 1'b0;
                    sda_drive_low <= 1'b0;
                    bit_count <= 3'd7;
                    if (!sda_in) begin
                        if (latched_rw) begin
                            rx_shift <= 8'd0;
                            state <= READ;
                        end else begin
                            tx_shift <= data_in;
                            sda_drive_low <= ~data_in[7];
                            state <= WRITE_DRIVE;
                        end
                    end else begin
                        ack_err <= 1'b1;
                        state <= ERROR;
                    end
                end

                WRITE_DRIVE: begin
                    busy <= 1'b1;
                    scl_drive_low <= 1'b0;
                    sda_drive_low <= ~tx_shift[7];
                    state <= WRITE_SAMPLE;
                end

                WRITE_SAMPLE: begin
                    busy <= 1'b1;
                    scl_drive_low <= 1'b0;
                    if (tx_shift[7] && !sda_in) begin
                        ack_err <= 1'b1;
                        state <= ARBITRATION_LOST;
                    end else if (bit_count == 3'd0) begin
                        sda_drive_low <= 1'b0;
                        state <= ACK2;
                    end else begin
                        tx_shift <= {tx_shift[6:0], 1'b0};
                        bit_count <= bit_count - 1'b1;
                        scl_drive_low <= 1'b1;
                        sda_drive_low <= ~tx_shift[6];
                        state <= WRITE_DRIVE;
                    end
                end

                READ: begin
                    busy <= 1'b1;
                    scl_drive_low <= 1'b0;
                    sda_drive_low <= 1'b0;
                    rx_shift <= {rx_shift[6:0], sda_in};
                    if (bit_count == 3'd0) begin
                        data_out <= {rx_shift[6:0], sda_in};
                        state <= ACK2;
                    end else begin
                        bit_count <= bit_count - 1'b1;
                    end
                end

                ACK2: begin
                    busy <= 1'b1;
                    scl_drive_low <= 1'b0;
                    sda_drive_low <= 1'b0;
                    if (!latched_rw && sda_in) begin
                        ack_err <= 1'b1;
                    end
                    state <= STOP_A;
                end

                STOP_A: begin
                    busy <= 1'b1;
                    scl_drive_low <= 1'b0;
                    sda_drive_low <= 1'b1;
                    if (!scl_in) begin
                        state <= CLOCK_STRETCH;
                    end else begin
                        state <= STOP_B;
                    end
                end

                CLOCK_STRETCH: begin
                    busy <= 1'b1;
                    scl_drive_low <= 1'b0;
                    sda_drive_low <= 1'b1;
                    if (scl_in) begin
                        state <= STOP_B;
                    end
                end

                STOP_B: begin
                    busy <= 1'b0;
                    scl_drive_low <= 1'b0;
                    sda_drive_low <= 1'b0;
                    state <= IDLE;
                end

                ARBITRATION_LOST: begin
                    busy <= 1'b0;
                    scl_drive_low <= 1'b0;
                    sda_drive_low <= 1'b0;
                    ack_err <= 1'b1;
                    state <= IDLE;
                end

                default: begin
                    busy <= 1'b0;
                    scl_drive_low <= 1'b0;
                    sda_drive_low <= 1'b0;
                    ack_err <= 1'b1;
                    state <= IDLE;
                end
            endcase
        end
    end

endmodule