#!/usr/bin/env python3
# Copyright Advanced Micro Devices, Inc., or its affiliates.
# SPDX-License-Identifier: MIT

"""Run rocroller tests in parallel shards.

Usage: run-tests-sharded PROFILE [BUILD_DIR]

Number of shards is auto-detected as ncores/2.
"""

import argparse
import os
import pty
import random
import select
import subprocess
import sys
import time
from dataclasses import dataclass
from pathlib import Path

# Add lib directory to path for rrtest import
sys.path.insert(0, str(Path(__file__).parent / "lib"))
from rrtest import get_test_commands


@dataclass
class ShardInfo:
    """Information about a running test shard."""

    process: subprocess.Popen
    master_fd: int
    shard_tag: str
    shard_index: int
    test_type: str
    executable: str
    output_buffer: str = ""
    start_time: float = 0.0


def get_gpu_arch():
    """Detect GPU architecture (e.g., gfx1201) if possible."""
    try:
        result = subprocess.run(
            ["rocminfo"],
            capture_output=True,
            text=True,
            check=True,
            timeout=5,
        )
        # Look for gfx architecture in output
        for line in result.stdout.splitlines():
            if "Name:" in line and "gfx" in line.lower():
                # Extract gfx architecture (e.g., gfx1201)
                parts = line.lower().split()
                for part in parts:
                    if part.startswith("gfx"):
                        return part.strip()
    except:
        pass
    return None


def get_available_cpus():
    """Get the number of CPUs available, respecting cgroups and affinity."""
    try:
        result = subprocess.run(
            ["nproc"], capture_output=True, text=True, check=True, timeout=1
        )
        return int(result.stdout.strip())
    except:
        pass

    return os.cpu_count() or 4


def get_shuffle_seed():
    """Return one random seed for all shards (gtest and catch2).

    Uses TEST_SHUFFLE_SEED env var if set (for reproducibility);
    otherwise picks a random seed.

    Returns:
        (seed, from_env): seed is an int in [1, 2**31-1]; from_env is True if
        the seed came from TEST_SHUFFLE_SEED.
    """
    master = os.environ.get("TEST_SHUFFLE_SEED")
    if master is not None:
        return int(master), True
    return random.SystemRandom().randint(1, 2**31 - 1), False


def start_shard(
    shard_index,
    base_cmd,
    test_type,
    num_shards,
    build_dir,
    shard_seed=None,
):
    """Start a shard of tests and return process info.

    Args:
        shard_index: Index of this shard (0-based)
        base_cmd: Base command list from get_test_commands (includes executable and filters)
        test_type: Framework name (e.g., "gtest", "catch2", "ctest")
        num_shards: Total number of shards
        build_dir: Build directory path
        shard_seed: Optional pre-computed random seed for this shard (used for gtest/catch2).
    """
    # Extract executable name from command
    executable = base_cmd[0]

    # Generate shard tag prefix from framework name (first letter uppercase)
    test_prefix = test_type[0].upper()
    shard_tag = f"[{test_prefix}{shard_index:02d}]"
    print(f"{shard_tag} Starting shard", flush=True)

    # Set environment variables
    env = os.environ.copy()

    # Parallelization
    env["OPENBLAS_NUM_THREADS"] = "2"
    env["OMP_NUM_THREADS"] = "2"

    # Debug logging for failed kernel generation
    env["AMD_COMGR_REDIRECT_LOGS"] = "stdout"

    # Start with base command from profile and add framework-specific sharding arguments
    cmd = base_cmd.copy()

    if test_type == "gtest":
        if shard_seed is None:
            raise ValueError("shard_seed is required for gtest")
        env["GTEST_TOTAL_SHARDS"] = str(num_shards)
        env["GTEST_SHARD_INDEX"] = str(shard_index)
        cmd.extend(
            [
                "--gtest_shuffle",
                f"--gtest_random_seed={shard_seed}",
                f"--gtest_output=xml:test_report/gtest_shard_{shard_index}.xml",
            ]
        )
    elif test_type == "catch2":
        if shard_seed is None:
            raise ValueError("shard_seed is required for catch2")
        cmd.extend(
            [
                "--order",
                "rand",
                "--shard-count",
                str(num_shards),
                "--shard-index",
                str(shard_index),
                "--rng-seed",
                str(shard_seed),
                "--reporter",
                f"junit::out=test_report/catch2_shard_{shard_index}.xml",
                "--reporter",
                "console::out=-::colour-mode=ansi",
            ]
        )
    elif test_type == "ctest":
        # CTest doesn't have built-in sharding.  Let CTest use threads instead.
        num_threads = 4
        cmd.extend(
            [
                "--parallel",
                str(num_threads),
                "--output-on-failure",
                "--output-junit",
                f"test_report/ctest_shard_{shard_index}.xml",
            ]
        )

    # Log environment variables we set (added or overridden)
    orig_env = os.environ
    set_vars = [
        (k, env[k])
        for k in sorted(env.keys())
        if k not in orig_env or orig_env.get(k) != env[k]
    ]
    if set_vars:
        print(f"{shard_tag} Environment:", flush=True)
        for k, v in set_vars:
            print(f"{shard_tag}   {k}={v}", flush=True)

    # Log the command being executed
    cmd_str = " ".join(cmd)
    print(f"{shard_tag} Running command in {build_dir}:", flush=True)
    print(f"{shard_tag}   {cmd_str}", flush=True)

    # Use pseudo-TTY to ensure line-buffered output from subprocess
    master_fd, slave_fd = pty.openpty()

    process = subprocess.Popen(
        cmd,
        env=env,
        cwd=build_dir,
        stdout=slave_fd,
        stderr=slave_fd,
        close_fds=True,
    )

    # Close slave fd in parent process
    os.close(slave_fd)

    return ShardInfo(
        process=process,
        master_fd=master_fd,
        shard_tag=shard_tag,
        shard_index=shard_index,
        test_type=test_type,
        executable=executable,
        output_buffer="",
        start_time=time.time(),
    )


