`timescale 1ns / 1ps

module tb_digital_lock_keypad_fsm;
    reg clk;
    reg rst;
    reg key_valid;
    reg [3:0] key_code;
    wire unlock;

    integer errors = 0;
    integer tests_run = 0;
    integer i;
    integer rand_sel;
    reg [1:0] golden_state;

    localparam [1:0] S_IDLE = 2'd0;
    localparam [1:0] S_2    = 2'd1;
    localparam [1:0] S_24   = 2'd2;
    localparam [1:0] S_246  = 2'd3;
    localparam integer NUM_RANDOM_TESTS = 200;

    digital_lock_keypad_fsm uut (
        .clk(clk),
        .rst(rst),
        .key_valid(key_valid),
        .key_code(key_code),
        .unlock(unlock)
    );

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

    function [1:0] expected_next_state;
        input [1:0] state;
        input       valid;
        input [3:0] code;
        begin
            if (!valid) begin
                expected_next_state = state;
            end else if (code == 4'hF) begin
                expected_next_state = S_IDLE;
            end else begin
                case (state)
                    S_IDLE: expected_next_state = (code == 4'd2) ? S_2 : S_IDLE;
                    S_2:    expected_next_state = (code == 4'd4) ? S_24 :
                                                  (code == 4'd2) ? S_2  : S_IDLE;
                    S_24:   expected_next_state = (code == 4'd6) ? S_246 :
                                                  (code == 4'd2) ? S_2   : S_IDLE;
                    S_246:  expected_next_state = (code == 4'd2) ? S_IDLE : S_IDLE;
                    default: expected_next_state = S_IDLE;
                endcase
            end
        end
    endfunction

    function expected_unlock;
        input [1:0] state;
        input       valid;
        input [3:0] code;
        begin
            expected_unlock = valid && (state == S_246) && (code == 4'd2);
        end
    endfunction

    function [3:0] random_legal_key;
        input integer sel;
        begin
            case (sel % 11)
                0:  random_legal_key = 4'd0;
                1:  random_legal_key = 4'd1;
                2:  random_legal_key = 4'd2;
                3:  random_legal_key = 4'd3;
                4:  random_legal_key = 4'd4;
                5:  random_legal_key = 4'd5;
                6:  random_legal_key = 4'd6;
                7:  random_legal_key = 4'd7;
                8:  random_legal_key = 4'd8;
                9:  random_legal_key = 4'd9;
                default: random_legal_key = 4'hF;
            endcase
        end
    endfunction

    task apply_cycle;
        input [255:0] tag;
        input         next_rst;
        input         next_key_valid;
        input [3:0]   next_key_code;
        reg [1:0]     exp_state;
        reg           exp_unlock;
        begin
            rst = next_rst;
            key_valid = next_key_valid;
            key_code = next_key_code;

            if (next_rst) begin
                exp_state = S_IDLE;
                exp_unlock = 1'b0;
            end else begin
                exp_state = expected_next_state(golden_state, next_key_valid, next_key_code);
                exp_unlock = expected_unlock(golden_state, next_key_valid, next_key_code);
            end

            @(posedge clk);
            #1;

            tests_run = tests_run + 1;

            if (unlock !== exp_unlock) begin
                $display("ERROR [%s]: unlock mismatch. Expected %b, got %b (golden_state=%0d, key_valid=%0b, key_code=%0h)",
                         tag, exp_unlock, unlock, golden_state, next_key_valid, next_key_code);
                errors = errors + 1;
            end

            golden_state = exp_state;
        end
    endtask

    initial begin
        rst = 1'b0;
        key_valid = 1'b0;
        key_code = 4'd0;
        golden_state = S_IDLE;

        $display("===========================================");
        $display("  Digital Lock with Keypad FSM Testbench");
        $display("===========================================");

        // Reset behavior, including reset priority over an input key.
        apply_cycle("reset_idle",        1'b1, 1'b0, 4'd0);
        apply_cycle("reset_with_key",    1'b1, 1'b1, 4'd2);
        apply_cycle("post_reset_idle",   1'b0, 1'b0, 4'd9);

        // Basic success case.
        apply_cycle("basic_2",           1'b0, 1'b1, 4'd2);
        apply_cycle("basic_4",           1'b0, 1'b1, 4'd4);
        apply_cycle("basic_6",           1'b0, 1'b1, 4'd6);
        apply_cycle("basic_unlock",      1'b0, 1'b1, 4'd2);
        apply_cycle("post_unlock_idle",  1'b0, 1'b0, 4'd0);

        // Wrong digit after zero matched digits.
        apply_cycle("wrong_idle",        1'b0, 1'b1, 4'd7);
        apply_cycle("idle_after_wrong",  1'b0, 1'b0, 4'd0);

        // Wrong digit after one matched digit.
        apply_cycle("one_match_2",       1'b0, 1'b1, 4'd2);
        apply_cycle("one_match_wrong",   1'b0, 1'b1, 4'd5);

        // Wrong digit after two matched digits.
        apply_cycle("two_match_2",       1'b0, 1'b1, 4'd2);
        apply_cycle("two_match_4",       1'b0, 1'b1, 4'd4);
        apply_cycle("two_match_wrong",   1'b0, 1'b1, 4'd9);

        // Wrong digit after three matched digits.
        apply_cycle("three_match_2",     1'b0, 1'b1, 4'd2);
        apply_cycle("three_match_4",     1'b0, 1'b1, 4'd4);
        apply_cycle("three_match_6",     1'b0, 1'b1, 4'd6);
        apply_cycle("three_match_wrong", 1'b0, 1'b1, 4'd8);

        // Restart-on-2 overlap behavior.
        apply_cycle("overlap_2",         1'b0, 1'b1, 4'd2);
        apply_cycle("overlap_4",         1'b0, 1'b1, 4'd4);
        apply_cycle("overlap_restart",   1'b0, 1'b1, 4'd2);
        apply_cycle("overlap_4_again",   1'b0, 1'b1, 4'd4);
        apply_cycle("overlap_6",         1'b0, 1'b1, 4'd6);
        apply_cycle("overlap_unlock",    1'b0, 1'b1, 4'd2);

        // CLEAR behavior.
        apply_cycle("clear_2",           1'b0, 1'b1, 4'd2);
        apply_cycle("clear_4",           1'b0, 1'b1, 4'd4);
        apply_cycle("clear_key",         1'b0, 1'b1, 4'hF);
        apply_cycle("clear_idle",        1'b0, 1'b0, 4'd3);
        apply_cycle("clear_retry_2",     1'b0, 1'b1, 4'd2);
        apply_cycle("clear_retry_4",     1'b0, 1'b1, 4'd4);
        apply_cycle("clear_retry_6",     1'b0, 1'b1, 4'd6);
        apply_cycle("clear_unlock",      1'b0, 1'b1, 4'd2);

        // key_valid low must ignore key_code, including CLEAR-like values.
        apply_cycle("hold_2",            1'b0, 1'b1, 4'd2);
        apply_cycle("hold_invalid",      1'b0, 1'b0, 4'hF);
        apply_cycle("hold_4",            1'b0, 1'b1, 4'd4);
        apply_cycle("hold_invalid2",     1'b0, 1'b0, 4'd1);
        apply_cycle("hold_6",            1'b0, 1'b1, 4'd6);
        apply_cycle("hold_unlock",       1'b0, 1'b1, 4'd2);

        // Back-to-back successful attempts.
        apply_cycle("double_1_2",        1'b0, 1'b1, 4'd2);
        apply_cycle("double_1_4",        1'b0, 1'b1, 4'd4);
        apply_cycle("double_1_6",        1'b0, 1'b1, 4'd6);
        apply_cycle("double_1_unlock",   1'b0, 1'b1, 4'd2);
        apply_cycle("double_2_2",        1'b0, 1'b1, 4'd2);
        apply_cycle("double_2_4",        1'b0, 1'b1, 4'd4);
        apply_cycle("double_2_6",        1'b0, 1'b1, 4'd6);
        apply_cycle("double_2_unlock",   1'b0, 1'b1, 4'd2);

        // Randomized regression using the same scoreboard model.
        for (i = 0; i < NUM_RANDOM_TESTS; i = i + 1) begin
            rand_sel = $urandom_range(0, 10);
            apply_cycle("random", 1'b0, $urandom_range(0, 1), random_legal_key(rand_sel));
        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