############################################################################
# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
############################################################################
"""Test the --runtime job submission option and its runtime plugins."""

import json

import pytest

import atf


# Setup
@pytest.fixture(scope="module", autouse=True)
def setup():
    atf.require_version(
        (26, 11),
        "bin/srun",
        reason="Issue 50190: --runtime plugin ownership added in 26.11",
    )
    atf.require_version(
        (26, 11),
        "bin/salloc",
        reason="Issue 50190: --runtime plugin ownership added in 26.11",
    )
    atf.require_version(
        (26, 11),
        "bin/sbatch",
        reason="Issue 50190: --runtime plugin ownership added in 26.11",
    )
    atf.require_version(
        (26, 11),
        "bin/scontrol",
        reason="Issue 50190: --runtime plugin ownership added in 26.11",
    )
    atf.require_version(
        (26, 11),
        "bin/squeue",
        reason="Issue 50190: squeue -O Runtime added in 26.11",
    )
    atf.require_auto_config("deletes oci.conf directly on teardown")
    # salloc rejects --container without an interactive step to pass it to.
    atf.require_config_parameter_includes("LaunchParameters", "use_interactive_step")
    # A het job's components are co-scheduled, so a running one needs a node
    # per component.
    atf.require_nodes(2)
    atf.require_slurm_running()


@pytest.mark.parametrize("command", ["salloc", "sbatch", "srun"])
def test_runtime_recorded(command):
    """salloc/sbatch/srun record --runtime on the job, normalized to
    runtime/<name>."""

    job_id = atf.submit_job(
        command, "--runtime=oci", "sleep 60", background=True, fatal=True
    )
    atf.wait_for_job_state(job_id, "RUNNING", fatal=True)
    assert (
        atf.get_job_parameter(job_id, "Runtime", fatal=True) == "runtime/oci"
    ), f"{command}: --runtime not recorded as runtime/oci"
    runtime = atf.run_command_output(
        f"squeue -j {job_id} -O Runtime --noheader", fatal=True
    ).strip()
    assert runtime == "runtime/oci", f"{command}: squeue Runtime should be runtime/oci"


def test_runtime_json_key_rename():
    """The job runtime is keyed 'runtime' in v0.0.46 and 'container_type' in
    v0.0.45 (the same struct member was renamed)."""

    job_id = atf.submit_job_sbatch('--runtime=oci --hold --wrap "true"', fatal=True)

    job = json.loads(
        atf.run_command_output(f"scontrol --json show job {job_id}", fatal=True)
    )["jobs"][0]
    assert job.get("runtime") == "runtime/oci", "v0.0.46 'runtime' key missing"
    assert "container_type" not in job, "v0.0.46 should not use 'container_type'"

    job45 = json.loads(
        atf.run_command_output(f"scontrol --json=v0.0.45 show job {job_id}", fatal=True)
    )["jobs"][0]
    assert (
        job45.get("container_type") == "runtime/oci"
    ), "v0.0.45 'container_type' key missing"


def test_runtime_bogus_fails():
    """An unknown runtime plugin fails when the step is launched.

    Only srun is exercised: it launches a step under slurmstepd, which loads
    the runtime plugin and fails. salloc and sbatch do not validate the
    runtime plugin name at submission time, so they are not tested here.
    """

    error = atf.run_command_error("srun --runtime=bogus true", xfail=True, fatal=True)
    assert (
        "launch failed" in error
    ), "srun should report a step launch failure for an unknown runtime plugin"


def test_runtime_none_ignores_container():
    """runtime/none runs tasks on the host, ignoring the --container value."""

    # runtime/none ignores --container (even an invalid bundle) and runs the
    # task directly on the host, exiting 0.
    atf.run_command("srun --runtime=none --container=invalid hostname", fatal=True)


@pytest.fixture
def oci_enabled():
    """Enable OCI via a minimal runc oci.conf and remove it on teardown.

    OCI support is global (oci.conf enables --container processing for every
    step), so it is scoped to this fixture to keep it away from the other tests,
    which assume OCI is disabled.
    """

    atf.require_config_file(
        "oci.conf",
        'EnvExclude="^(SLURM_CONF|SLURM_CONF_SERVER)="\n'
        'RunTimeQuery="runc --rootless=true --root=/run/user/%U/ state %n.%u.%j.%s.%t"\n'
        'RunTimeKill="runc --rootless=true --root=/run/user/%U/ kill -a %n.%u.%j.%s.%t"\n'
        'RunTimeDelete="runc --rootless=true --root=/run/user/%U/ delete --force %n.%u.%j.%s.%t"\n'
        'RunTimeRun="runc --rootless=true --root=/run/user/%U/ run %n.%u.%j.%s.%t -b %b"\n',
    )
    yield
    atf.run_command(
        f"rm -f {atf.properties['slurm-config-dir']}/oci.conf",
        user=atf.properties["slurm-user"],
        fatal=True,
        quiet=True,
    )


def test_runtime_oci_uses_container(oci_enabled):
    """runtime/oci engages the --container value, unlike runtime/none.

    runtime/oci does not ignore --container, so an invalid bundle fails the
    step launch where runtime/none exits 0 on the same input. A full
    containerized run additionally needs a valid OCI bundle and is out of scope
    here.
    """

    atf.run_command_error(
        "srun --runtime=oci --container=invalid hostname", xfail=True, fatal=True
    )


