############################################################################
# Copyright (C) SchedMD LLC.
############################################################################
import logging
import re

import pytest

import atf

pytestmark = pytest.mark.slow

sock_0_reserved_cores = {0, 1, 2, 3}
sock_0_regular_cores = {4, 5}
sock_1_reserved_cores = {6, 7, 8, 9}
sock_1_regular_cores = {10, 11}


# Setup
@pytest.fixture(scope="module", autouse=True)
def setup():
    atf.require_auto_config("wants to create custom gpu files and custom gres")
    atf.require_config_parameter("SelectType", "select/cons_tres")
    atf.require_config_parameter(
        "SelectTypeParameters", "CR_CORE,MULTIPLE_SHARING_GRES_PJ"
    )
    atf.require_config_parameter("GresTypes", "gpu,shard")
    atf.require_nodes(
        2,
        [
            ("Gres", "gpu:4,shard:40"),
            ("Sockets", 2),
            ("CoresPerSocket", 6),
            ("ThreadsPerCore", 1),
            ("RestrictedCoresPerGPU", 2),
        ],
    )

    for tty_num in range(8):
        atf.require_tty(tty_num)

    atf.require_config_parameter(
        "Name",
        "gpu Cores=0-5 File=/dev/tty[0-1]\nName=gpu Cores=6-11 File=/dev/tty[2-3]",
        source="gres",
    )
    atf.require_slurm_running()


def test_config():
    expected = "2(0-3,6-9)"
    res_cores = atf.run_command_output(
        "scontrol show nodes node1 | awk -F 'RestrictedCoresPerGPU=' '{print $2}' | tr -d ' \n'",
        fatal=True,
    )
    assert (
        expected == res_cores
    ), f"RestrictedCoresPerGPU config not what expected. {res_cores}(scontrol output) != {expected}(expected)\n"


