############################################################################
# Copyright (C) SchedMD LLC.
############################################################################
import atf
import pytest
import os

# import re
# import time
import pwd

max_mem_cpu = 2
max_mem_node = 1
min_num_nodes = 2
max_num_nodes = 1
max_time = 1
node_list = []
current_policy = ""
previous_limit = ""
p1_node_str = ""

# Limits in development that are not included in this test:
# AllocNodes, AllowGroups, QOS usage threshold

limits_dict = {
    "MaxMemPerCPU": {
        "flag": "--mem-per-cpu=",
        "fail": max_mem_cpu + 1,
        "pass": max_mem_cpu,
        "_set": max_mem_cpu,
    },
    "MaxMemPerNode": {
        "flag": "--mem=",
        "fail": max_mem_node + 1,
        "pass": max_mem_node,
        "_set": max_mem_node,
    },
    "MinNodes": {
        "flag": "-N",
        "fail": min_num_nodes - 1,
        "pass": min_num_nodes,
        "_set": min_num_nodes,
    },
    "MaxNodes": {
        "flag": "-N",
        "fail": max_num_nodes + 1,
        "pass": max_num_nodes,
        "_set": max_num_nodes,
    },
    "MaxTime": {"flag": "-t", "fail": max_time + 1, "pass": max_time, "_set": max_time},
    "AllowAccounts": {
        "flag": "",
        "fail": "--account=bad_account",
        "pass": "--account=good_account",
        "_set": "good_account",
    },
    "AllowQos": {
        "flag": "",
        "fail": "--qos=bad_qos",
        "pass": "--qos=good_qos",
        "_set": "good_qos",
    },
}


# Setup
@pytest.fixture(scope="module", autouse=True)
def setup():
    global node_list, p1_node_str

    atf.require_auto_config("wants to change partitions, modify/create nodes")
    atf.require_accounting(modify=True)
    atf.require_config_parameter("AccountingStorageEnforce", "limits")

    # Gather a list of nodes that meet the RealMemory Requirement (node_list)
    atf.require_nodes(4, [("RealMemory", max_mem_node + 1)])

    p1_node_str = "node1,node2"
    atf.require_config_parameter(
        "PartitionName",
        {
            "p1": {"Nodes": p1_node_str, "Default": "NO", "State": "UP"},
            "p2": {
                "Nodes": "node3,node4",
                "State": "UP",
                "MaxTime": "INFINITE",
            },
        },
    )

    atf.require_slurm_running()


@pytest.fixture(scope="module")
def setup_account():
    test_user = pwd.getpwuid(os.getuid())[0]
    atf.run_command(
        "sacctmgr -vi add account good_account",
        user=atf.properties["slurm-user"],
        fatal=True,
    )
    atf.run_command(
        f"sacctmgr -vi add user {test_user} account=good_account",
        user=atf.properties["slurm-user"],
        fatal=True,
    )
    atf.run_command(
        "sacctmgr -vi add qos good_qos,bad_qos",
        user=atf.properties["slurm-user"],
        fatal=True,
    )
    atf.run_command(
        f"sacctmgr -vi modify user {test_user} set qos=good_qos,bad_qos",
        user=atf.properties["slurm-user"],
        fatal=True,
    )


@pytest.fixture(scope="function")
def cancel_jobs():
    """Cancel all jobs after each test"""
    yield
    atf.cancel_all_jobs(fatal=True)


# Helper funcs:


def set_enforce_part_limits_policy(policy):
    global current_policy

    if current_policy != policy:
        atf.set_config_parameter("EnforcePartLimits", policy)
        current_policy = policy


def set_partition_limit(limit_name, limit_value):
    global previous_limit

    if previous_limit:
        atf.set_partition_parameter("p1", previous_limit, None)
    atf.set_partition_parameter("p1", limit_name, limit_value)
    previous_limit = limit_name


