#!/usr/bin/env python3
# Copyright (c) Advanced Micro Devices, Inc.
# SPDX-License-Identifier:  MIT
"""Run the rocprofiler-compute project's installed CTest suite.

This runner is installed with the rocprofiler-compute test payload and executes
the CTest suite generated by that project. It is intended for CI and
super-repo test orchestration that installs the project's test component.
For custom install layouts, pass both the ROCm prefix and CTest payload path
explicitly.
"""

import argparse
import logging
import os
import shlex
import subprocess
from pathlib import Path
from typing import Dict, List, Optional

PROJECT_NAME = "rocprofiler-compute"

EXCLUDED_TESTS = [
    "test_profile_live_attach_detach",
]

QUICK_TESTS = [
    "test_utils",
    "test_num_xcds_cli_output",
    "test_num_xcds_spec_class",
    "test_L1_cache_counters",
    "test_analyze_workloads",
    "test_analyze_commands",
    "test_metric_validation",
    "test_profile_iteration_multiplexing_1",
]

logging.basicConfig(level=logging.INFO)

HELP_EPILOG = f"""\
This script runs the {PROJECT_NAME} project test suite from an installed CTest
payload. In the CI/super-repo install layout, it discovers:

  <rocm-prefix>/<libexecdir>/{PROJECT_NAME}/CTestTestfile.cmake

For custom install layouts, pass --ctest-dir for the top-level CTest payload
directory and --rocm-path for the ROCm install prefix used by the tests.
"""


def format_command(cmd: List[str]) -> str:
    return " ".join(shlex.quote(str(arg)) for arg in cmd)


def path_from_env(name: str) -> Optional[Path]:
    value = os.getenv(name)
    if not value:
        return None
    return Path(value)


def installed_ctest_dir(script_path: Path) -> Optional[Path]:
    if script_path.parent.name == PROJECT_NAME:
        return script_path.parent.resolve()
    return None


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(
        description=(
            f"Run the {PROJECT_NAME} project test suite from its installed "
            "CTest payload."
        ),
        epilog=HELP_EPILOG,
        formatter_class=argparse.RawDescriptionHelpFormatter,
    )
    parser.add_argument(
        "--rocm-path",
        type=Path,
        default=path_from_env("ROCM_PATH"),
        help=(
            "ROCm install prefix used by the rocprofiler-compute tests. "
            "Defaults to ROCM_PATH or the installed runner/payload location."
        ),
    )
    parser.add_argument(
        "--ctest-dir",
        "--test-dir",
        dest="ctest_dir",
        type=Path,
        default=path_from_env("ROCPROFILER_COMPUTE_TEST_DIR"),
        help=(
            "Top-level installed rocprofiler-compute CTest payload directory "
            "containing CTestTestfile.cmake. Defaults to "
            "ROCPROFILER_COMPUTE_TEST_DIR or this runner's installed location."
        ),
    )
    return parser.parse_args()


def get_rocm_path(
    script_path: Path, rocm_path: Optional[Path], ctest_dir: Optional[Path]
) -> Path:
    if rocm_path is not None:
        return rocm_path.resolve()

    # Installed at <prefix>/<libexecdir>/rocprofiler-compute.
    installed_dir = installed_ctest_dir(script_path)
    if installed_dir is not None:
        return installed_dir.parents[1].resolve()

    if ctest_dir is not None:
        resolved_ctest_dir = ctest_dir.resolve()
        if resolved_ctest_dir.name == PROJECT_NAME:
            return resolved_ctest_dir.parents[1].resolve()

    raise RuntimeError("ROCM_PATH is required when the runner is not installed.")


def prepend_env_path(env: Dict[str, str], key: str, paths: List[Path]) -> None:
    existing = env.get(key)
    values = [str(path) for path in paths]
    if existing:
        values.append(existing)
    env[key] = os.pathsep.join(values)


def rocm_lib_dirs(rocm_path: Path) -> List[Path]:
    candidates = [rocm_path / "lib", rocm_path / "lib64"]
    lib_dirs = [path for path in candidates if path.is_dir()]
    if not lib_dirs:
        raise FileNotFoundError(
            f"Could not find ROCm library directory under {rocm_path}"
        )
    sysdeps_dirs = [
        lib_dir / "rocm_sysdeps" / "lib"
        for lib_dir in lib_dirs
        if (lib_dir / "rocm_sysdeps" / "lib").is_dir()
    ]
    return lib_dirs + sysdeps_dirs


def setup_env(rocm_path: Path) -> Dict[str, str]:
    env = os.environ.copy()
    env["ROCM_PATH"] = str(rocm_path)

    # Multi-GPU CI runners inject a GPU_DEVICE_ORDINAL alongside
    # HIP_VISIBLE_DEVICES via their GPU-isolation env-file. Newer HIP treats the
    # mismatch as a fatal agent-visibility conflict and exposes zero GPU agents,
    # so profiling collects no kernel data. HIP_VISIBLE_DEVICES supersedes it, so
    # drop the ordinal. Mirrors the fix in TheRock's test_rocprofiler_sdk.py.
    if env.get("HIP_VISIBLE_DEVICES"):
        env.pop("GPU_DEVICE_ORDINAL", None)

    prepend_env_path(env, "PATH", [rocm_path / "bin"])
    prepend_env_path(
        env,
        "LD_LIBRARY_PATH",
        rocm_lib_dirs(rocm_path),
    )

    return env


def build_ctest_command() -> List[str]:
    shard_index = int(os.getenv("SHARD_INDEX", "1")) - 1
    total_shards = int(os.getenv("TOTAL_SHARDS", "1"))

    cmd = [
        "ctest",
        "--output-on-failure",
        "--verbose",
        "--exclude-regex",
        "|".join(EXCLUDED_TESTS),
        "--tests-information",
        f"{shard_index},,{total_shards}",
    ]

    if os.getenv("TEST_TYPE", "full") == "quick":
        cmd.extend(["--tests-regex", "|".join(QUICK_TESTS)])

    return cmd


def main() -> None:
    args = parse_args()
    script_path = Path(__file__).resolve()
    rocm_path = get_rocm_path(script_path, args.rocm_path, args.ctest_dir)
    ctest_dir = (
        args.ctest_dir.resolve()
        if args.ctest_dir is not None
        else installed_ctest_dir(script_path) or rocm_path / "libexec" / PROJECT_NAME
    )
    ctest_file = ctest_dir / "CTestTestfile.cmake"
    if not ctest_file.is_file():
        raise FileNotFoundError(
            f"Could not find rocprofiler-compute CTest payload: {ctest_file}"
        )
    env = setup_env(rocm_path)
    cmd = build_ctest_command()

    logging.info("++ Exec [%s]$ %s", ctest_dir, format_command(cmd))
    subprocess.run(cmd, cwd=ctest_dir, check=True, env=env)


if __name__ == "__main__":
    main()