ok_test_parameters = [
    (
        "-N1 -n1 ",
        0,
        1,
        0,
        0,
        "",
    ),
    (
        "-N1 -n4 ",
        0,
        2,
        0,
        2,
        "",
    ),
    (
        "-N2 -n8 ",
        0,
        2,
        0,
        2,
        "",  # Second node
    ),
    (
        "-N2 -n3 -c2 ",
        0,
        2,
        0,
        0,
        "",  # Second node
    ),
    (
        "-N1 -n1 --gres=gpu:1 ",
        1,
        0,
        0,
        0,
        "gpu:1(IDX:0)",
    ),
    (
        "-N1 -n4 --gres=gpu:1 ",
        2,
        2,
        0,
        0,
        "gpu:1(IDX:0)",
    ),
    (  # Try to allocate all Cores
        "-N1 -c12 --gres=gpu:4 ",
        4,
        2,
        4,
        2,
        "gpu:4(IDX:0-3)",
    ),
    (
        "-N1 -n4 --tres-per-task=gres/gpu:1 ",
        2,
        0,
        2,
        0,
        "gpu:4(IDX:0-3)",
    ),
    (
        "-N2 -n5 --tres-per-task=gres/gpu:1 ",
        1,
        0,
        0,
        0,
        "gpu:1(IDX:0)",  # Second node
    ),
    (
        "-N1 -n1 -c3 --tres-per-task=gres/gpu:1 ",
        2,
        1,
        0,
        0,
        "gpu:1(IDX:0)",
    ),
    (
        "-N1 -n3 -c2 --tres-per-task=gres/gpu:1 ",
        2,
        1,
        3,  # May warrant a fix. Ideally should be 2,0,4,0
        0,
        "gpu:3(IDX:0,2-3)",
    ),
    # Note: There is no guarantee that reserved cores are used in proportion to shared gres allocations
    pytest.param(  # Test allocating all shards with 1 cores (see Note above)
        "-N1 --gres=shard:40 ",
        1,
        0,
        0,
        0,
        "shard:40(10/10,10/10,10/10,10/10)",
        marks=pytest.mark.xfail(
            atf.get_version() < (25, 11),
            reason="Ticket 22391: ResCoresPerGPU with shard support",
        ),
    ),
    pytest.param(  # Test allocating all cores with 4 shards
        "-N1 -n1 -c12 --gres=shard:4 ",
        4,
        2,
        4,
        2,
        "shard:4(4/10,0/10,0/10,0/10)",
        marks=pytest.mark.xfail(
            atf.get_version() < (25, 11),
            reason="Ticket 22391: ResCoresPerGPU with shard support",
        ),
    ),
    pytest.param(
        "-N1 -n1 -c1 --tres-per-task=gres/shard:1 ",
        1,
        0,
        0,
        0,
        "shard:1(1/10,0/10,0/10,0/10)",
        marks=pytest.mark.xfail(
            atf.get_version() < (25, 11),
            reason="Ticket 22391: ResCoresPerGPU with shard support",
        ),
    ),
    pytest.param(
        "-N1 -n3 -c2 --tres-per-task=gres/shard:5 ",
        4,
        2,
        0,
        0,
        "shard:15(10/10,5/10,0/10,0/10)",
        marks=pytest.mark.xfail(
            atf.get_version() < (25, 11),
            reason="Ticket 22391: ResCoresPerGPU with shard support",
        ),
    ),
    pytest.param(
        "-N1 -n1 -c8 --tres-per-task=gres/shard:10 ",
        4,
        2,
        2,
        0,
        "shard:10(10/10,0/10,0/10,0/10)",
        marks=pytest.mark.xfail(
            atf.get_version() < (25, 11),
            reason="Ticket 22391: ResCoresPerGPU with shard support",
        ),
    ),
    pytest.param(
        "-N1 -n1 -c1 --gres=shard:1 ",
        1,
        0,
        0,
        0,
        "shard:1(1/10,0/10,0/10,0/10)",
        marks=pytest.mark.xfail(
            atf.get_version() < (25, 11),
            reason="Ticket 22391: ResCoresPerGPU with shard support",
        ),
    ),
    pytest.param(
        "-N1 -n2 -c3 --gres=shard:3 ",
        4,
        2,
        0,
        0,
        "shard:3(3/10,0/10,0/10,0/10)",
        marks=pytest.mark.xfail(
            atf.get_version() < (25, 11),
            reason="Ticket 22391: ResCoresPerGPU with shard support",
        ),
    ),
    pytest.param(
        "-N2 -n2 --gres=shard:5 ",
        1,
        0,
        0,
        0,
        "shard:5(5/10,0/10,0/10,0/10)",  # Second node
        marks=pytest.mark.xfail(
            atf.get_version() < (25, 11),
            reason="Ticket 22391: ResCoresPerGPU with shard support",
        ),
    ),
]


@pytest.mark.parametrize(
    "job_args,s0_res,s0_reg,s1_res,s1_reg,gres",
    ok_test_parameters,
    ids=[
        param.values[0].strip() if hasattr(param, "values") else param[0].strip()
        for param in ok_test_parameters
    ],
)
def test_ok(job_args, s0_res, s0_reg, s1_res, s1_reg, gres):
    job_str = f'{job_args} --wrap "sleep infinity"'
    job_id = atf.submit_job_sbatch(job_str, fatal=True, quiet=False)
    atf.wait_for_job_state(job_id, "RUNNING", fatal=True)
    job_dict = atf.get_job(job_id)
    atf.cancel_all_jobs(quiet=True)
    atf.wait_for_job_state(job_id, "CANCELLED")

    assert str(job_dict["GRES"]) == gres, f"The job should have this gres: {gres}"

    cores = set(atf.range_to_list(job_dict["CPU_IDs"]))

    logging.debug(f"Values: {s0_res}, {s0_reg}, {s1_res}, {s1_reg}")

    assert s0_res == len(
        cores & sock_0_reserved_cores
    ), f"The job should have {s0_res} reserved cores on socket 0. We got CPU_IDs={job_dict['CPU_IDs']}"
    assert s0_reg == len(
        cores & sock_0_regular_cores
    ), f"The job should have {s0_reg} not reserved cores on socket 0. We got CPU_IDs={job_dict['CPU_IDs']}"
    assert s1_res == len(
        cores & sock_1_reserved_cores
    ), f"The job should have {s1_res} reserved cores on socket 1. We got CPU_IDs={job_dict['CPU_IDs']}"
    assert s1_reg == len(
        cores & sock_1_regular_cores
    ), f"The job should have {s1_reg} not reserved cores on socket 1. We got CPU_IDs={job_dict['CPU_IDs']}"


