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

import argparse
import logging
import os
import resource
import shlex
import subprocess
import sys
import time
from pathlib import Path
from typing import Optional, Tuple

logging.basicConfig(
    level=logging.INFO, format="%(asctime)s - %(levelname)s: %(message)s"
)
logger = logging.getLogger(__name__)


def parse_arguments() -> argparse.Namespace:
    """
    Parse command-line arguments.
    """
    parser = argparse.ArgumentParser(
        description="Run ROCm Debug Agent tests.",
        formatter_class=argparse.RawDescriptionHelpFormatter,
        epilog="""
Environment Variables:
  OUTPUT_ARTIFACTS_DIR     ROCm tree root directory (optional)
  ROCM_PATH                ROCm installation path (used with --try-rocm-path)
        """,
    )
    parser.add_argument(
        "--try-rocm-path",
        action="store_true",
        help="Use ROCM_PATH environment variable as ROCm tree root (fatal error if path cannot be resolved; takes priority over OUTPUT_ARTIFACTS_DIR).",
    )
    parser.add_argument(
        "--max-retries",
        type=int,
        default=3,
        help="Maximum number of test retry attempts (default: 3).",
    )
    parser.add_argument(
        "--retry-delay",
        type=int,
        default=5,
        help="Base delay in seconds between retries (default: 5).",
    )

    return parser.parse_args()


def set_core_dump_limit() -> None:
    """
    Set core dump size limit to 0 (equivalent to ulimit -c 0).
    """
    try:
        resource.setrlimit(resource.RLIMIT_CORE, (0, 0))
        logger.info("[✓] Core dump limit set to 0.")
    except (ValueError, OSError) as e:
        logger.warning(f"[!] Failed to set core dump limit: {e}")
        logger.warning("Core files may be generated and consume disk space.")


def validate_path(path: Path, path_type: str, must_exist: bool = True) -> Path:
    """
    Validate and resolve a path.

    Args:
        path: Path to validate.
        path_type: Description of the path (for error messages).
        must_exist: Whether the path must exist.

    Returns:
        Resolved path.

    Raises:
        SystemExit: If path validation fails.
    """
    try:
        resolved = path.resolve(strict=must_exist)
        if must_exist and not resolved.exists():
            logger.error(f"[X] Error: {path_type} does not exist: {resolved}")
            sys.exit(1)
        return resolved
    except (OSError, RuntimeError) as e:
        logger.error(f"[X] Error: Could not resolve {path_type} '{path}': {e}")
        sys.exit(1)


def get_rocm_tree_root(try_rocm_path: bool = False) -> Path:
    """
    Determine the ROCm tree root directory.

    Priority:
    1. If try_rocm_path is True: Use ROCM_PATH environment variable if set
       (fatal error if path can't be resolved)
    2. If ROCM_PATH is not set: Use OUTPUT_ARTIFACTS_DIR environment variable if set
    3. Derive from script location (go up 2 directory levels)

    Args:
        try_rocm_path: If True, use ROCM_PATH with fatal error on resolution failure.

    Returns:
        Path to ROCm tree root.

    Raises:
        SystemExit: If ROCm tree root cannot be determined.
    """
    # Try ROCM_PATH first if requested.
    if try_rocm_path:
        rocm_path_str = os.getenv("ROCM_PATH")
        if rocm_path_str is not None:
            logger.info("Using ROCM_PATH as ROCm tree root.")
            # Fatal error if path can't be resolved when --try-rocm-path is used
            rocm_root = validate_path(Path(rocm_path_str), "ROCM_PATH", must_exist=True)
            return rocm_root

    # Check OUTPUT_ARTIFACTS_DIR.
    artifacts_dir_str = os.getenv("OUTPUT_ARTIFACTS_DIR")
    if artifacts_dir_str is not None:
        logger.info("Using OUTPUT_ARTIFACTS_DIR as ROCm tree root.")
        rocm_root = validate_path(
            Path(artifacts_dir_str), "OUTPUT_ARTIFACTS_DIR", must_exist=True
        )
    else:
        # Derive ROCm tree root from script location.
        # Script is at <root>/.github/scripts/test_rocr-debug-agent.py
        # Go up 2 levels: scripts/ -> .github/ -> root/
        script_path = Path(__file__).resolve()
        rocm_root = script_path.parent.parent.parent
        logger.info(f"Derived ROCm tree root from script location: {rocm_root}")

    return rocm_root


def get_default_paths(try_rocm_path: bool = False) -> Tuple[Path, Path]:
    """
    Get default paths for test binary and script.

    Args:
        try_rocm_path: If True, try ROCM_PATH environment variable first.

    Returns:
        Tuple of (test_bin, test_script) paths.

    Raises:
        SystemExit: If paths cannot be resolved or don't exist.
    """
    rocm_root = get_rocm_tree_root(try_rocm_path)

    # Both test binary and script are in <root>/tests/rocm-debug-agent/
    test_dir = rocm_root / "tests" / "rocm-debug-agent"
    test_bin = test_dir / "rocm-debug-agent-test"
    test_script = test_dir / "run-test.py"

    # Validate that the paths exist.
    if not test_bin.exists():
        logger.error(f"[X] Error: Test binary not found: {test_bin}")
        sys.exit(1)

    if not test_script.exists():
        logger.error(f"[X] Error: Test script not found: {test_script}")
        sys.exit(1)

    return test_bin, test_script


