"""Test SLURM module."""

from pathlib import Path

from glicid_spawner.slurm import (
    SlurmCluster,
    SlurmCpu,
    SlurmGpu,
    SlurmNode,
    SlurmPartition,
    gres,
    sinfo,
    sinfo_filter,
    sinfo_from_file,
    sinfo_reader,
    sinfo_run,
    subprocess,
)

DATA = Path(__file__).parent / 'data'
SINFO_FILE = DATA / 'sinfo.txt'
SINFO_CONTENT = SINFO_FILE.read_text()


def test_slurm_dataclasses():
    """Test SLURM dataclasses formatter."""
    # CPU
    cpu = SlurmCpu(1, '2', 4.0)

    assert cpu.allocated == 1
    assert cpu.idle == 2
    assert cpu.total == 4

    assert isinstance(cpu.allocated, int)
    assert isinstance(cpu.idle, int)
    assert isinstance(cpu.total, int)

    # GPU
    gpu = SlurmGpu('fOo', '1')

    assert gpu  # __bool__
    assert gpu == str(gpu) == gpu.name == 'Foo'
    assert gpu.nb == 1

    assert isinstance(gpu.name, str)
    assert isinstance(gpu.nb, int)

    # Default values
    gpu = SlurmGpu()

    assert not gpu  # __bool__
    assert str(gpu) == gpu.name == 'None'
    assert gpu.nb == 0

    # Node
    node = SlurmNode(*'nautilus standard cnode001 completing 0/96/0/96 384000 (null)'.split())

    assert node.cluster == 'nautilus'
    assert node.partition == 'standard'
    assert node == str(node) == node.hostname == 'cnode001'
    assert node.state == 'completing'
    assert node.cpu.allocated == 0
    assert node.cpu.idle == 96
    assert node.cpu.total == 96
    assert node.mem == 384
    assert node.gpu.name == 'None'

    # Partition
    partition = SlurmPartition('standard', [node])

    assert partition == str(partition) == partition.name == 'standard'
    assert len(partition) == 1

    for _node in partition:
        assert _node == 'cnode001'

    _node = partition['cnode001']  # __getitem__
    assert isinstance(_node, SlurmNode)
    assert _node == 'cnode001'

    assert partition.gpus == 'None'
    assert partition.max_idle_cpu == 96
    assert partition.max_mem == 384

    # Cluster
    cluster = SlurmCluster('nautilus', [partition])

    assert cluster == str(cluster) == cluster.name == 'nautilus'
    assert len(cluster) == 1

    for _partition in cluster:
        assert _partition == 'standard'

    _partition = cluster['standard']  # __getitem__
    assert isinstance(_partition, SlurmPartition)
    assert _partition == 'standard'


def test_slurm_sinfo_run(monkeypatch):
    """Test SLURM SINFO run command."""
    monkeypatch.setattr(subprocess, 'check_output', lambda cmd: ' '.join(cmd).encode())

    assert sinfo_run() == (
        'sinfo '
        '--federation '
        '--noheader '
        '--responding '
        '--cluster=all '
        '--Format=Cluster,PartitionName,NodeHost,StateLong,CPUsState,Memory,Gres'
    )

    assert sinfo_run(username='john-doe') == (
        'su - john-doe -c '
        'sinfo '
        '--federation '
        '--noheader '
        '--responding '
        '--cluster=all '
        '--Format=Cluster,PartitionName,NodeHost,StateLong,CPUsState,Memory,Gres'
    )


