############################################################################
# Copyright (C) SchedMD LLC.
############################################################################
import os
import re
import time

import pytest

import atf

pytestmark = pytest.mark.slow


# Setup
@pytest.fixture(scope="module", autouse=True)
def setup():
    atf.require_auto_config("wants to set the mail program")
    atf.require_slurm_running()


@pytest.fixture(scope="module")
def mail_program_out():
    mail_tmp_dir = str(atf.module_tmp_path)
    mail_program = str(atf.module_tmp_path / "mail_program.sh")
    mail_program_out = str(atf.module_tmp_path / "mail_program.out")

    os.chmod(mail_tmp_dir, 0o777)
    atf.make_bash_script(
        mail_program,
        rf"""echo SLURM_JOB_ID=$SLURM_JOB_ID SLURM_JOB_USER=$SLURM_JOB_USER SLURM_JOB_MAIL_TYPE=$SLURM_JOB_MAIL_TYPE >> {mail_program_out}""",
    )
    os.chmod(mail_program, 0o777)
    atf.set_config_parameter("MailProg", str(mail_program))

    return mail_program_out


def test_mail_type_and_mail_user(mail_program_out):
    """Test of mail options (--mail-type and --mail-user options)."""

    slurm_user = "slurm-user"
    job_id = atf.submit_job_srun(
        f"--mail-type=all --mail-user={atf.properties[slurm_user]} id"
    )
    atf.wait_for_file(mail_program_out, fatal=True)
    output = atf.run_command_output(f"cat {mail_program_out}")
    assert (
        re.findall(
            rf"SLURM_JOB_ID={job_id} SLURM_JOB_USER={atf.properties[slurm_user]} SLURM_JOB_MAIL_TYPE=Began",
            output,
        )
        is not None
    ), "Start mail not sent"
    assert (
        re.findall(
            rf"SLURM_JOB_ID={job_id} SLURM_JOB_USER={atf.properties[slurm_user]} SLURM_JOB_MAIL_TYPE=Ended",
            output,
        )
        is not None
    ), "End mail not sent"

    job_id = atf.submit_job_sbatch(
        f'--wrap="srun --mail-type=END --mail-user={atf.properties[slurm_user]} sleep 100"'
    )
    atf.wait_for_job_state(job_id, "RUNNING")
    atf.run_command(f"scancel {job_id}")
    output = atf.run_command_output(f"cat {mail_program_out}")
    assert (
        re.findall(
            rf"SLURM_JOB_ID=\d+ SLURM_JOB_USER={atf.properties[slurm_user]} SLURM_JOB_MAIL_TYPE=Ended",
            output,
        )
        is not None
    ), "End mail not sent after job was canceled"

    job_id = atf.submit_job_srun(
        f"-t1 --mail-type=ALL,TIME_LIMIT,TIME_LIMIT_90,TIME_LIMIT_80,TIME_LIMIT_50 --mail-user={atf.properties[slurm_user]} sleep 300",
        timeout=120,
        xfail=True,
    )
    time.sleep(5)
    output = atf.run_command_output(f"cat {mail_program_out}")
    assert (
        re.findall(
            rf"SLURM_JOB_ID={job_id} SLURM_JOB_USER={atf.properties[slurm_user]} SLURM_JOB_MAIL_TYPE=Began",
            output,
        )
        is not None
    ), "Start mail not sent for timeout test"
    assert (
        re.findall(
            rf"SLURM_JOB_ID={job_id} SLURM_JOB_USER={atf.properties[slurm_user]} SLURM_JOB_MAIL_TYPE=Failed",
            output,
        )
        is not None
    ), "Failed mail not sent for timeout test"
    assert (
        re.findall(
            rf"SLURM_JOB_ID={job_id} SLURM_JOB_USER={atf.properties[slurm_user]} SLURM_JOB_MAIL_TYPE=Reached time limit",
            output,
        )
        is not None
    ), "Time limit mail not sent for timeout test"
    assert (
        re.findall(
            rf"SLURM_JOB_ID={job_id} SLURM_JOB_USER={atf.properties[slurm_user]} SLURM_JOB_MAIL_TYPE=Reached 90% of time limit",
            output,
        )
        is not None
    ), "Reached 90% of time limit mail not sent for timeout test"
    assert (
        re.findall(
            rf"SLURM_JOB_ID={job_id} SLURM_JOB_USER={atf.properties[slurm_user]} SLURM_JOB_MAIL_TYPE=Reached 80% of time limit",
            output,
        )
        is not None
    ), "Reached 80% of time limit mail not sent for timeout test"
    assert (
        re.findall(
            rf"SLURM_JOB_ID={job_id} SLURM_JOB_USER={atf.properties[slurm_user]} SLURM_JOB_MAIL_TYPE=Reached 50% of time limit",
            output,
        )
        is not None
    ), "Reached 50% of time limit mail not sent for timeout test"