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


# Globals
user_name = atf.get_user_name()
test_name = os.path.splitext(os.path.basename(__file__))[0]
acct1 = f"{test_name}_acct1"
acct2 = f"{test_name}_acct2"
name1 = f"{test_name}_job1"
name2 = f"{test_name}_job2"
nodes = ["node1", "node2", "node3"]
partitions = [f"{test_name}_part1", f"{test_name}_part2", f"{test_name}_part3"]
qos1 = f"{test_name}_qos1"
qos2 = f"{test_name}_qos2"
reservations = [f"{test_name}_resv1", f"{test_name}_resv2", f"{test_name}_resv3"]
wckey = f"{test_name}_wckey1"


@pytest.fixture(scope="module", autouse=True)
def setup():
    # Test needs to run a 2 het job with 3 components, 2 jobs, 2 arrays of 3 jobs, and 2 arrays with 2 jobs.
    # We need 18 nodes to be able to run with select_linear, and 7+ CPUs per node
    atf.require_nodes(18, [("CPUs", 8), ("RealMemory", 100)])
    atf.require_accounting(True)
    atf.require_config_parameter("TrackWcKey", "yes")
    atf.require_config_parameter("TrackWcKey", "yes", source="slurmdbd")
    # Reducing bf_interval makes the test faster when using het jobs.
    atf.add_config_parameter_value("SchedulerParameters", "bf_interval=2")
    atf.require_slurm_running()

    # Basic account and user setup
    sacctmgr_acct = f"account {acct1} {acct2}"
    sacctmgr_user = f"user {user_name} account={acct1},{acct2}"
    atf.run_command(
        f"sacctmgr -i add {sacctmgr_acct}",
        user=atf.properties["slurm-user"],
        fatal=True,
    )
    atf.run_command(
        f"sacctmgr -i add {sacctmgr_user}",
        user=atf.properties["slurm-user"],
        fatal=True,
    )

    yield

    atf.run_command(
        f"sacctmgr -i del {sacctmgr_user} ",
        user=atf.properties["slurm-user"],
    )
    atf.run_command(
        f"sacctmgr -i del {sacctmgr_acct} ",
        user=atf.properties["slurm-user"],
    )


# Cancel all jobs before and after the test
@pytest.fixture(scope="function", autouse=True)
def cancel_jobs():
    atf.cancel_all_jobs()
    yield
    atf.cancel_all_jobs()


#
# Main fixtures
#
@pytest.fixture
def setup_partitions():
    for part in partitions:
        atf.run_command(
            f"scontrol create partitionname={part} Nodes=ALL",
            user=atf.properties["slurm-user"],
            fatal=True,
        )

    yield

    atf.cancel_all_jobs()
    for part in partitions:
        atf.run_command(
            f"scontrol delete partitionname={part}",
            user=atf.properties["slurm-user"],
        )


@pytest.fixture
def setup_qos():
    atf.run_command(
        f"sacctmgr -i add qos {qos1} {qos2}",
        user=atf.properties["slurm-user"],
        fatal=True,
    )
    atf.run_command("sacctmgr show assoc tree")
    atf.run_command(
        f"sacctmgr -i mod user {user_name} set qos+={qos1},{qos2}",
        user=atf.properties["slurm-user"],
        fatal=True,
    )

    yield

    atf.cancel_all_jobs()
    atf.run_command(
        f"sacctmgr -i mod user {user_name} set qos-={qos1},{qos2}",
        user=atf.properties["slurm-user"],
    )
    atf.run_command(
        f"sacctmgr -i del qos {qos1} {qos2}",
        user=atf.properties["slurm-user"],
    )


@pytest.fixture
def setup_reservations():
    for resv in reservations:
        # Hetjobs need 3 nodes in the reservation
        atf.run_command(
            f"scontrol create reservationname={resv} starttime=now duration=5 users={user_name} nodecnt=3",
            user=atf.properties["slurm-user"],
            fatal=True,
        )

    yield

    atf.cancel_all_jobs()
    for resv in reservations:
        atf.run_command(
            f"scontrol delete reservationname={resv}",
            user=atf.properties["slurm-user"],
        )


@pytest.fixture
def setup_wckeys():
    atf.run_command(
        f"sacctmgr -i add user {user_name} set wckey={wckey}",
        user=atf.properties["slurm-user"],
        fatal=True,
    )

    yield

    atf.cancel_all_jobs()
    atf.run_command(
        f"sacctmgr -i del user {user_name} wckey={wckey}",
        user=atf.properties["slurm-user"],
    )


