############################################################################
# Copyright (C) SchedMD LLC.
############################################################################
import atf
import pytest
from concurrent.futures import ThreadPoolExecutor, as_completed
import time

max_threads = 10


@pytest.fixture(scope="module", autouse=True)
def setup():
    atf.require_nodes(max_threads)
    atf.require_slurm_running()


# Defines the work done by each executor thread in the stress test
def run_stress_task(thread_id, iterations, sleep_time):
    """
    Simulates the work of a single stress thread.
    Each thread runs a series of Slurm commands (sinfo, sbatch, squeue) in a loop.
    This function is designed to be run by an executor thread (ThreadPoolExecutor).

    Args:
        thread_id (int): A unique identifier for this run/thread.
        iterations (int): The number of times to loop through the command sequence.
        sleep_time (float): The duration in seconds to sleep between command executions.

    Returns:
        tuple: A tuple containing the thread_id and an exit code (0 for success, non-zero for failure).
    """
    # Loop for the specified number of iterations
    for i in range(1, iterations + 1):
        # Run sinfo command
        result_sinfo = atf.run_command("sinfo")
        if result_sinfo["exit_code"] != 0:
            return thread_id, result_sinfo["exit_code"]

        # Pause execution for the specified sleep time
        time.sleep(sleep_time)

        # Construct the sbatch command to execute /bin/true.
        # Job name includes thread_id and iteration number for uniqueness.
        # Node and task counts for sbatch vary with the iteration number.
        sbatch_args = f"--job-name=test_task{thread_id}_{i} -N1-{i} -n{i} -O -s -t1 --wrap='/bin/true'"
        job_id = atf.submit_job_sbatch(sbatch_args)
        if job_id == 0:
            return thread_id, 1

        # Pause execution
        time.sleep(sleep_time)

        # Run squeue command
        result_squeue = atf.run_command("squeue")
        if result_squeue["exit_code"] != 0:
            return thread_id, result_squeue["exit_code"]

        # Pause execution
        time.sleep(sleep_time)

    # If all iterations complete successfully, return 0 for success
    return thread_id, 0


@pytest.mark.parametrize(
    "threads_count, iterations, sleep_time",
    [
        (2, 3, 1),  # A smaller run for quick testing
        (
            max_threads,
            int(max_threads / 2),
            1,
        ),  # Parameters mimicking the original expect test's default values
        (5, 2, 0.5),  # A quicker, moderate run
    ],
)
def test_stress_slurm_commands(threads_count, iterations, sleep_time):
    """
    Stress test multiple simultaneous Slurm commands via multiple threads.
    Each thread executes the `run_stress_task` function, which runs sinfo,
    sbatch /bin/true, and squeue in a loop.
    The test verifies that all concurrently run tasks complete successfully.
    """
    # Calculate a timeout for the entire test.
    # This is based on the original expect script's timeout logic: max_job_delay * iterations * thread_cnt.
    # Assuming max_job_delay is 120 seconds as often seen in Slurm expect tests.
    total_timeout = 120 * iterations * threads_count

    # Initialize counters and lists for tracking task results
    successful_tasks = 0
    failed_thread_details = []

    # Use ThreadPoolExecutor to run tasks concurrently
    with ThreadPoolExecutor(max_workers=threads_count) as executor:
        # Create a dictionary to map future objects to their thread_ids
        # This helps in identifying which thread corresponds to a completed future
        futures = {
            executor.submit(
                run_stress_task,
                thread_id,
                iterations,
                sleep_time,
            ): thread_id
            for thread_id in range(
                threads_count
            )  # Create `threads_count` number of threads
        }

        try:
            # Wait for threads to complete, with an overall timeout for the block
            # as_completed yields futures as they complete (or raise exceptions)
            for future in as_completed(futures, timeout=total_timeout):
                # Get the original thread_id for this future
                thread_id_from_future_map = futures[future]
                try:
                    # Get the result from the completed future. This will
                    # re-raise any exception that occurred in the worker thread.
                    returned_thread_id, exit_code = future.result()
                    if exit_code == 0:
                        successful_tasks += 1
                    else:
                        # If the thread reported a non-zero exit code, record it
                        # as a failure
                        error_msg = f"Thread {returned_thread_id} failed with exit code {exit_code}."
                        failed_thread_details.append(error_msg)
                except Exception as e:
                    # If future.result() raises an exception (e.g., an unhandled
                    # error in run_stress_task), record this as a failure.
                    error_msg = f"Thread {thread_id_from_future_map} generated an exception: {type(e).__name__} - {e}"
                    failed_thread_details.append(error_msg)
        except TimeoutError:
            # If as_completed times out, it means not all tasks finished within total_timeout
            error_msg = f"Stress test timed out after {total_timeout:.2f} seconds. Not all tasks completed."
            failed_thread_details.append(error_msg)
            # Attempt to cancel any tasks that are still running
            for fut in futures:
                if not fut.done():
                    fut.cancel()
        except Exception as e:
            # Catch any other unexpected errors during as_completed processing
            error_msg = f"An unexpected error occurred during task execution: {type(e).__name__} - {e}"
            failed_thread_details.append(error_msg)

    # Assert that all tasks completed successfully
    # If not, provide a detailed error message listing the failures.
    assert (
        successful_tasks == threads_count
    ), f"Only {successful_tasks} of {threads_count} tasks passed. Failures: {'; '.join(failed_thread_details)}"