`timescale 1ns/1ps

package bh_assert;
    int bh_assert_pass_count = 0;
    int bh_assert_fail_count = 0;
    int bh_assert_warn_count = 0;

    localparam BH_ASSERT_LOG_LEVEL_FAIL = 0;
    localparam BH_ASSERT_LOG_LEVEL_ASSERT = 1;
    localparam BH_ASSERT_LOG_LEVEL_WARN = 2;
    localparam BH_ASSERT_LOG_LEVEL_INFO = 3;

    int bh_assert_log_level = BH_ASSERT_LOG_LEVEL_INFO;  // 0 = errors only, 1 = all assertions, 2 = warnings, 3 = info

    localparam COLOR_RED = "\033[31m";
    localparam COLOR_YELLOW = "\033[33m";
    localparam COLOR_GREEN = "\033[32m";
    localparam COLOR_NORMAL = "\033[0;39m";

    task bh_assert_equal(logic [127:0] val, logic [127:0] expected, string description);
        // display results
        $timeformat(-9, 2, " ns", 20);
        $display(
            "%t: %s: %8d %s %8d - %s",
            $time,
            val === expected ? {COLOR_GREEN, "PASS", COLOR_NORMAL} : {COLOR_RED, "FAIL", COLOR_NORMAL},
            val,
            val === expected ? "==" : "!=",
            expected,
            description
        );

        // update statistics
        if (val === expected) begin
            bh_assert_pass_count = bh_assert_pass_count + 1;
        end else begin
            bh_assert_fail_count = bh_assert_fail_count + 1;
        end
    endtask

    task bh_assert_stats;
        $timeformat(-9, 2, " ns", 20);
        $display("%t: DONE: %1d pass, %1d fail, %1d warn", $time, bh_assert_pass_count, bh_assert_fail_count, bh_assert_warn_count);
        if (bh_assert_pass_count + bh_assert_fail_count == 0) begin
            $display("%sERROR%s: no assertions found", COLOR_YELLOW, COLOR_NORMAL);
            // TODO: error in a better way?
        end else if (bh_assert_fail_count > 0) begin
            $display("%sERROR%s: some tests failed", COLOR_RED, COLOR_NORMAL);
        end else begin
            $display("%sSUCCESS%s: all tests passed", COLOR_GREEN, COLOR_NORMAL);
        end
    endtask

    task bh_info(string description);
        if (bh_assert_log_level >= BH_ASSERT_LOG_LEVEL_INFO) begin
            $timeformat(-9, 2, " ns", 20);
            $display("%t: INFO: %s", $time, description);
        end
    endtask

    task bh_warn(string description);
        if (bh_assert_log_level >= BH_ASSERT_LOG_LEVEL_WARN) begin
            $timeformat(-9, 2, " ns", 20);
            $display("%t: %sWARN%s: %s", $time, COLOR_YELLOW, COLOR_NORMAL, description);
        end
        bh_assert_warn_count = bh_assert_warn_count + 1;
    endtask

endpackage