#
# Custom fixtures, depending on main ones
#
@pytest.fixture
def setup_partition2_down(setup_partitions):
    """Set partitions[2] down so we ensure that matching jobs will run on
    partition[0] and mismatching ones in partition[1]"""
    atf.run_command(
        f"scontrol update partitionname={partitions[2]} state=down",
        user=atf.properties["slurm-user"],
        fatal=True,
    )


@pytest.fixture
def setup_reservations_used(setup_reservations):
    """'Run an exclusive job in reservations[1] to ensure that jobs run in
    reservations[0]"""
    job_id = atf.submit_job_sbatch(
        f"--exclusive --reservation={reservations[1]} --wrap 'srun sleep infinity'",
        fatal=True,
    )
    atf.wait_for_job_state(job_id, "RUNNING", fatal=True, timeout=60)


#
# Helper functions
#
def submit_job(job_list, opt, opt_val, expected_state):
    test_opt = ""
    other_opts = ""

    if opt is not None:
        test_opt = f"--{opt}={opt_val}"
    if expected_state == "PENDING":
        other_opts = "--hold"

    job_list.append(
        (
            atf.submit_job_sbatch(
                f"-n1 -c1 --mem=10 {test_opt} {other_opts} --wrap 'srun sleep infinity'",
                fatal=True,
            )
        )
    )


def submit_het_job(job_list, opt, opt_val_1, opt_val_2, expected_state):
    het_opts_1 = "-n1 -c1 --mem=10"
    het_opts_2 = "-n1 -c1 --mem=10"
    other_opts = ""

    if expected_state == "PENDING":
        other_opts += " --hold"

    # Nodelist only for the first component as each component needs a different node
    if opt != "nodelist":
        het_opts_1 += f" --{opt}={opt_val_1}"
        het_opts_2 += f" --{opt}={opt_val_2}"
    else:
        het_opts_1 += f" --{opt}={opt_val_1}"

    # We want to cancel the 2nd component because canceling the first one cancels all components
    job_id = atf.submit_job_sbatch(
        f"{other_opts} {het_opts_2} : {het_opts_1} --wrap 'srun sleep infinity'",
        fatal=True,
    )
    job_list.append(job_id)
    job_list.append(job_id + 1)


def wait_for_jobs(job_list, expected_state, timeout):
    for job_id in job_list:
        atf.wait_for_job_state(job_id, expected_state, fatal=True, timeout=timeout)


def run_scancel(opt, opt_filter, job_list):
    jobs_str = " ".join(map(str, job_list))
    atf.run_command(f"scancel --ctld --{opt}={opt_filter} {jobs_str}")