def satisfy_pending_job_limit(job_id, limit_name, val_pass):
    atf.wait_for_job_state(job_id, "PENDING", poll_interval=0.1, fatal=True, quiet=True)

    # Update partition limit to comply
    atf.run_command(
        f"scontrol update partitionname=p1 {limit_name}={val_pass}",
        user=atf.properties["slurm-user"],
        fatal=True,
        quiet=True,
    )

    # Allow time for job to requeue and complete
    atf.wait_for_job_state(job_id, "DONE", poll_interval=0.5, fatal=True, quiet=True)


# Test functions
def enforce_ALL(limit_name, flag, val_fail, val_pass):
    # Must account for the higher value needed to pass first assert when using
    # MaxMemPerCPU as it will allocate more cpus to fill the memory requirement
    # If it can't allocate more cpus it will fail, so the limit will still be imposed
    # on p2 (test 3) in a way because the MaxMemPerCPU will exceed the nodes RealMemory
    # see https://slurm.schedmd.com/slurm.conf.html#OPT_MaxMemPerCPU

    custom_val_fail = val_fail
    if limit_name == "MaxMemPerCPU":
        custom_val_fail = val_fail - 1
    else:
        if limit_name == "AllowAccounts":
            custom_val_fail = ""

    # 1 Reject p1,p2 with no p1 limit met
    assert (
        atf.submit_job_sbatch(
            f'-p p1,p2 {flag}{val_fail} --wrap "hostname" -o /dev/null'
        )
        == 0
    ), f"Job should fail on p1,p2 due to {limit_name} limit not met on the required partition p1 with EnforcePartLimits=ALL"

    # 2 Reject p1 no limit met
    assert (
        atf.submit_job_sbatch(f'-p p1 {flag}{val_fail} --wrap "hostname" -o /dev/null')
        == 0
    ), f"Job should fail on p1 due to {limit_name} limit not met on the required partition p1 with EnforcePartLimits=ALL"

    # 3 Accept p2 no limit met ** This one and the first have a memory conflict
    assert (
        atf.submit_job_sbatch(
            f'-p p2 {flag}{custom_val_fail} --wrap "hostname" -o /dev/null'
        )
        != 0
    ), f"Job should pass on p2 despite {limit_name} limit not met on the required partition p1 with EnforcePartLimits=ALL"

    # 4 Accept p1 with limit met
    assert (
        atf.submit_job_sbatch(f'-p p1 {flag}{val_pass} --wrap "hostname" -o /dev/null')
        != 0
    ), f"Job should pass on p1 due to {limit_name} limit met on the required partition p1 with EnforcePartLimits=ALL"

    # 5 Accept p1,p2 with p1 limit met
    assert (
        atf.submit_job_sbatch(
            f'-p p1,p2 {flag}{val_pass} --wrap "hostname" -o /dev/null'
        )
        != 0
    ), f"Job should pass on p1,p2 due to {limit_name} limit met on the required partition p1 with EnforcePartLimits=ALL"


def enforce_ANY(limit_name, flag, val_fail, val_pass):
    # Must account for the higher value needed to pass first assert (from 'ALL') when using
    # MaxMemPerCPU as it will allocate more cpus to fill the memory requirement

    custom_val_fail = val_fail
    if limit_name == "MaxMemPerCPU":
        custom_val_fail = val_fail - 1
    else:
        if limit_name == "AllowAccounts":
            custom_val_fail = ""

    # 1 Accept p1,p2 with no p1 limit met
    assert (
        atf.submit_job_sbatch(
            f'-p p1,p2 {flag}{custom_val_fail} --wrap "hostname" -o /dev/null'
        )
        != 0
    ), f"Job should pass on p1,p2 despite {limit_name} not met on p1 with EnforcePartLimits=ANY"

    # 2 Reject p1 no limit met
    assert (
        atf.submit_job_sbatch(f'-p p1 {flag}{val_fail} --wrap "hostname" -o /dev/null')
        == 0
    ), f"Job should fail on p1 due to {limit_name} limit not met on the required partition p1 with EnforcePartLimits=ANY"

    # 3 Accept p2 no limit met ** This one and the first have a memory conflict
    assert (
        atf.submit_job_sbatch(
            f'-p p2 {flag}{custom_val_fail} --wrap "hostname" -o /dev/null'
        )
        != 0
    ), f"Job should pass on p2 despite {limit_name} limit not met on the required partition p1 with EnforcePartLimits=ANY"

    # 4 Accept p1 with limit met
    assert (
        atf.submit_job_sbatch(f'-p p1 {flag}{val_pass} --wrap "hostname" -o /dev/null')
        != 0
    ), f"Job should pass on p1 due to {limit_name} limit met on the required partition p1 with EnforcePartLimits=ANY"

    # 5 Accept p1,p2 with p1 limit met
    assert (
        atf.submit_job_sbatch(
            f'-p p1,p2 {flag}{val_pass} --wrap "hostname" -o /dev/null'
        )
        != 0
    ), f"Job should pass on p1,p2 due to {limit_name} limit met on the required partition p1 with EnforcePartLimits=ANY"


