`timescale 1ns / 1ps

module tb_rle_encoder;
    reg clk;
    reg rst;
    reg in_valid;
    reg [7:0] in_data;
    reg in_last;

    wire out_valid;
    wire [15:0] out_data;
    wire out_last;

    integer errors = 0;
    integer tests_run = 0;
    integer i;

    reg golden_have_run;
    reg golden_pending_flush;
    reg [7:0] golden_run_value;
    reg [7:0] golden_run_count;

    rle_encoder uut (
        .clk(clk),
        .rst(rst),
        .in_valid(in_valid),
        .in_data(in_data),
        .in_last(in_last),
        .out_valid(out_valid),
        .out_data(out_data),
        .out_last(out_last)
    );

    initial begin
        clk = 1'b0;
        forever #5 clk = ~clk;
    end

    task apply_cycle;
        input [255:0] tag;
        input next_rst;
        input next_in_valid;
        input [7:0] next_in_data;
        input next_in_last;
        reg exp_out_valid;
        reg [15:0] exp_out_data;
        reg exp_out_last;
        reg next_have_run;
        reg next_pending_flush;
        reg [7:0] next_run_value;
        reg [7:0] next_run_count;
        begin
            rst = next_rst;
            in_valid = next_in_valid;
            in_data = next_in_data;
            in_last = next_in_last;

            exp_out_valid = 1'b0;
            exp_out_data = 16'h0000;
            exp_out_last = 1'b0;

            next_have_run = golden_have_run;
            next_pending_flush = golden_pending_flush;
            next_run_value = golden_run_value;
            next_run_count = golden_run_count;

            if (next_rst) begin
                next_have_run = 1'b0;
                next_pending_flush = 1'b0;
                next_run_value = 8'h00;
                next_run_count = 8'h00;
            end else if (golden_pending_flush) begin
                exp_out_valid = 1'b1;
                exp_out_data = {golden_run_count, golden_run_value};
                exp_out_last = 1'b1;
                next_have_run = 1'b0;
                next_pending_flush = 1'b0;
                next_run_value = 8'h00;
                next_run_count = 8'h00;
            end else if (!golden_have_run) begin
                if (next_in_valid) begin
                    if (next_in_last) begin
                        exp_out_valid = 1'b1;
                        exp_out_data = {8'd1, next_in_data};
                        exp_out_last = 1'b1;
                        next_have_run = 1'b0;
                    end else begin
                        next_have_run = 1'b1;
                        next_run_value = next_in_data;
                        next_run_count = 8'd1;
                    end
                end
            end else begin
                if (next_in_valid) begin
                    if (next_in_data == golden_run_value) begin
                        if (golden_run_count == 8'd254) begin
                            exp_out_valid = 1'b1;
                            exp_out_data = {8'd255, golden_run_value};
                            exp_out_last = next_in_last;
                            next_have_run = 1'b0;
                            next_pending_flush = 1'b0;
                            next_run_value = 8'h00;
                            next_run_count = 8'h00;
                        end else if (next_in_last) begin
                            exp_out_valid = 1'b1;
                            exp_out_data = {golden_run_count + 8'd1, golden_run_value};
                            exp_out_last = 1'b1;
                            next_have_run = 1'b0;
                            next_pending_flush = 1'b0;
                            next_run_value = 8'h00;
                            next_run_count = 8'h00;
                        end else begin
                            next_run_count = golden_run_count + 8'd1;
                        end
                    end else begin
                        exp_out_valid = 1'b1;
                        exp_out_data = {golden_run_count, golden_run_value};
                        exp_out_last = 1'b0;
                        next_have_run = 1'b1;
                        next_run_value = next_in_data;
                        next_run_count = 8'd1;
                        next_pending_flush = next_in_last;
                    end
                end
            end

            @(posedge clk);
            #1;

            tests_run = tests_run + 1;

            if (out_valid !== exp_out_valid) begin
                $display("ERROR [%s]: out_valid mismatch. Expected %b, got %b",
                         tag, exp_out_valid, out_valid);
                errors = errors + 1;
            end

            if (exp_out_valid) begin
                if (out_data !== exp_out_data) begin
                    $display("ERROR [%s]: out_data mismatch. Expected %h, got %h",
                             tag, exp_out_data, out_data);
                    errors = errors + 1;
                end
                if (out_last !== exp_out_last) begin
                    $display("ERROR [%s]: out_last mismatch. Expected %b, got %b",
                             tag, exp_out_last, out_last);
                    errors = errors + 1;
                end
            end else if (out_last !== 1'b0) begin
                $display("ERROR [%s]: out_last should be 0 when out_valid is 0", tag);
                errors = errors + 1;
            end

            golden_have_run = next_have_run;
            golden_pending_flush = next_pending_flush;
            golden_run_value = next_run_value;
            golden_run_count = next_run_count;
        end
    endtask

    task idle_cycles;
        input [255:0] tag;
        input integer count;
        integer local_idx;
        begin
            for (local_idx = 0; local_idx < count; local_idx = local_idx + 1)
                apply_cycle(tag, 1'b0, 1'b0, 8'h00, 1'b0);
        end
    endtask

    task send_stream_same_value;
        input [255:0] tag;
        input [7:0] value;
        input integer length;
        integer idx;
        begin
            for (idx = 0; idx < length; idx = idx + 1)
                apply_cycle(tag, 1'b0, 1'b1, value, (idx == length - 1));
        end
    endtask

    task send_random_stream;
        input [255:0] tag;
        input integer length;
        integer idx;
        reg [7:0] cur_value;
        begin
            cur_value = $urandom;
            for (idx = 0; idx < length; idx = idx + 1) begin
                if ((idx != 0) && ($urandom_range(0, 3) == 0))
                    cur_value = $urandom;
                apply_cycle(tag, 1'b0, 1'b1, cur_value, (idx == length - 1));
                if ($urandom_range(0, 3) == 0)
                    idle_cycles("random_gap", 1);
            end
            if (golden_pending_flush)
                idle_cycles("random_flush", 1);
        end
    endtask

    initial begin
        rst = 1'b0;
        in_valid = 1'b0;
        in_data = 8'h00;
        in_last = 1'b0;

        golden_have_run = 1'b0;
        golden_pending_flush = 1'b0;
        golden_run_value = 8'h00;
        golden_run_count = 8'h00;

        $display("===========================================");
        $display("  Streaming Run-Length Encoder Testbench");
        $display("===========================================");

        // Reset behavior.
        apply_cycle("reset_0", 1'b1, 1'b0, 8'h00, 1'b0);
        apply_cycle("reset_1", 1'b1, 1'b1, 8'hAA, 1'b1);
        apply_cycle("post_reset", 1'b0, 1'b0, 8'h00, 1'b0);

        // Single-run stream.
        send_stream_same_value("single_run", 8'h33, 4);

        // Alternating values, ending with a different final pixel.
        apply_cycle("alt_0", 1'b0, 1'b1, 8'h10, 1'b0);
        apply_cycle("alt_1", 1'b0, 1'b1, 8'h20, 1'b0);
        apply_cycle("alt_2", 1'b0, 1'b1, 8'h30, 1'b1);
        idle_cycles("alt_flush", 1);

        // Idle cycle inside a run.
        apply_cycle("idle_mid_0", 1'b0, 1'b1, 8'h55, 1'b0);
        idle_cycles("idle_mid_gap", 2);
        apply_cycle("idle_mid_1", 1'b0, 1'b1, 8'h55, 1'b1);

        // Different final pixel creates a pending flush.
        apply_cycle("diff_last_0", 1'b0, 1'b1, 8'h44, 1'b0);
        apply_cycle("diff_last_1", 1'b0, 1'b1, 8'h44, 1'b0);
        apply_cycle("diff_last_2", 1'b0, 1'b1, 8'h99, 1'b1);
        idle_cycles("diff_last_flush", 1);

        // Exact 255-count run.
        send_stream_same_value("run255", 8'hA7, 255);

        // 255 rollover into a new run of the same value.
        for (i = 0; i < 257; i = i + 1)
            apply_cycle("rollover_same", 1'b0, 1'b1, 8'h5C, (i == 256));

        // Mixed stream with several runs.
        apply_cycle("mixed_0", 1'b0, 1'b1, 8'h11, 1'b0);
        apply_cycle("mixed_1", 1'b0, 1'b1, 8'h11, 1'b0);
        apply_cycle("mixed_2", 1'b0, 1'b1, 8'h22, 1'b0);
        apply_cycle("mixed_3", 1'b0, 1'b1, 8'h22, 1'b0);
        apply_cycle("mixed_4", 1'b0, 1'b1, 8'h22, 1'b0);
        apply_cycle("mixed_5", 1'b0, 1'b1, 8'h33, 1'b0);
        apply_cycle("mixed_6", 1'b0, 1'b1, 8'h44, 1'b1);
        idle_cycles("mixed_flush", 1);

        // Back-to-back streams.
        send_stream_same_value("stream_a", 8'h77, 3);
        apply_cycle("stream_b_0", 1'b0, 1'b1, 8'h88, 1'b0);
        apply_cycle("stream_b_1", 1'b0, 1'b1, 8'h99, 1'b1);
        idle_cycles("stream_b_flush", 1);

        // Reset during activity.
        apply_cycle("pre_reset_run", 1'b0, 1'b1, 8'hCC, 1'b0);
        apply_cycle("pre_reset_run2", 1'b0, 1'b1, 8'hCC, 1'b0);
        apply_cycle("mid_reset", 1'b1, 1'b0, 8'h00, 1'b0);
        apply_cycle("after_reset_idle", 1'b0, 1'b0, 8'h00, 1'b0);
        send_stream_same_value("after_reset_stream", 8'h12, 2);

        // Randomized regression.
        for (i = 0; i < 40; i = i + 1) begin
            send_random_stream("random_stream", $urandom_range(1, 30));
            idle_cycles("random_stream_gap", $urandom_range(0, 2));
        end

        $display("");
        $display("===========================================");
        $display("  Tests Run: %0d", tests_run);
        $display("===========================================");

        if (errors == 0) begin
            $display("TEST_RESULT: PASS");
        end else begin
            $display("TEST_RESULT: FAIL (%0d errors)", errors);
        end

        $finish;
    end
endmodule