def run_shards(shard_infos):
    """Run multiple shards and multiplex their I/O using select.

    Returns: (failed_shards, timing_info)
        failed_shards: list of (shard_index, test_type, exit_code) tuples
        timing_info: list of (shard_index, test_type, executable, elapsed_seconds) tuples
    """
    active_shards = {info.master_fd: info for info in shard_infos}
    failed_shards = []
    timing_info = []

    while active_shards:
        fds = list(active_shards.keys())

        ready, _, _ = select.select(fds, [], [], 0.1)

        for fd in ready:
            info = active_shards[fd]

            try:
                chunk = os.read(fd, 4096).decode("utf-8", errors="replace")
                if chunk:
                    info.output_buffer += chunk

                    # Process complete lines
                    while "\n" in info.output_buffer:
                        line, info.output_buffer = info.output_buffer.split("\n", 1)
                        line = line.rstrip("\r")
                        if line:
                            print(f"{info.shard_tag} {line}", flush=True)
            except OSError:
                # File descriptor closed or error
                pass

        # Check which processes have finished
        finished_fds = []
        for fd, info in active_shards.items():
            if info.process.poll() is not None:
                finished_fds.append(fd)

        # Handle finished processes
        for fd in finished_fds:
            info = active_shards[fd]

            # Read any remaining output
            try:
                while True:
                    chunk = os.read(fd, 4096).decode("utf-8", errors="replace")
                    if not chunk:
                        break
                    info.output_buffer += chunk
            except OSError:
                pass

            # Print any remaining buffered output
            if info.output_buffer:
                for line in info.output_buffer.split("\n"):
                    line = line.rstrip("\r")
                    if line:
                        print(f"{info.shard_tag} {line}", flush=True)

            # Close file descriptor
            os.close(fd)

            # Get return code and elapsed time
            return_code = info.process.wait()
            elapsed_time = time.time() - info.start_time

            # Record timing information
            timing_info.append(
                (info.shard_index, info.test_type, info.executable, elapsed_time)
            )

            if return_code == 0:
                print(
                    f"{info.shard_tag} Completed successfully in {elapsed_time:.1f}s",
                    flush=True,
                )
            else:
                print(
                    f"{info.shard_tag} Failed with exit code {return_code} after {elapsed_time:.1f}s",
                    flush=True,
                )
                failed_shards.append((info.shard_index, info.test_type, return_code))

            # Remove from active shards
            del active_shards[fd]

    return failed_shards, timing_info