@pytest.mark.parametrize("command", ["salloc", "sbatch", "srun"])
@pytest.mark.parametrize(
    "options,parameter,expected,inherited",
    [
        ("--runtime=oci", "Runtime", "runtime/oci", None),
        (
            "--container=/nonexistent/bundle",
            "Container",
            "/nonexistent/bundle",
            None,
        ),
        ("--container-id=abc123", "ContainerID", "abc123", "abc123"),
    ],
    ids=["runtime", "container", "container-id"],
)
def test_hetjob_component_option_propagation(
    command, options, parameter, expected, inherited
):
    """A heterogeneous component that omits an option records what that
    option's propagation rule says, not whatever the previous component
    asked for.

    --runtime and --container are reset for each component, so the second
    records neither. --container-id is propagated, so the second inherits
    it. heterogeneous_jobs.html lists which options fall on which side.

    Only an omitted option can inherit. A component that names the option
    explicitly always wins, so giving every component its own value would
    prove nothing either way.

    The job is held and never launched, so this asserts on the submitted job
    records. A component that asked for neither records neither - the runtime
    falls back to DefRuntimePlugin later, when slurmstepd loads the plugin.
    """

    if command == "sbatch":
        job_id = atf.submit_job_sbatch(
            f"-N1 {options} --hold --wrap=true : -N1", fatal=True
        )
    else:
        # A held salloc or srun never returns, so it has to be backgrounded.
        job_id = atf.submit_job(
            command,
            f"-N1 {options} --hold : -N1",
            "true",
            background=True,
            fatal=True,
        )
    jobs = list(atf.get_jobs(job_id, fatal=True).values())

    assert len(jobs) == 2, "the het job should have two components"
    components = {job["HetJobOffset"]: job for job in jobs}
    assert (
        components[0][parameter] == expected
    ), f"component 0 should record the {parameter} it asked for"
    assert (
        components[1].get(parameter) == inherited
    ), f"component 1 asked for no {parameter} and should record {inherited}"


def test_hetjob_component_runtime_at_step_launch(oci_enabled):
    """Each heterogeneous component launches its step under its own runtime.

    With DefRuntimePlugin=runtime/none, the component naming --runtime=oci
    engages its --container value and fails on an invalid bundle, while the
    component naming no runtime ignores --container and succeeds. Exactly one
    of the two steps failing is what separates per-component plugin selection
    from either component's runtime being used for both.

    The job runs, unlike the submission tests above, because the plugin is
    chosen in slurmstepd rather than at submission.
    """

    atf.set_config_parameter("DefRuntimePlugin", "runtime/none")

    error = atf.run_command_error(
        "srun -N1 --runtime=oci --container=invalid : -N1 --container=invalid true",
        fatal=True,
    )
    assert (
        error.count("Task launch for StepId=") == 1
    ), "only the component that asked for runtime/oci should fail its step launch"


def test_runtime_env_vars():
    """SLURM_RUNTIME is exported to the job and SBATCH_RUNTIME sets --runtime."""

    output = atf.run_command_output("srun --runtime=none env", fatal=True)
    assert (
        "SLURM_RUNTIME=runtime/none" in output
    ), "SLURM_RUNTIME should be exported to the job environment"

    job_id = atf.submit_job_sbatch(
        '--hold --wrap "true"', env_vars="SBATCH_RUNTIME=oci", fatal=True
    )
    assert (
        atf.get_job_parameter(job_id, "Runtime", fatal=True) == "runtime/oci"
    ), "SBATCH_RUNTIME should set the job runtime"


def test_runtime_env_var_input_and_negative():
    """SRUN_RUNTIME sets --runtime, and SLURM_RUNTIME is only exported when
    --runtime is specified."""

    output = atf.run_command_output(
        "srun env", env_vars="SRUN_RUNTIME=none", fatal=True
    )
    assert (
        "SLURM_RUNTIME=runtime/none" in output
    ), "SRUN_RUNTIME should set the job runtime"

    output = atf.run_command_output("srun env", fatal=True)
    assert (
        "SLURM_RUNTIME=" not in output
    ), "SLURM_RUNTIME should not be exported when --runtime is not specified"


def test_def_runtime_plugin_reported():
    """scontrol show config reports DefRuntimePlugin, defaulting to runtime/oci."""

    atf.set_config_parameter("DefRuntimePlugin", None)
    assert (
        atf.get_config_parameter("DefRuntimePlugin", live=True) == "runtime/oci"
    ), "an unset DefRuntimePlugin should report the documented default runtime/oci"

    atf.set_config_parameter("DefRuntimePlugin", "runtime/none")
    assert (
        atf.get_config_parameter("DefRuntimePlugin", live=True) == "runtime/none"
    ), "DefRuntimePlugin should reflect the configured runtime/none"


def test_def_runtime_plugin_takes_effect(oci_enabled):
    """A job with no --runtime uses the configured DefRuntimePlugin.

    With oci.conf enabled, DefRuntimePlugin=runtime/oci engages the
    --container value and an invalid bundle fails the step launch, while
    DefRuntimePlugin=runtime/none ignores --container and the same job
    succeeds. The jobs request no --runtime, so the contrast proves the
    configured default selects the plugin.
    """

    atf.set_config_parameter("DefRuntimePlugin", "runtime/oci")
    atf.run_command_error("srun --container=invalid hostname", xfail=True, fatal=True)

    atf.set_config_parameter("DefRuntimePlugin", "runtime/none")
    atf.run_command("srun --container=invalid hostname", fatal=True)