array_test_parameters = [
    (  # Test with no gres, should use only regular cores
        "--array=[0-100] -N1 -n1 --nodelist=node1 ",
        4,
    ),
    (  # Test with gres gpu, should use all gpus
        "--array=[0-100] -N1 -n1 --nodelist=node1 --gres=gpu:1 ",
        4,
    ),
    pytest.param(  # Test with gres shard, should use all cores
        "--array=[0-100] -N1 -n1 --nodelist=node1 --gres=shard:1 ",
        12,
        marks=pytest.mark.xfail(
            atf.get_version() < (25, 11),
            reason="Ticket 22391: ResCoresPerGPU with shard support",
        ),
    ),
    pytest.param(  # Test with gres shard, should use all shards
        "--array=[0-100] -N1 -n1 --nodelist=node1 --gres=shard:5 ",
        8,
        marks=pytest.mark.xfail(
            atf.get_version() < (25, 11),
            reason="Ticket 22391: ResCoresPerGPU with shard support",
        ),
    ),
]


@pytest.mark.parametrize(
    "job_args,running_count",
    array_test_parameters,
    ids=[
        param.values[0].strip() if hasattr(param, "values") else param[0].strip()
        for param in array_test_parameters
    ],
)
def test_array(job_args, running_count):
    job_str = f'{job_args} --wrap "sleep infinity"'
    job_id = atf.submit_job_sbatch(job_str, fatal=True, quiet=False)
    atf.wait_for_node_state("node1", ["ALLOCATED", "MIXED"], fatal=True)
    count = int(atf.run_command_output("squeue --noheader --state=Running| wc -l"))
    atf.cancel_all_jobs(quiet=True)
    atf.wait_for_job_state(job_id, "CANCELLED")

    logging.info(f"Count: {count} / {running_count}")

    assert (
        count == running_count
    ), f"The job array should have {running_count} running tasks"


fail_test_parameters = [
    (
        "-N1 -n3 -c2",
        r"srun: error: .+ Requested node configuration is not available",
    ),
    (  # Try allocating more cores than GPUs * ReservedCoresPerGPU + regular cores
        "-N1 -n4 -c2 --gres=gpu:1",
        r"srun: error: .+ Requested node configuration is not available",
    ),
    (
        "-N1 -n5 -c1",
        r"srun: error: .+ Requested node configuration is not available",
    ),
    (
        "-N2 -n5 -c2",
        r"srun: error: .+ Requested node configuration is not available",
    ),
    (
        "-N2 -n9 -c1",
        r"srun: error: .+ Requested node configuration is not available",
    ),
    # You need at least 1 shard per GPU if you want all Cores despite the fact
    # that they may be on the same GPU. Fully requiring GPU usage to get the
    # cores would be difficult and beyond current scope.
    # This is a quirk of the current implementation of shard allocation with
    # ReservedCoresPerGPU
    (
        "-N1 -n1 -c12 --gres=shard:1",
        r"srun: error: .+ Requested node configuration is not available",
    ),
]


@pytest.mark.parametrize(
    "job_args,expected_msg",
    fail_test_parameters,
    ids=[p[0].strip() for p in fail_test_parameters],
)
def test_fail(job_args, expected_msg):
    output = atf.run_command(
        f"srun {job_args} true",
        xfail=True,
    )
    assert output["exit_code"] != 0, (
        f"Expected non-zero exit code for job_args: '{job_args}', because of ReservedCoresPerGPU.\n"
        f"STDOUT: {output['stdout']!r}\n"
        f"STDERR: {output['stderr']!r}"
    )
    assert re.search(expected_msg, output["stderr"]), (
        f"Expected error should be in stderr, got:\n"
        f"job_args: {job_args}\n"
        f"exit_code: {output['exit_code']}\n"
        f"Expected regex: {expected_msg}\n"
        f"STDOUT: {output['stdout']!r}\n"
        f"STDERR: {output['stderr']!r}"
    )