# Copyright 2023 Google LLC

# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at

#     https://www.apache.org/licenses/LICENSE-2.0

# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

"""Rules to build matcha test targets"""

load(
    "//rules:matcha.bzl",
    "DV_CORE_TARGETS",
    "flash_binary",
    "opentitan_rom_binary",
)

def matcha_verilator_test(
        name,
        sec_flash_binary,
        smc_flash_binary = "@//sw/device/tests/smc:simple_smc",
        ml_flash_binary = None,
        rom_img = "@//sw/device/lib/testing/test_rom:matcha_test_rom_sim_verilator_scr_vmem",
        otp_img = "@//hw/top_matcha/data:otp_img_rma",
        verilator_testbench = "@//hw:matcha_verilator",
        timeout = "moderate",
        tags = [],
        **kwargs):
    """A helper macro for generating matcha verilator tests.

    This macro is mostly a wrapper around `sh_test`.

    Args:
      name: The name of this rule.
      rom_img: The boot rom image target to use.
      otp_img: The otp image target to use.
      sec_flash_binary: sec_flash_binary target to be loaded to ram.
      smc_flash_binary: smc_flash_binary target to be loaded to ram_smc.
      ml_flash_binary: ml_flash_binary target to be loaded to ml_dmem.
      verilator_testbench: The verilated HW target.
      timeout: test timeout. Default to moderate (5 min).
      tags: Additional test tags to attach.
      **kwargs: Arguments to forward to `sh_test`.

    This macro emits the following rule:
      sh_test                named: {name}
    """

    verilator_suffix = "_sim_verilator_vmem"

    sec_verilator_binary = sec_flash_binary + verilator_suffix
    smc_verilator_binary = smc_flash_binary + verilator_suffix
    ml_verilator_binary = ml_flash_binary if ml_flash_binary else None

    sh_test_tags = ["verilator"] + tags
    sh_test_runner = "@//util:run_chip_verilator_sim.sh"
    sh_test_args = [
        "$(location %s)" % (verilator_testbench),
        "$(location %s)" % (rom_img),
        "$(location %s)" % (sec_verilator_binary),
        "$(location %s)" % (otp_img),
        "$(location %s)" % (smc_verilator_binary),
    ]
    if ml_verilator_binary:
        sh_test_args.append("$(location %s)" % (ml_verilator_binary))

    sh_test_data = [
        verilator_testbench,
        rom_img,
        sec_verilator_binary,
        otp_img,
        smc_verilator_binary,
    ]
    if ml_verilator_binary:
        sh_test_data.append(ml_verilator_binary)

    native.sh_test(
        name = name,
        srcs = [sh_test_runner],
        args = sh_test_args,
        data = sh_test_data,
        timeout = timeout,
        tags = sh_test_tags,
        **kwargs
    )

def dv_params(
        # Base Parameters
        args = ["$(location {dvsim_config})"],
        data = [],
        local = True,
        otp = "//hw/top_matcha/data:otp_img_rma",
        rom = "//sw/device/lib/testing/test_rom:matcha_test_rom_sim_dv_scr_vmem",
        tags = [],
        timeout = "moderate",
        test_runner = "//util:dvsim_matcha_test_runner.sh",
        # DV-specific Parameters
        bootstrap_sw = False,  # Default to backdoor loading.
        dvsim_config = "@//hw/top_matcha/dv:chip_sim_cfg.hjson",
        **kwargs):
    """A macro to create DV sim parameters for Matcha dv tests.

    This macro emits a dictionary of parameters which are pasted into the DV
    simulation specific test rule.

    Args:
        args: Extra arguments to pass to the test runner (`dvsim.py`).
        data: Data dependencies of the test.
        local: Whether the test should be run locally without sandboxing.
        otp: The OTP image to use.
        rom: The ROM image to use.
        tags: The test tags to apply to the test rule.
        timeout: The timeout to apply to the test rule.
        test_runner: sh script to invoke the test.
        bootstrap_sw: Whether to load test image with bootstrap.
        dvsim_config: The dvsim.py Hjson config file for the toplevel.
        **kwargs: return.
    Returns:
        **kwargs: build parameters for DV targets.
    """
    required_args = [
        "-i",
        "chip_sw_{name}",
        "--",
    ]
    required_data = [
        dvsim_config,
        "@lowrisc_opentitan//util/dvsim",
        "@//hw:all_files",
    ]
    required_tags = ["dv"]
    kwargs.update(
        args = args + required_args,
        data = required_data + data,
        local = local,
        otp = otp,
        rom = rom,
        tags = required_tags + tags,
        test_runner = test_runner,
        timeout = timeout,
        bootstrap_sw = bootstrap_sw,
        dvsim_config = dvsim_config,
    )
    return kwargs

def _format_list(param_name, list1, datadict, **kwargs):
    """Concatenate and format list items.

    This is used to prepare substitutions in user-supplied args to the
    various test invocations (ie: the location of flash).
    Args:
        @param param_name: The name of the item in `datadict`.
        @param list1: A list of items to prepend to the list item from datadict.
        @param datadict: A dictionary of per-test parameters.
        @param **kwargs: Values to pass to the format function.
    Returns:
        list[str]
    """
    return [x.format(**kwargs) for x in list1 + datadict.pop(param_name, [])]