def enforce_NO(limit_name, flag, val_fail, val_pass):
    # Must account for the higher value needed to pass first assert (from 'ALL') when using
    # MaxMemPerCPU as it will allocate more cpus to fill the memory requirement

    custom_val_fail = val_fail
    if limit_name == "MaxMemPerCPU":
        custom_val_fail = val_fail - 1
    else:
        if limit_name == "AllowAccounts":
            custom_val_fail = ""

    # 1 Submit -> pend on p1,p2 with bad p1 limit set -> complete with p1 limit met
    job_id = atf.submit_job_sbatch(
        f'-p p1,p2 {flag}{custom_val_fail} --wrap "hostname&" -o /dev/null'
    )
    satisfy_pending_job_limit(job_id, limit_name, custom_val_fail)
    assert (
        atf.get_job_parameter(job_id, "JobState", quiet=True) == "COMPLETED"
    ), f"Job should submit, pend, then complete on p1,p2 with updated limit {limit_name} on partition p1 to passing valueswith EnforcePartLimits=NO"

    # 2 Submit -> pend on just p1 with bad limit, then complete with good limit
    job_id = atf.submit_job_sbatch(
        f'-p p1 {flag}{custom_val_fail} --wrap "hostname&" -o /dev/null'
    )
    satisfy_pending_job_limit(job_id, limit_name, custom_val_fail)
    assert (
        atf.get_job_parameter(job_id, "JobState", quiet=True) == "COMPLETED"
    ), f"Job should submit, pend, then complete on p1 with updated limit {limit_name} on partition p1 to passing values with EnforcePartLimits=NO"

    # 3 Submit -> complete on p2 with no limit set
    assert (
        atf.submit_job_sbatch(
            f'-p p2, {flag}{custom_val_fail} --wrap "hostname" -o /dev/null'
        )
        != 0
    ), f"Job should pass on p2 despite {limit_name} limit not met on the required partition p1 with EnforcePartLimits=NO"

    # 4 Submit -> complete on p1 with p1 limit met
    assert (
        atf.submit_job_sbatch(f'-p p1 {flag}{val_pass} --wrap "hostname" -o /dev/null')
        != 0
    ), f"Job should pass on p1 due to {limit_name} limit met on the required partition p1 with EnforcePartLimits=NO"

    # 5 Submit -> complete on p1,p2 with p1 limit met
    assert (
        atf.submit_job_sbatch(
            f'-p p1,p2 {flag}{val_pass} --wrap "hostname" -o /dev/null'
        )
        != 0
    ), f"Job should pass on p1,p2 due to {limit_name} limit met on the required partition p1 with EnforcePartLimits=NO"


