import logging
import os
import re
import shlex
import subprocess
from pathlib import Path
import sys
import platform

logging.basicConfig(level=logging.INFO)
THEROCK_BIN_DIR_STR = os.getenv("THEROCK_BIN_DIR")
if THEROCK_BIN_DIR_STR is None:
    logging.info(
        "++ Error: env(THEROCK_BIN_DIR) is not set. Please set it before executing tests."
    )
    sys.exit(1)
THEROCK_BIN_DIR = Path(THEROCK_BIN_DIR_STR)
SCRIPT_DIR = Path(__file__).resolve().parent
THEROCK_DIR = SCRIPT_DIR.parent.parent.parent
THEROCK_TEST_DIR = Path(THEROCK_DIR) / "build"

ROCDECODE_TEST_PATH = str(
    Path(THEROCK_BIN_DIR).resolve().parent / "share" / "rocdecode" / "test"
)
if not os.path.isdir(ROCDECODE_TEST_PATH):
    logging.info(f"++ Error: rocdecode tests not found in {ROCDECODE_TEST_PATH}")
    sys.exit(1)
else:
    logging.info(f"++ INFO: rocdecode tests found in {ROCDECODE_TEST_PATH}")
env = os.environ.copy()


# set env variables required for tests
def setup_env(env):
    ROCM_PATH = Path(THEROCK_BIN_DIR).resolve().parent
    env["ROCM_PATH"] = str(ROCM_PATH)
    logging.info(f"++ rocdecode setting ROCM_PATH={ROCM_PATH}")
    if platform.system() == "Linux":
        HIP_LIB_PATH = Path(THEROCK_BIN_DIR).resolve().parent / "lib"
        logging.info(f"++ rocdecode setting LD_LIBRARY_PATH={HIP_LIB_PATH}")
        if "LD_LIBRARY_PATH" in env:
            env["LD_LIBRARY_PATH"] = f"{HIP_LIB_PATH}:{env['LD_LIBRARY_PATH']}"
        else:
            env["LD_LIBRARY_PATH"] = str(HIP_LIB_PATH)
        ROCM_SYSDEPS_LIB_PATH = ROCM_PATH / "lib" / "rocm_sysdeps" / "lib"
        LD_PRELOAD_LIBS = [
            str(ROCM_SYSDEPS_LIB_PATH / "librocm_sysdeps_va.so.2"),
            str(ROCM_SYSDEPS_LIB_PATH / "librocm_sysdeps_va-drm.so.2"),
        ]
        LD_PRELOAD_VALUE = ":".join(LD_PRELOAD_LIBS)
        logging.info(f"++ rocdecode setting LD_PRELOAD={LD_PRELOAD_VALUE}")
        env["LD_PRELOAD"] = LD_PRELOAD_VALUE
        logging.info(f"++ rocdecode setting LIBVA_DRIVERS_PATH={ROCM_SYSDEPS_LIB_PATH}")
        env["LIBVA_DRIVERS_PATH"] = str(ROCM_SYSDEPS_LIB_PATH)
    else:
        logging.info(f"++ rocdecode tests only supported on Linux")
        sys.exit(0)


def execute_tests(env):
    ROCDECODE_TEST_DIR = Path(THEROCK_TEST_DIR) / "rocdecode-test"

    ROCDECODE_TEST_DIR.mkdir(parents=True, exist_ok=True)

    # rocdecode tests are shipped as CMake source and must be built on the target
    # machine. This serves two purposes:
    # 1. Verifies that the installed rocdecode headers and libraries are functional.
    # 2. Some test dependencies (e.g. video codec libraries) are not bundled in the
    #    TheRock artifacts and must be linked from the system at build time.
    cmd = [
        "cmake",
        "-GNinja",
        "-DENABLE_EXTENDED_TESTS=ON",
        ROCDECODE_TEST_PATH,
    ]
    logging.info(f"++ Exec [{ROCDECODE_TEST_DIR}]$ {shlex.join(cmd)}")
    subprocess.run(cmd, cwd=ROCDECODE_TEST_DIR, check=True, env=env)

    cmd = [
        "ctest",
        "-N",
    ]
    logging.info(f"++ Exec [{ROCDECODE_TEST_DIR}]$ {shlex.join(cmd)}")
    ctest_list = subprocess.run(
        cmd,
        cwd=ROCDECODE_TEST_DIR,
        check=True,
        env=env,
        capture_output=True,
        text=True,
    )
    logging.info(ctest_list.stdout)
    match = re.search(r"Total Tests:\s*(\d+)", ctest_list.stdout)
    if match is None:
        raise RuntimeError(
            "Failed to determine CTest test count from `ctest -N` output"
        )
    if int(match.group(1)) == 0:
        raise RuntimeError("CTest discovered zero rocdecode tests")

    cmd = [
        "ctest",
        "--extra-verbose",
        "--output-on-failure",
    ]
    logging.info(f"++ Exec [{ROCDECODE_TEST_DIR}]$ {shlex.join(cmd)}")
    subprocess.run(cmd, cwd=ROCDECODE_TEST_DIR, check=True, env=env)


if __name__ == "__main__":
    setup_env(env)
    execute_tests(env)