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

# import re

dw_wlm_cli_client = f"{atf.properties['slurm-sbin-dir']}/dw_wlm_cli"
dwstat_client = f"{atf.properties['slurm-sbin-dir']}/dwstat"


def require_datawarp_clients():
    # For the time being, we are installing the clients if necessary,
    # but not removing them afterwards. Later on, we could generalize
    # {backup,restore}_config_file to work for any file

    if (
        atf.run_command_exit(
            f"test -f {dw_wlm_cli_client}",
            user=atf.properties["slurm-user"],
            quiet=True,
        )
        != 0
    ):
        if atf.properties["auto-config"]:
            atf.run_command(
                f"cp {atf.properties['slurm-source-dir']}/src/plugins/burst_buffer/datawarp/dw_wlm_cli {dw_wlm_cli_client}",
                user=atf.properties["slurm-user"],
                fatal=True,
                quiet=True,
            )
        else:
            pytest.skip(
                f"This test requires the dw_wlm_cli client to be installed in {atf.properties['slurm-sbin-dir']}",
                allow_module_level=True,
            )

    if (
        atf.run_command_exit(
            f"test -f {dwstat_client}", user=atf.properties["slurm-user"], quiet=True
        )
        != 0
    ):
        if atf.properties["auto-config"]:
            atf.run_command(
                f"cp {atf.properties['slurm-source-dir']}/src/plugins/burst_buffer/datawarp/dwstat {dwstat_client}",
                user=atf.properties["slurm-user"],
                fatal=True,
                quiet=True,
            )
        else:
            pytest.skip(
                f"This test requires the dwstat client to be installed in {atf.properties['slurm-sbin-dir']}",
                allow_module_level=True,
            )


@pytest.fixture(scope="module", autouse=True)
def setup():
    atf.require_config_parameter("BurstBufferType", "burst_buffer/datawarp")
    # atf.require_config_parameter('DebugFlags', 'burstbuffer')
    atf.require_config_parameter(
        "AllowUsers", f"{atf.properties['test-user']}", source="burst_buffer"
    )
    atf.require_config_parameter_includes("Flags", "EmulateCray", source="burst_buffer")
    atf.require_config_parameter_includes(
        "Flags", "EnablePersistent", source="burst_buffer"
    )
    require_datawarp_clients()
    atf.require_config_parameter(
        "GetSysState", dw_wlm_cli_client, source="burst_buffer"
    )
    atf.require_config_parameter("GetSysStatus", dwstat_client, source="burst_buffer")
    atf.require_slurm_running()


def test_create_and_use_burst_buffer(tmp_path):
    """Create and use burst buffer"""

    bb_use_script = str(tmp_path / "bb_use.sh")
    bb_use_output = str(tmp_path / "bb_use.out")
    atf.make_bash_script(
        bb_use_script,
        """#DW persistentdw name=test_buffer
scontrol show burst
""",
    )
    bb_create_script = str(tmp_path / "bb_create.sh")
    bb_create_output = str(tmp_path / "bb_create.out")
    atf.make_bash_script(
        bb_create_script,
        """#BB create_persistent name=test_buffer capacity=48 access=striped type=scratch
scontrol show burst
""",
    )

    # Submit a job to use the persisent burst buffer
    bb_use_job_id = atf.submit_job_sbatch(
        f"-N1 -t1 -o {bb_use_output} {bb_use_script}", fatal=True
    )

    # Submit a job to create the persisent burst buffer
    bb_create_job_id = atf.submit_job_sbatch(
        f"-N1 -t1 -o {bb_create_output} {bb_create_script}", fatal=True
    )

    # The burst buffer creation should complete first
    assert atf.wait_for_job_state(
        bb_create_job_id, "DONE", timeout=660
    ), f"Burst buffer creation job ({bb_create_job_id}) did not run"

    # The burst buffer usage should then complete
    assert atf.wait_for_job_state(
        bb_use_job_id, "DONE", timeout=660
    ), f"Burst buffer usage job ({bb_create_job_id}) did not run"

    # Ensure the burst buffer usage job ran after the buffer creation
    assert atf.wait_for_file(bb_use_output, timeout=660)
    assert (
        atf.run_command_exit(f"grep test_buffer {bb_use_output}") == 0
    ), "Job using burst buffer ran before bb creation"

    assert atf.wait_for_file(bb_create_output, timeout=660)
    assert (
        atf.run_command_exit(f"grep test_buffer {bb_create_output}") == 0
    ), "Job creating burst buffer failed to do so"


def test_remove_burst_buffer(tmp_path):
    """Delete burst buffer"""

    bb_delete_script = str(tmp_path / "bb_delete.sh")
    bb_delete_output = str(tmp_path / "bb_delete.out")
    atf.make_bash_script(
        bb_delete_script,
        """#BB destroy_persistent name=test_buffer
scontrol show burst
""",
    )

    # Submit a job to delete the persisent burst buffer
    bb_delete_job_id = atf.submit_job_sbatch(
        f"-N1 -t1 -o {bb_delete_output} {bb_delete_script}", fatal=True
    )

    # The burst buffer deletion job should complete
    assert atf.wait_for_job_state(
        bb_delete_job_id, "DONE", timeout=660
    ), f"Burst buffer deletion job ({bb_delete_job_id}) did not run"

    assert atf.wait_for_file(bb_delete_output, timeout=660)
    assert (
        atf.run_command_exit(f"grep test_buffer {bb_delete_output}") != 0
    ), "Job deleting burst buffer failed to do so"