# Parameters of test_filter and test_filter_hetjob
parameters = [
    # Test --account
    (None, "account", acct1, "account", acct1, acct2, "RUNNING", "RUNNING"),
    # Test --jobname
    (None, "jobname", name1, "job-name", name1, name2, "RUNNING", "RUNNING"),
    # Test --nodelist with single node (only works for running)
    (
        None,
        "nodelist",
        nodes[0],
        "nodelist",
        nodes[0],
        nodes[1],
        "RUNNING",
        "RUNNING",
    ),
    # Test multi-node nodelist (scancel --nodelist only works for running)
    (
        None,
        "nodelist",
        nodes[0],
        "nodelist",
        f"{nodes[0]},{nodes[1]}",
        f"{nodes[1]},{nodes[2]}",
        "RUNNING",
        "RUNNING",
    ),
    # Test qos
    ("setup_qos", "qos", qos1, "qos", qos1, qos2, "RUNNING", "RUNNING"),
    # Test --wckey
    # TODO: This wckey testing is very simple. We should also test with DefaultWckey.
    ("setup_wckeys", "wckey", wckey, "wckey", wckey, "", "RUNNING", "RUNNING"),
    # Test --reservation with a single reservation for pending jobs
    (
        "setup_reservations",
        "reservation",
        reservations[0],
        "reservation",
        reservations[0],
        reservations[1],
        "PENDING",
        "PENDING",
    ),
    # Test --reservation with a single reservation for running jobs
    (
        "setup_reservations",
        "reservation",
        reservations[0],
        "reservation",
        reservations[0],
        reservations[1],
        "RUNNING",
        "RUNNING",
    ),
    # Test --reservation with multiple reservations for pending jobs
    (
        "setup_reservations",
        "reservation",
        reservations[0],
        "reservation",
        reservations[0] + "," + reservations[1],
        reservations[1] + "," + reservations[2],
        "PENDING",
        "PENDING",
    ),
    # Test --reservation with multiple reservation for running jobs
    # Needs an exclusive job in reservation[1] to ensure that job runs into [0]
    (
        "setup_reservations_used",
        "reservation",
        reservations[0],
        "reservation",
        reservations[0] + "," + reservations[1],
        reservations[2],
        "RUNNING",
        "RUNNING",
    ),
    # Test --partition with a single partition for pending jobs
    (
        "setup_partitions",
        "partition",
        partitions[0],
        "partition",
        partitions[0],
        partitions[1],
        "PENDING",
        "PENDING",
    ),
    # Test --partition with a single partition for running jobs
    (
        "setup_partitions",
        "partition",
        partitions[0],
        "partition",
        partitions[0],
        partitions[1],
        "RUNNING",
        "RUNNING",
    ),
    # Test --partition with multiple partitions for pending jobs
    (
        "setup_partitions",
        "partition",
        partitions[0],
        "partition",
        partitions[0] + "," + partitions[1],
        partitions[1] + "," + partitions[2],
        "PENDING",
        "PENDING",
    ),
    # Test --partition with multiple partitions for running jobs
    # Needs partition[2] down to ensure that matching jobs run in [0] and
    # mismatching in [1].
    (
        "setup_partition2_down",
        "partition",
        partitions[0],
        "partition",
        partitions[0] + "," + partitions[2],
        partitions[1] + "," + partitions[2],
        "RUNNING",
        "RUNNING",
    ),
    # Test --state for pending in partition[2] and running in partition[0]
    # Needs partition[2] down to ensure that matching jobs run in [0] and
    # mismatching in [1].
    (
        "setup_partition2_down",
        "state",
        "PENDING",
        "partition",
        partitions[2],
        partitions[0],
        "PENDING",
        "RUNNING",
    ),
    # Test --state for running in partition[0] and pending in partition[2]
    # Needs partition[2] down to ensure that matching jobs run in [0] and
    # mismatching in [1].
    (
        "setup_partition2_down",
        "state",
        "RUNNING",
        "partition",
        partitions[0],
        partitions[2],
        "RUNNING",
        "PENDING",
    ),
]


@pytest.mark.parametrize(
    "fixture,scancel_opt,scancel_val,sbatch_opt,sbatch_match_val,sbatch_mismatch_val,match_job_state,mismatch_job_state",
    parameters,
)
def test_filter(
    request,
    fixture,
    scancel_opt,
    scancel_val,
    sbatch_opt,
    sbatch_match_val,
    sbatch_mismatch_val,
    match_job_state,
    mismatch_job_state,
):
    # Custom fixture may be necessary
    if fixture is not None:
        request.getfixturevalue(fixture)

    matching_jobs = []
    mismatching_jobs = []
    # matching_hetjobs = []
    # mismatching_hetjobs = []
    # het_jobs = []

    logging.info(
        f"Test scancel --ctld --{scancel_opt}={scancel_val}, sbatch --{sbatch_opt}, job match:{sbatch_match_val}, job mismatch:{sbatch_mismatch_val}"
    )

    # Submit jobs
    submit_job(matching_jobs, sbatch_opt, sbatch_match_val, match_job_state)
    submit_job(mismatching_jobs, sbatch_opt, sbatch_mismatch_val, mismatch_job_state)

    # Waiting for jobs
    wait_for_jobs(matching_jobs, match_job_state, 60)
    wait_for_jobs(mismatching_jobs, mismatch_job_state, 60)

    # Just to show jobs in the logs
    atf.run_command_output(
        "echo ''; squeue --Format=JobId,Account,Name,Nodelist,Partition,Qos,Reservation,State,UserName,WCKey"
    )

    # Cancel jobs with --scancel_opt=scancel:val
    run_scancel(scancel_opt, scancel_val, [])

    # Wait for jobs matching the signal filter to finish
    # And verify that jobs not matching the signal filter are still in their expected state
    wait_for_jobs(matching_jobs, "CANCELLED", 10)
    wait_for_jobs(mismatching_jobs, mismatch_job_state, 60)


