# Copyright Advanced Micro Devices, Inc.
# SPDX-License-Identifier: MIT

import logging
import os
import shlex
import subprocess
from pathlib import Path

THEROCK_BIN_DIR = os.getenv("THEROCK_BIN_DIR")
AMDGPU_FAMILIES = os.getenv("AMDGPU_FAMILIES")
SCRIPT_DIR = Path(__file__).resolve().parent
THEROCK_DIR = SCRIPT_DIR.parent.parent.parent

# GTest sharding
SHARD_INDEX = os.getenv("SHARD_INDEX", 1)
TOTAL_SHARDS = os.getenv("TOTAL_SHARDS", 1)
environ_vars = os.environ.copy()
# For display purposes in the GitHub Action UI, the shard array is 1th indexed. However for shard indexes, we convert it to 0th index.
environ_vars["GTEST_SHARD_INDEX"] = str(int(SHARD_INDEX) - 1)
environ_vars["GTEST_TOTAL_SHARDS"] = str(TOTAL_SHARDS)

# Enable GTest "brief" output: only show failures and the final results
environ_vars["GTEST_BRIEF"] = str(1)

# Some of our runtime kernel compilations have been relying on either ROCM_PATH being set, or ROCm being installed at
# /opt/rocm. Neither of these is true in TheRock so we need to supply ROCM_PATH to our tests.
ROCM_PATH = Path(THEROCK_BIN_DIR).resolve().parent
environ_vars["ROCM_PATH"] = str(ROCM_PATH)

logging.basicConfig(level=logging.INFO)

# If quick tests are enabled, we run quick tests only.
# Otherwise, we run the standard test suite.
test_type = os.getenv("TEST_TYPE", "standard")

# TODO(#2823): Re-enable test once flaky issue is resolved
TESTS_TO_IGNORE = ["unpack_util_test", "contamination_test", "map_util_test"]

test_subdir = ""
# CTest --timeout is per-test (seconds), not wall-clock for the whole shard.
# A value near the GitHub step limit lets one hung test burn the entire job (ROCM-24171).
# rocWMMA unit/gemm binaries should finish well under this on healthy runners; a stuck
# test then fails with a clear CTest timeout instead of an opaque workflow cancel at the
# GitHub Actions step limit.
_PER_TEST_TIMEOUT_FULL_SEC = 1800
_PER_TEST_TIMEOUT_QUICK_SEC = 720
timeout = str(_PER_TEST_TIMEOUT_FULL_SEC)
if test_type == "quick":
    # The emulator regression tests are very fast.
    # If we need something even faster we can use "/smoke" here.
    test_subdir = "/regression"
    timeout = str(_PER_TEST_TIMEOUT_QUICK_SEC)
elif test_type == "regression":
    test_subdir = "/regression"
    timeout = str(_PER_TEST_TIMEOUT_QUICK_SEC)

# Make per-device adjustments; per-GPU overrides can be added below.
# Example: if AMDGPU_FAMILIES == "gfx1153": ctest_parallelism = "2"
ctest_parallelism = "1"

cmd = [
    "ctest",
    "--test-dir",
    f"{THEROCK_BIN_DIR}/rocwmma{test_subdir}",
    "--output-on-failure",
    "--parallel",
    ctest_parallelism,
    "--timeout",
    timeout,
    "--exclude-regex",
    "|".join(TESTS_TO_IGNORE),
]
logging.info(f"++ Exec [{THEROCK_DIR}]$ {shlex.join(cmd)}")

subprocess.run(
    cmd,
    cwd=THEROCK_DIR,
    check=True,
    env=environ_vars,
)