def print_section(
    title: str,
    border_char: str = "=",
    width: int = 80,
    center: bool = True,
    inline: bool = False,
    color: Optional[str] = None,
) -> None:
    """
    Print a visually distinct section header for console output.

    Supports two modes:
    1. Full multi-line section:
        ==============================
              Section Title
        ==============================
    2. Inline single-line section:
        -------- Section Title --------

    Args:
        title (str):
            The text to display inside the section header.
        border_char (str, optional):
            Character used for the border line (default: "=").
        width (int, optional):
            Total width of the header including borders (default: 80).
        center (bool, optional):
            Whether to center the title text for multi-line sections (default: True).
        inline (bool, optional):
            If True, print a single-line header with title inline (default: False).
        color (str, optional):
            ANSI color escape code applied to both title and borders (default: None, no color).

            Examples:
                "\033[92m" → Green
                "\033[93m" → Yellow
                "\033[94m" → Blue
                "\033[91m" → Red
                "\033[0m" resets color

    Example:
        print_section("EXCLUSIVE FAILING TESTS COMPARISON")
        print_section("Clang/Clang++", border_char="-", inline=True)
        print_section("WARNING", border_char="!", width=50, color="\033[93m")

    Notes:
        - Works in standard terminals and logs.
    """
    reset = "\033[0m"
    apply_color = (
        (lambda text: f"{color}{text}{reset}") if color else (lambda text: text)
    )

    # Always add a newline to the beginning of the section.
    logger.info("")

    if inline:
        # Prepare inline title.
        title_str = f" {title} "
        remaining = width - len(title_str)
        if remaining < 0:
            remaining = 0
        left = border_char * (remaining // 2)
        right = border_char * (remaining - len(left))
        logger.info(apply_color(f"{left}{title_str}{right}"))
    else:
        # Multi-line section style.
        border = border_char * width
        title_line = f"{title:^{width}}" if center else title
        logger.info(apply_color(border))
        logger.info(apply_color(title_line))
        logger.info(apply_color(border))


def run_tests(
    test_script: Path,
    test_bin_dir: Path,
    max_retries: int = 3,
    retry_delay: int = 5,
) -> None:
    """
    Runs the testsuite with a retry mechanism.

    Args:
        test_script: Path to test script.
        test_bin_dir: Directory containing test binaries.
        max_retries: Maximum number of retry attempts.
        retry_delay: Base delay in seconds between retries.

    Raises:
        SystemExit: If all retry attempts fail.
    """
    cmd = [sys.executable, str(test_script), str(test_bin_dir)]

    for attempt in range(1, max_retries + 1):
        print_section(f"Running tests (Attempt {attempt}/{max_retries})")

        logger.info(f"Exec [{test_bin_dir}]$ {shlex.join(cmd)}")

        start_time = time.perf_counter()
        try:
            subprocess.run(cmd, cwd=str(test_bin_dir), check=True)

            duration = time.perf_counter() - start_time

            print_section(
                f"[✓] Tests succeeded on attempt {attempt}. Duration: {duration:.2f}s"
            )
            return

        except subprocess.CalledProcessError as e:
            duration = time.perf_counter() - start_time
            logger.error(
                f"[X] Attempt {attempt}/{max_retries} failed with exit code {e.returncode} "
                f"after {duration:.2f}s"
            )

            if attempt < max_retries:
                # Exponential-ish backoff: multiply base delay by attempt number
                wait_time = attempt * retry_delay
                logger.info(f"Retrying in {wait_time}s...")
                time.sleep(wait_time)
            else:
                # We failed all the attempts.
                print_section(f"[X] All {max_retries} attempts failed.")
                sys.exit(1)


def main() -> None:
    """
    Main entry point for the script.
    """
    # Parse command-line arguments.
    args = parse_arguments()

    print_section("Path discovery")
    # Discover paths using automatic logic.
    test_bin, test_script = get_default_paths(try_rocm_path=args.try_rocm_path)
    test_bin_dir = test_bin.parent

    logger.info(f"Test Binary: {test_bin}")
    logger.info(f"Test Script: {test_script}")
    logger.info(f"Test Bin Dir: {test_bin_dir}")

    print_section("Disabling core file generation")

    # Set core dump limit to 0 (ulimit -c 0).
    set_core_dump_limit()

    # Run tests.
    run_tests(
        test_script=test_script,
        test_bin_dir=test_bin_dir,
        max_retries=args.max_retries,
        retry_delay=args.retry_delay,
    )


if __name__ == "__main__":
    main()