@pytest.mark.parametrize(
    "fixture,scancel_opt,scancel_val,sbatch_opt,sbatch_match_val,sbatch_mismatch_val,match_job_state,mismatch_job_state",
    parameters,
)
def test_filter_hetjobs(
    request,
    fixture,
    scancel_opt,
    scancel_val,
    sbatch_opt,
    sbatch_match_val,
    sbatch_mismatch_val,
    match_job_state,
    mismatch_job_state,
):
    # Custom fixture may be necessary
    if fixture is not None:
        request.getfixturevalue(fixture)

    matching_hetjobs = []
    mismatching_hetjobs = []
    het_jobs = []

    logging.info(
        f"Test scancel --ctld --{scancel_opt}={scancel_val}, sbatch --{sbatch_opt}, job match:{sbatch_match_val}, job mismatch:{sbatch_mismatch_val} for hetjobs"
    )

    # Submit hetjobs
    # There are two special cases to avoid:
    # - When canceling with --nodelist we may cancel first component and it will cancel the whole hetjob
    # - When canceling with --state we cannot distinguish between components
    if scancel_opt != "nodelist":
        submit_het_job(
            matching_hetjobs,
            sbatch_opt,
            sbatch_match_val,
            sbatch_match_val,
            match_job_state,
        )
        submit_het_job(
            mismatching_hetjobs,
            sbatch_opt,
            sbatch_mismatch_val,
            sbatch_mismatch_val,
            mismatch_job_state,
        )
    if scancel_opt != "state":
        submit_het_job(
            het_jobs, sbatch_opt, sbatch_match_val, sbatch_mismatch_val, match_job_state
        )

    # Waiting for hetjobs
    if scancel_opt != "nodelist":
        wait_for_jobs(matching_hetjobs, match_job_state, 60)
        wait_for_jobs(mismatching_hetjobs, mismatch_job_state, 60)
    if scancel_opt != "state":
        wait_for_jobs([het_jobs[1]], match_job_state, 60)
        wait_for_jobs([het_jobs[0]], mismatch_job_state, 60)

    # Just to show jobs in the logs
    atf.run_command_output(
        "echo ''; squeue --Format=JobId,Account,Name,Nodelist,Partition,Qos,Reservation,State,UserName,WCKey"
    )

    # Cancel jobs with --scancel_opt=scancel:val
    run_scancel(scancel_opt, scancel_val, [])

    # Same for hetjobs
    if scancel_opt != "nodelist":
        wait_for_jobs(matching_hetjobs, "CANCELLED", 10)
        wait_for_jobs(mismatching_hetjobs, mismatch_job_state, 60)
    if scancel_opt != "state":
        wait_for_jobs([het_jobs[1]], "CANCELLED", 10)
        wait_for_jobs([het_jobs[0]], mismatch_job_state, 60)


def test_signal_job_ids():
    test_jobs = []
    other_jobs = []
    submissions = []

    logging.info("Test scancel --signal with --ctld")

    script = "--wrap 'srun sleep infinity'"
    job_opt = "-n1 -c1 --mem=10"

    het_job_submit = f"{job_opt} : {job_opt} : {job_opt}"
    array_submit1 = f"{job_opt} --array=1-3"
    array_submit2 = f"{job_opt} --array=1-10:5"

    submissions.append(f"{job_opt} {script}")
    submissions.append(f"{het_job_submit} {script}")
    submissions.append(f"{array_submit1} {script}")
    submissions.append(f"{array_submit2} {script}")

    for s in submissions:
        test_jobs.append(atf.submit_job_sbatch(s, fatal=True))
        other_jobs.append(atf.submit_job_sbatch(s, fatal=True))

    wait_for_jobs(test_jobs, "RUNNING", 35)
    wait_for_jobs(other_jobs, "RUNNING", 35)

    run_scancel("signal", "SIGSTOP", test_jobs)
    wait_for_jobs(test_jobs, "STOPPED", 35)
    wait_for_jobs(other_jobs, "RUNNING", 35)

    run_scancel("signal", "SIGCONT", test_jobs)
    wait_for_jobs(test_jobs, "RUNNING", 35)
    wait_for_jobs(other_jobs, "RUNNING", 35)

    # SIGTERM will make the job end in FAILED, not CANCELLED
    run_scancel("signal", "SIGTERM", test_jobs)
    wait_for_jobs(test_jobs, "FAILED", 35)
    wait_for_jobs(other_jobs, "RUNNING", 35)

    # SIGKILL will make the job end in CANCELLED
    run_scancel("signal", "SIGKILL", other_jobs)
    wait_for_jobs(other_jobs, "CANCELLED", 35)

    # Sanity check that there are no running jobs
    output = atf.run_command_output("squeue --noheader")
    assert len(output) == 0


# TODO: Test --user with another test user, and specific job ids.
# TODO: Test specific job ids for a mix of pending and running jobs, also in arrays:
#       (1) the whole array job is pending
#       (2) the array job is partially pending and partially running: --array=1-10%3
#       (3) the whole array job is running
#       Also with different arrays expressions:
#       (a) 10: the entire array
#       (b) 10_5: single job of an array
#       (b) 10_[1,4-7,8-10]: complex expression