############################################################################
# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
############################################################################
import pytest

import atf

line_received = "Line received:"


def setup_module():
    atf.require_version(
        (26, 5),
        "bin/sattach",
        reason="Ticket 22180: SLUID availability added in 26.05+",
    )
    atf.require_nodes(1)
    atf.require_slurm_running()


@pytest.fixture(scope="function")
def interactive_step_id(request):
    """
    Submits a job with a single step waiting for 'trigger' input, and returns
    the StepId based on the request.param, so it can be JobId.0 or SLUID.0
    """

    file_out = "file.out"
    ready_str = "ready to read input"
    step_script = "interactive_step.sh"
    atf.make_bash_script(
        step_script,
        f"""
echo {ready_str}
read line
echo "{line_received} $line"
""",
    )

    job_script = "job_with_interactive_step.sh"
    atf.make_bash_script(
        job_script,
        # We need to keep the stdin of the step open or an EOF will close it
        # and sattach won't be able to use it.
        f"""srun {step_script} < <(sleep infinity)""",
    )

    job_id = atf.submit_job_sbatch(f"-t1 -o {file_out} {job_script}", fatal=True)

    # Wait until the interactive step is ready to read the line
    atf.wait_for_file(file_out)
    for t in atf.timer():
        output = atf.run_command_output(f"cat {file_out}")
        if ready_str in output:
            break
    else:
        pytest.fail(f"File {file_out} wasn't ready")

    # Return the StepID using either the JobId or SLUID based on the param
    return f"{atf.get_job_parameter(job_id, request.param)}.0"


@pytest.mark.parametrize("interactive_step_id", ["JobId", "SLUID"], indirect=True)
def test_sattach(interactive_step_id):
    """Verify sattach receives output and can send input to the step, using JobId or SLUID."""

    result = atf.run_command(f"sattach {interactive_step_id}", input="trigger")
    assert (
        f"{line_received} trigger" in result["stdout"]
    ), f"sattach should attach IO streams and get '{line_received} trigger', got: {result['stdout']}"
    assert result["exit_code"] == 0, "sattach should end normally"