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


# Setup
@pytest.fixture(scope="module", autouse=True)
def setup():
    atf.require_auto_config('wants to create sh files')
    atf.require_config_parameter('GresTypes', 'gpu')
    atf.require_config_parameter('SelectType', 'select/cons_tres')
    atf.require_nodes(2, [('Gres', 'gpu:2'), ('CPUs', 4)])

    gpu_file_pattern = make_gpu_files(4)
    atf.require_config_parameter('Name', {'gpu' : {'File' :
        gpu_file_pattern}}, source='gres')

    atf.make_bash_script('step.sh', '''
            echo 'NODE_ID:'$SLURM_NODEID 'STEP_ID:'$SLURM_STEP_ID
            sleep 1
            exit 0''')
    step_path = f"{atf.module_tmp_path}/step.sh"
    atf.make_bash_script('job.sh', f"""
            echo 'step 1'
            srun --exact -n2 --gpus=2 --gpus-per-node=0 --mem=0 {step_path} &
            echo 'step 2'
            srun --exact -n2 --gpus=2 --gpus-per-node=0 --mem=0 {step_path} &
            wait
            exit 0""")
    atf.require_slurm_running()


def make_gpu_files(count):
    """Make files in the tmp path for gpu's to point to
    Returns pattern TMP/gpu[1-COUNT]"""

    for i in range(1, count + 1):
        atf.run_command(f"touch {atf.module_tmp_path}/gpu{i}")
    return f"{atf.module_tmp_path}/gpu[1-{count}]"


def validate_job(job_id, out_file):
    """Both jobs have the same validation process"""

    assert job_id != 0, 'Expect job to submit properly'
    atf.wait_for_job_state(job_id, 'DONE')
    with open(out_file, 'r') as f:
        out = f.read()
        assert out.startswith('step 1'), 'Expect step 1 to start without issues'
        assert 'step 2' in out, 'Expect step 2 to start'
        assert 'NODE_ID:0 STEP_ID:0' in out, \
                'Expect the first node to be in first step'
        assert 'NODE_ID:1 STEP_ID:0' in out, \
                'Expect the second node to be in first step'
        assert 'NODE_ID:0 STEP_ID:1' in out, \
                'Expect the first node to be in second step'
        assert 'NODE_ID:1 STEP_ID:1' in out, \
                'Expect the second node to be in second step'


def test_exact_gpu_full_resources():
    """Test --exact with all of node resources"""

    out_file = f"{atf.module_tmp_path}/out"
    job_path = f"{atf.module_tmp_path}/job.sh"
    job_id = atf.submit_job_sbatch(f"--cpus-per-gpu=2 --gpus-per-node=2 -N2 \
        -n4 -t1 --output={out_file} {job_path}")
    validate_job(job_id, out_file)


def test_exact_gpu_parial_resources():
    """Test --exact with partial node resources"""

    job_path = f"{atf.module_tmp_path}/job.sh"
    out_file = f"{atf.module_tmp_path}/out"
    job_id = atf.submit_job_sbatch(
        f"--cpus-per-gpu=1 --gpus-per-node=2 -N2 -n4 -t1 \
          --output={out_file} {job_path}")
    validate_job(job_id, out_file)
