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


@pytest.fixture(scope="module", autouse=True)
def setup():
    atf.require_accounting()
    atf.require_config_parameter("PriorityType", "priority/multifactor")
    atf.require_slurm_running()


@pytest.fixture(scope="module")
def default_partition():
    return atf.default_partition()


@pytest.fixture(scope="module")
def queued_job(default_partition):
    nodes_dict = atf.get_nodes()
    total_cpus = 0

    for node_name in nodes_dict:
        node_dict = nodes_dict[node_name]
        node_partitions = node_dict["partitions"]
        if default_partition in node_partitions and "IDLE" in node_dict["state"]:
            total_cpus += node_dict["cpus"]

    atf.submit_job_sbatch(
        f'--output=/dev/null --error=/dev/null -n {total_cpus} --exclusive --wrap="sleep 600"',
        fatal=True,
    )
    queued_job_id = atf.submit_job_sbatch(
        f'--output=/dev/null --error=/dev/null -n {total_cpus} --exclusive --wrap="sleep 600"',
        fatal=True,
    )

    return queued_job_id


def test_noheader(queued_job, default_partition):
    """Verify sprio --noheader option"""

    output = atf.run_command_output(f"sprio --noheader -j {queued_job}", fatal=True)
    assert re.search(r"JOBID|PARTITION|PRIORITY|SITE", output) is None
    assert re.search(rf"{queued_job}\s+{default_partition}", output) is not None


def test_jobs(queued_job, default_partition):
    """Verify sprio --jobs option"""

    output = atf.run_command_output(f"sprio --jobs {queued_job}", fatal=True)
    assert re.search(r"JOBID\s+PARTITION\s+PRIORITY\s+SITE", output) is not None
    assert re.search(rf"{queued_job}\s+{default_partition}", output) is not None


def test_long(queued_job, default_partition):
    """Verify sprio --long option"""

    user_name = atf.get_user_name()
    output = atf.run_command_output(f"sprio --long -j {queued_job}", fatal=True)
    assert (
        re.search(
            r"JOBID\s+PARTITION\s+USER\s+ACCOUNT\s+PRIORITY\s+SITE\s+AGE\s+ASSOC\s+FAIRSHARE\s+JOBSIZE\s+PARTITION\s+QOSNAME\s+QOS\s+NICE\s+TRES[\r\n]",
            output,
        )
        is not None
    )
    assert (
        re.search(
            rf"{queued_job}\s+{default_partition}\s+{user_name}\s+\S+(?:\s+\d+){{7,}}\s+\S+(?:\s+\d+){{2,}}\s*\S*",
            output,
        )
        is not None
    )


def test_norm(queued_job, default_partition):
    """Verify sprio --norm option"""

    output = atf.run_command_output(f"sprio --norm -j {queued_job}", fatal=True)
    assert re.search(r"JOBID\s+PARTITION\s+PRIORITY", output) is not None
    assert (
        re.search(rf"{queued_job}\s+{default_partition}\s+[\d\.]+", output) is not None
    )


def test_format(queued_job, default_partition):
    """Verify sprio --norm option"""

    user_name = atf.get_user_name()
    output = atf.run_command_output(
        f'sprio --format "%.15i %.8u %.10y %.10Y %.10S %.10a %.10A %.10b %.10B %.10f %.10F %.10j %.10J %.10p %.10P %.10q %.10Q %.6N" -j {queued_job}',
        fatal=True,
    )
    assert (
        re.search(
            r"JOBID\s+USER\s+PRIORITY\s+PRIORITY\s+SITE\s+AGE\s+AGE\s+ASSOC\s+ASSOC",
            output,
        )
        is not None
    )
    assert (
        re.search(
            r"FAIRSHARE\s+FAIRSHARE\s+JOBSIZE\s+JOBSIZE\s+PARTITION\s+PARTITION\s+QOS\s+QOS\s+NICE",
            output,
        )
        is not None
    )
    assert (
        re.search(
            rf"{queued_job}\s+{user_name}\s+[\d\.]+\s+\d+\s+\d+\s+[\d\.]+\s+\d+\s+[\d\.]+\s+\d+\s+[\d\.]+\s+\d+\s+[\d\.]+\s+\d+\s+[\d\.]+\s+\d+\s+\d+",
            output,
        )
        is not None
    )


def test_user(queued_job, default_partition):
    """Verify sprio --user option"""

    user_name = atf.get_user_name()
    output = atf.run_command_output(f"sprio --user {user_name}", fatal=True)
    assert re.search(r"JOBID\s+PARTITION\s+USER", output) is not None
    assert (
        re.search(rf"{queued_job}\s+{default_partition}\s+{user_name}", output)
        is not None
    )


def test_verbose(queued_job, default_partition):
    """Verify sprio --verbose option"""

    output = atf.run_command_output(f"sprio --verbose -j {queued_job}", fatal=True)
    assert re.search(r"format\s+= \(null\)", output) is not None
    assert re.search(r"job_flag\s+= 1", output) is not None
    assert re.search(rf"jobs\s+= {queued_job}", output) is not None
    assert re.search(r"partition\s+= \(null\)", output) is not None
    assert re.search(r"users\s+= \(null\)", output) is not None
    assert re.search(r"verbose\s+= 1", output) is not None
    assert re.search(r"JOBID\s+PARTITION\s+PRIORITY", output) is not None
    assert re.search(rf"{queued_job}\s+{default_partition}\s+\d+", output) is not None