def matcha_dv_test(
        name,
        targets = ["dv"],
        args = [],
        data = [],
        test_in_rom = False,
        signed = False,
        dv = None,
        test_binary = None,
        per_device_deps = {"sim_dv": [DV_CORE_TARGETS.get("secure_core")]},
        word_size = 64,
        **kwargs):
    """A helper macro for generating Matcha functional tests.

    This macro is mostly a wrapper around flash_binary, but creates
    testing artifacts for each of the hardware targets in `targets`. The testing
    artifacts are then given to an `sh_test` rule which dispatches the test to
    the corresponding hardware target via //util/dvsim_matcha_test_runner.sh.
    Args:
      name: The name of this rule.
      targets: A list of hardware targets on which to dispatch tests.
      args: Extra arguments (in addition to those defined in the target-
                   specific parameter dictionary) to pass to the test runner.
      data: Extra data dependencies (in addition to those defined in the
                   target-specific parameter dictionary) needed while executing
                   the test.
      test_in_rom: Whether to run the test from ROM, runs from flash by
                          default.
      signed: Whether to sign the test image. Unsigned by default.
      dv: DV test parameters.
      test_binary: Use the named binary as the test program rather than building one from srcs/deps.
      per_device_deps: Target deps that are forwared to flash_binary. Default to build for secure
                       core DV only.
      word_size: The word size that is forwared to flash_binary. Default to 64 for secure core.
      **kwargs: Arguments to forward to `flash_binary`.

    This macro emits the following rules:
        flash_binary named: {name}_prog (and all emitted rules).
        sh_test                named: dv_{name}
    """

    # Generate flash artifacts for test.
    deps = kwargs.pop("deps", [])
    if test_in_rom:
        opentitan_rom_binary(
            name = name + "_rom_prog",
            deps = deps,
            **kwargs
        )
    if not test_binary:
        test_binary = name + "_prog"
        flash_binary(
            name = test_binary,
            signed = signed,
            per_device_deps = per_device_deps,
            word_size = word_size,
            deps = deps,
            **kwargs
        )

    all_tests = []

    target_params = {
        "sim_dv": dv_params() if not dv else dv,
    }

    for target, params in target_params.items():
        if target.split("_")[-1] not in targets:
            continue

        # Set test name.
        test_name = "{}_{}".format(name, target)
        if "manual" not in params.get("tags"):
            all_tests.append(test_name)

        sw_logs_db = []

        # Set flash image.
        if target in ["sim_dv"]:
            flash = "{}_{}_scr_vmem".format(test_binary, target)
            sw_logs_db.append("{}_{}_logs_db".format(test_binary, target))
        else:
            flash = "{}_{}_bin".format(test_binary, target)
        if signed:
            flash += "_signed"

        # If the (flash) test image is to be loaded via bootstrap in the DV
        # simulation environment, then we need to use a special VMEM image
        # that has been split into SPI flash frames. Currently, signed
        # images loaded via bootstrap in DV sim are not supported.
        # TODO: support signed bootstap images in DV sim.
        if target == "sim_dv" and params.pop("bootstrap_sw"):
            if test_in_rom:
                fail("Tests that run in ROM cannot be bootstrapped.")
            if signed:
                fail("A signed test cannot be bootstrapped in DV sim.")
            flash = "{}_{target}_frames_vmem".format(test_binary, target)

        # Set ROM image.
        rom = params.pop("rom")
        if test_in_rom:
            rom = "{}_rom_prog_{}_scr_vmem".format(name, target)
        if target in ["sim_dv"]:
            sw_logs_db.append(rom.replace("_scr_vmem", "_logs_db"))

        # Set OTP image.
        otp = params.pop("otp")

        # Success and failure strings.
        exit_strings_kwargs = {}

        # Retrieve remaining device-agnostic params.
        test_runner = params.pop("test_runner")

        # Retrieve device-specific params.
        dvsim_config = None
        if target == "sim_dv":
            dvsim_config = params.pop("dvsim_config")

        # Concatenate args / data passed into the matcha_dv_test macro
        # with args / data from device-specific params.
        # TODO(lowRISC/opentitan:#11779): remove this concatenation action
        concat_args = _format_list(
            "args",
            args,
            params,
            dvsim_config = dvsim_config,
            flash = flash,
            name = name,
            otp = otp,
            rom = rom,
            rom_kind = None,
            bitstream = None,
            **exit_strings_kwargs
        )
        concat_data = _format_list(
            "data",
            data,
            params,
            flash = flash,
            bitstream = None,
        )

        # Environment variables to pass into sh_test
        env = {}

        native.sh_test(
            name = test_name,
            srcs = [test_runner],
            args = concat_args,
            data = [
                flash,
                rom,
                otp,
            ] + concat_data + sw_logs_db,
            env = env,
            **params
        )