def test_slurm_sinfo_reader():
    """Test SLURM SINFO reader."""
    nodes = sinfo_reader(SINFO_CONTENT)

    for node in nodes:
        assert isinstance(node, SlurmNode)

    node = nodes[0]

    assert node.cluster == 'cluster'
    assert node.partition == 'Devel'
    assert node.hostname == 'nazare001'
    assert node.state == 'idle'
    assert node.cpu.allocated == 0
    assert node.cpu.idle == node.cpu.total == 20
    assert node.mem == 128
    assert not node.gpu

    assert [node.cluster for node in nodes] == 7 * ['cluster'] + 12 * ['nautilus'] + 6 * ['waves']

    assert len([node for node in nodes if node.state in ('idle', 'mixed')]) == 10

    for node in nodes:
        if node.state == 'idle':
            assert node.cpu.allocated == 0
            assert node.cpu.idle > 0
        elif node.state == 'mixed':
            assert node.cpu.allocated > 0
            assert node.cpu.idle > 0
        elif node.state == 'allocated':
            assert node.cpu.allocated > 0
            assert node.cpu.idle == 0

    assert sum(node.mem for node in nodes) == 7_792

    assert [node.gpu.name for node in nodes if node.gpu] == [
        'T4',
        'A40',
        'P100',
        'K80',
        'P100',
        'A100',
        'A100',
        'A100_2g.10gb',
        'A100',
        'A100',
        'A100_2g.10gb',
        'A100',
        'A100',
    ]

    assert [node.gpu.nb for node in nodes if node.gpu] == [2, 2, 2, 4, 1, 1, 2, 6, 1, 2, 6, 2, 2]


def test_slurm_sinfo_filter(monkeypatch):
    """Test SLURM SINFO filtered resources."""
    resources = sinfo_reader(SINFO_CONTENT)

    clusters = sinfo_filter(resources)

    assert isinstance(clusters, dict)
    assert len(clusters) == 3
    assert list(clusters) == ['cluster', 'nautilus', 'waves']  # __eq__ on cluster.name

    assert [len(partitions) for partitions in clusters.values()] == [2, 4, 2]

    nautilus = clusters['nautilus']

    assert isinstance(nautilus, SlurmCluster)
    assert len(nautilus) == 4
    assert nautilus.partitions == ['standard', 'gpu', 'visu', 'all']  # __eq__ on partition.name

    gpu = nautilus['gpu']

    assert len(gpu) == 2
    assert [node.hostname for node in gpu] == ['gnode1', 'gnode2']
    assert [node.cpu.allocated for node in gpu] == [4, 0]
    assert [node.cpu.idle for node in gpu] == [92, 96]
    assert [node.mem for node in gpu] == [768, 256]
    assert [node.gpu.name for node in gpu] == ['A100', 'A100']
    assert [node.gpu.nb for node in gpu] == [1, 2]

    # Get only `idle` nodes
    clusters = sinfo_filter(resources, with_states=('idle'))

    assert list(clusters) == ['cluster', 'nautilus']
    assert [len(partitions) for partitions in clusters.values()] == [1, 3]

    # Discard clusters without partition available
    clusters = sinfo_filter(resources, with_states=('completing'))

    assert list(clusters) == ['nautilus']
    assert [len(partitions) for partitions in clusters.values()] == [2]


def test_slurm_sinfo_from_file(monkeypatch):
    """Test SLURM SINFO resources from file."""
    resources = sinfo_from_file(SINFO_FILE, with_states=('idle'))

    assert [
        node.hostname
        for cluster in resources.values()
        for partition in cluster
        for node in partition
    ] == ['nazare001', 'gnode2', 'visu1', 'gnode2', 'visu1']


def test_slurm_sinfo_resources(monkeypatch):
    """Test SLURM SINFO resources."""
    monkeypatch.setattr(subprocess, 'check_output', lambda _: SINFO_CONTENT.encode())

    clusters = sinfo(username='john-doe', with_states=('completing'))

    assert 'nautilus' in clusters
    assert list(clusters) == ['nautilus']

    nautilus = clusters['nautilus']

    assert isinstance(nautilus, SlurmCluster)
    assert 'standard' in nautilus
    assert list(nautilus) == ['standard', 'all']

    standard = nautilus['standard']

    assert isinstance(standard, SlurmPartition)
    assert 'cnode001' in standard
    assert list(standard) == ['cnode001']

    cnode = standard['cnode001']

    assert isinstance(cnode, SlurmNode)
    assert cnode == 'cnode001'


def test_slurm_gres():
    """Test SLURM GPU resources extraction."""
    resources = sinfo_from_file(SINFO_FILE, with_states=('idle', 'idle~', 'mixed', 'allocated'))

    gpus = gres(resources)

    # Sorted and without duplicates
    assert gpus == ['A100', 'A100_2g.10gb', 'A40', 'None', 'P100', 'T4']