`timescale 1ns/1ps

module tb_i2c_master;
    reg clk;
    reg rst_n;
    reg [6:0] addr;
    reg rw;
    reg [7:0] data_in;
    reg start_tx;

    tri1 scl;
    tri1 sda;

    reg slave_scl_low;
    reg slave_sda_low;
    reg nack_first_ack;
    reg arbitration_request;
    reg [7:0] read_value;
    reg saw_arbitration_lost;

    wire [7:0] data_out;
    wire busy;
    wire ack_err;

    integer errors;
    integer checks;

    localparam IDLE             = 4'd0;
    localparam ADDR_DRIVE       = 4'd3;
    localparam ADDR_SAMPLE      = 4'd4;
    localparam ACK1             = 4'd5;
    localparam READ             = 4'd8;
    localparam ACK2             = 4'd9;
    localparam STOP_A           = 4'd10;
    localparam CLOCK_STRETCH    = 4'd12;
    localparam ARBITRATION_LOST = 4'd13;

    assign scl = slave_scl_low ? 1'b0 : 1'bz;
    assign sda = slave_sda_low ? 1'b0 : 1'bz;

    i2c_master dut (
        .clk(clk),
        .rst_n(rst_n),
        .addr(addr),
        .rw(rw),
        .data_in(data_in),
        .start_tx(start_tx),
        .data_out(data_out),
        .busy(busy),
        .ack_err(ack_err),
        .scl(scl),
        .sda(sda)
    );

    always #5 clk = ~clk;

    always @(*) begin
        slave_sda_low = 1'b0;

        if (dut.state == ACK1) begin
            if (!nack_first_ack) begin
                slave_sda_low = 1'b1;
            end
        end else if (dut.state == ACK2 && !dut.latched_rw) begin
            slave_sda_low = 1'b1;
        end else if (dut.state == READ) begin
            slave_sda_low = ~read_value[dut.bit_count];
        end

        if (arbitration_request &&
            (dut.state == ADDR_DRIVE || dut.state == ADDR_SAMPLE) &&
            dut.tx_shift[7]) begin
            slave_sda_low = 1'b1;
        end
    end

    always @(posedge clk) begin
        if (dut.state == ARBITRATION_LOST) begin
            saw_arbitration_lost <= 1'b1;
        end
    end

    task check;
        input cond;
        input string message;
        begin
            checks = checks + 1;
            if (!cond) begin
                errors = errors + 1;
                $display("CHECK FAILED: %0s", message);
            end
        end
    endtask

    task reset_dut;
        begin
            rst_n = 1'b0;
            start_tx = 1'b0;
            addr = 7'h00;
            rw = 1'b0;
            data_in = 8'h00;
            nack_first_ack = 1'b0;
            arbitration_request = 1'b0;
            slave_scl_low = 1'b0;
            read_value = 8'h00;
            repeat (3) @(posedge clk);
            rst_n = 1'b1;
            @(posedge clk);
        end
    endtask

    task launch_transfer;
        input [6:0] tx_addr;
        input tx_rw;
        input [7:0] tx_data;
        begin
            addr = tx_addr;
            rw = tx_rw;
            data_in = tx_data;
            start_tx = 1'b1;
            @(posedge clk);
            start_tx = 1'b0;
        end
    endtask

    task wait_for_idle;
        integer timeout;
        begin
            timeout = 0;
            while ((busy || dut.state != IDLE) && timeout < 200) begin
                @(posedge clk);
                timeout = timeout + 1;
            end
            check(timeout < 200, "I2C transfer timed out");
        end
    endtask

    initial begin
        clk = 1'b0;
        rst_n = 1'b0;
        addr = 7'h00;
        rw = 1'b0;
        data_in = 8'h00;
        start_tx = 1'b0;
        slave_scl_low = 1'b0;
        nack_first_ack = 1'b0;
        arbitration_request = 1'b0;
        read_value = 8'h00;
        saw_arbitration_lost = 1'b0;
        errors = 0;
        checks = 0;

        reset_dut();

        launch_transfer(7'h42, 1'b0, 8'hA5);
        wait_for_idle();
        check(ack_err == 1'b0, "write transaction should complete without ack_err");
        check(busy == 1'b0, "busy should deassert after write");
        check(scl === 1'b1 && sda === 1'b1, "bus should return high after STOP");

        read_value = 8'h3C;
        launch_transfer(7'h42, 1'b1, 8'h00);
        wait_for_idle();
        check(data_out == 8'h3C, "read transaction should capture slave data");
        check(ack_err == 1'b0, "read transaction should complete without ack_err");

        nack_first_ack = 1'b1;
        launch_transfer(7'h42, 1'b0, 8'h55);
        wait_for_idle();
        check(ack_err == 1'b1, "NACK on first ACK should set ack_err");
        nack_first_ack = 1'b0;

        fork
            begin
                wait (dut.state == STOP_A);
                slave_scl_low = 1'b1;
                repeat (3) @(posedge clk);
                slave_scl_low = 1'b0;
            end
        join_none
        launch_transfer(7'h42, 1'b0, 8'hC3);
        wait_for_idle();
        check(ack_err == 1'b0, "clock stretch should not force ack_err");

        saw_arbitration_lost = 1'b0;
        fork
            begin
                wait (dut.state == ADDR_SAMPLE && dut.tx_shift[7]);
                arbitration_request = 1'b1;
                repeat (2) @(posedge clk);
                arbitration_request = 1'b0;
            end
        join_none
        launch_transfer(7'h60, 1'b0, 8'hF0);
        wait_for_idle();
        check(saw_arbitration_lost, "arbitration loss path should be reachable");
        check(ack_err == 1'b1, "arbitration loss should set ack_err");

        if (errors == 0) begin
            $display("Hint: Total mismatched samples is 0 out of %0d samples", checks);
        end else begin
            $display("Hint: Total mismatched samples is %0d out of %0d samples", errors, checks);
        end
        $display("Mismatches: %0d in %0d samples", errors, checks);
        $finish;
    end

endmodule