def enforce_NO_QOS(limit_name, flag, val_fail, val_pass):
    flag = "--qos="
    val_pass = "good_qos"
    val_fail = "bad_qos"

    # 1 Submit -> pend on p1,p2 with bad p1 limit set -> complete with p1 limit met
    job_id = atf.submit_job_sbatch(
        f'-p p1,p2 {flag}{val_fail} --wrap "hostname&" -o /dev/null'
    )
    satisfy_pending_job_limit(job_id, limit_name, f"{val_pass},{val_fail}")
    assert (
        atf.get_job_parameter(job_id, "JobState", quiet=True) == "COMPLETED"
    ), f"Job should submit, pend, then complete on p1,p2 with updated limit {limit_name} on partition p1 to passing valueswith EnforcePartLimits=NO"

    # Reset partition QOS
    atf.run_command(
        f"scontrol update partitionname=p1 {limit_name}={val_pass}",
        user=atf.properties["slurm-user"],
        fatal=True,
        quiet=True,
    )

    # 2 Submit -> pend on just p1 with bad limit, then complete with good limit
    job_id = atf.submit_job_sbatch(
        f'-p p1 {flag}{val_fail} --wrap "hostname&" -o /dev/null'
    )
    satisfy_pending_job_limit(job_id, limit_name, f"{val_pass},{val_fail}")
    assert (
        atf.get_job_parameter(job_id, "JobState", quiet=True) == "COMPLETED"
    ), f"Job should submit, pend, then complete on p1 with updated limit {limit_name} on partition p1 to passing values with EnforcePartLimits=NO"

    # Reset partition QOS
    atf.run_command(
        f"scontrol update partitionname=p1 {limit_name}={val_pass}",
        user=atf.properties["slurm-user"],
        fatal=True,
        quiet=True,
    )

    # 3 Submit -> complete on p2 with no limit set
    assert (
        atf.submit_job_sbatch(f'-p p2, {flag}{val_fail} --wrap "hostname" -o /dev/null')
        != 0
    ), f"Job should pass on p2 despite {limit_name} limit not met on the required partition p1 with EnforcePartLimits=NO"

    # 4 Submit -> complete on p1 with p1 limit met
    assert (
        atf.submit_job_sbatch(f'-p p1 {flag}{val_pass} --wrap "hostname" -o /dev/null')
        != 0
    ), f"Job should pass on p1 due to {limit_name} limit met on the required partition p1 with EnforcePartLimits=NO"

    # 5 Submit -> complete on p1,p2 with p1 limit met
    assert (
        atf.submit_job_sbatch(
            f'-p p1,p2 {flag}{val_pass} --wrap "hostname" -o /dev/null'
        )
        != 0
    ), f"Job should pass on p1,p2 due to {limit_name} limit met on the required partition p1 with EnforcePartLimits=NO"


# Tests:
@pytest.mark.parametrize("limit_name", limits_dict.keys())
def test_ALL(limit_name, setup_account, cancel_jobs):
    """Verify jobs are accepted and rejected with EnforePartLimits=ALL"""

    set_enforce_part_limits_policy("ALL")
    value = limits_dict[limit_name]
    set_partition_limit(limit_name, value["_set"])
    enforce_ALL(limit_name, value["flag"], value["fail"], value["pass"])


@pytest.mark.parametrize("limit_name", limits_dict.keys())
def test_ANY(limit_name, setup_account, cancel_jobs):
    """Verify jobs are accepted and rejected with EnforePartLimits=ANY"""

    set_enforce_part_limits_policy("ANY")
    value = limits_dict[limit_name]
    set_partition_limit(limit_name, value["_set"])
    enforce_ANY(limit_name, value["flag"], value["fail"], value["pass"])


@pytest.mark.parametrize("limit_name", limits_dict.keys())
def test_NO(limit_name, setup_account, cancel_jobs):
    """Verify jobs are accepted and rejected with EnforePartLimits=NO"""

    set_enforce_part_limits_policy("NO")
    value = limits_dict[limit_name]
    set_partition_limit(limit_name, value["_set"])

    if limit_name == "AllowQos":
        enforce_NO_QOS(limit_name, value["flag"], value["fail"], value["pass"])
    else:
        enforce_NO(limit_name, value["flag"], value["fail"], value["pass"])