def main():
    parser = argparse.ArgumentParser(
        description="Run rocroller tests in parallel shards"
    )
    parser.add_argument(
        "profile",
        help="Test profile to use from test-profiles.yaml (e.g., precheckin, codecov, smoke)",
    )
    parser.add_argument(
        "build_dir",
        nargs="?",
        default="build",
        help="Path to build directory (default: build)",
    )
    parser.add_argument(
        "--dry-run",
        action="store_true",
        help="List tests that would run with the given profile, but don't execute them",
    )

    args = parser.parse_args()

    # Auto-detect number of available CPUs (respecting cgroups) and use half
    num_cores = get_available_cpus()
    num_shards = max(1, num_cores // 2)

    # Cap shards to 8 on gfx1201 due to hanging issues
    gpu_arch = get_gpu_arch()
    if gpu_arch and "gfx1201" in gpu_arch:
        if num_shards > 8:
            print(f"Detected {gpu_arch}, capping shards to 8 (was {num_shards})")
            num_shards = 8

    # Cap shards to 48 maximum
    if num_shards > 48:
        print(f"Capping shards to 48 (was {num_shards})")
        num_shards = 48

    # Resolve build directory
    build_dir = Path(args.build_dir).resolve()
    if not build_dir.exists():
        print(f"ERROR: Build directory does not exist: {build_dir}")
        sys.exit(1)

    # Get test commands from profile using rrtest library
    try:
        commands = get_test_commands(args.profile)
        print(f"Using test profile: {args.profile}")
        print(f"Frameworks: {', '.join(commands.keys())}")
    except Exception as e:
        print(f"ERROR: Failed to load test profile '{args.profile}': {e}")
        sys.exit(1)

    print("=" * 50)
    print("Running sharded tests" + (" [DRY RUN]" if args.dry_run else ""))
    print(f"Build directory: {build_dir}")
    print(f"Test profile: {args.profile}")
    print(f"Available cores: {num_cores}")
    if gpu_arch:
        print(f"GPU architecture: {gpu_arch}")
    if not args.dry_run:
        print(f"Number of shards: {num_shards} (auto-detected)")
        shuffle_seed, from_env = get_shuffle_seed()
        if from_env:
            print(
                f"Shuffle seed: {shuffle_seed} (from TEST_SHUFFLE_SEED, reproducible)"
            )
        else:
            print(f"Shuffle seed: {shuffle_seed} (set TEST_SHUFFLE_SEED to reproduce)")
    if args.dry_run:
        print("Mode: DRY RUN - listing tests only")
    print("=" * 50)

    # Create test report directory
    test_report_dir = build_dir / "test_report"
    test_report_dir.mkdir(exist_ok=True)

    # Handle dry-run mode: list tests instead of executing
    if args.dry_run:
        print("\n=== DRY RUN MODE: Listing tests that would be executed ===\n")

        for framework, cmd_list in commands.items():
            for cmd in cmd_list:
                print(f"--- {framework.upper()} Test Suite ---")
                list_cmd = cmd.copy()

                # Add framework-specific list arguments
                if framework == "gtest":
                    list_cmd.append("--gtest_list_tests")
                elif framework == "catch2":
                    list_cmd.append("--list-tests")
                elif framework == "ctest":
                    list_cmd.extend(["--show-only=json-v1"])

                result = subprocess.run(list_cmd, cwd=build_dir)
                if result.returncode != 0:
                    print(
                        f"ERROR: Failed to list {framework} tests (exit code {result.returncode})"
                    )
                    sys.exit(1)
                print()

        print("=== DRY RUN COMPLETE ===")
        sys.exit(0)

    # Run all shards for each framework
    failed_shards = []
    all_timing_info = []

    # Process each framework in order
    for framework, cmd_list in commands.items():
        for cmd in cmd_list:
            # Determine number of shards based on framework
            # CTest runs tests in parallel internally, so only run 1 "shard"
            framework_num_shards = 1 if framework == "ctest" else num_shards

            print(f"\n--- Running {framework.upper()} shards ---")
            shard_infos = []

            for i in range(framework_num_shards):
                try:
                    seed = shuffle_seed if framework in ("gtest", "catch2") else None
                    info = start_shard(
                        i,
                        cmd,
                        framework,
                        framework_num_shards,
                        build_dir,
                        shard_seed=seed,
                    )
                    shard_infos.append(info)
                except Exception as e:
                    print(
                        f"ERROR: Failed to start {framework} shard {i}: {e}", flush=True
                    )
                    failed_shards.append((i, framework, 1))

            # Run all shards for this framework and wait for completion
            framework_failures, framework_timing = run_shards(shard_infos)
            failed_shards.extend(framework_failures)
            all_timing_info.extend(framework_timing)

    print()

    # Print timing summary
    if all_timing_info:
        print("=" * 50)
        print("TIMING SUMMARY")
        print("=" * 50)

        # Group timings by executable
        executable_timings = {}
        for idx, test_type, executable, elapsed in all_timing_info:
            if executable not in executable_timings:
                executable_timings[executable] = {"framework": test_type, "timings": []}
            executable_timings[executable]["timings"].append((idx, elapsed))

        # Print timings for each executable
        for executable in sorted(executable_timings.keys()):
            data = executable_timings[executable]
            timings = data["timings"]
            framework = data["framework"]
            prefix = framework[0].upper()

            print(f"\n{executable} ({framework}):")
            for idx, elapsed in sorted(timings):
                print(f"  [{prefix}{idx:02d}] {elapsed:6.1f}s")

            total = sum(t for _, t in timings)
            avg = total / len(timings)
            min_time = min(t for _, t in timings)
            max_time = max(t for _, t in timings)
            print(
                f"  Total: {total:.1f}s, Avg: {avg:.1f}s, Min: {min_time:.1f}s, Max: {max_time:.1f}s"
            )

        print("=" * 50)
        print()

    # Report results
    if failed_shards:
        print("FAILED: One or more test shards failed:")
        for shard_index, test_type, exit_code in failed_shards:
            print(f"  - {test_type} shard {shard_index}: exit code {exit_code}")
        sys.exit(1)
    else:
        print("SUCCESS: All test shards completed successfully.")
        print()
        print(f"Test results written to: {test_report_dir}/")

        # List generated XML files
        xml_files = sorted(test_report_dir.glob("*.xml"))
        if xml_files:
            for xml_file in xml_files:
                size = xml_file.stat().st_size
                print(f"  - {xml_file.name} ({size:,} bytes)")

        sys.exit(0)


if __name__ == "__main__":
    main()