############################################################################
# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved
############################################################################
import itertools

import pytest

import atf

topology_yaml = """
- topology: topo1
  cluster_default: true
  torus3d:
    toruses:
      - name: pod1
        dims:
          x: 4
          y: 4
          z: 2
        regions:
          - anchor: {x: 0, y: 0, z: 0}
            dims: {x: 4, y: 4, z: 1}
            nodes: node[1-16]
          - anchor: {x: 0, y: 0, z: 1}
            dims: {x: 4, y: 4, z: 1}
            nodes: node[17-32]
        placements:
          - dims:
              x: 2
              y: 2
              z: 2
          - dims:
              x: 2
              y: 2
              z: 1
          - dims:
              x: 4
              y: 4
              z: 2
"""

# All valid anchored 2x2x1 sub-cubes in the 4x4x2 torus (anchor_spacing=dims)
VALID_4NODE = [
    {"node1", "node2", "node5", "node6"},
    {"node3", "node4", "node7", "node8"},
    {"node9", "node10", "node13", "node14"},
    {"node11", "node12", "node15", "node16"},
    {"node17", "node18", "node21", "node22"},
    {"node19", "node20", "node23", "node24"},
    {"node25", "node26", "node29", "node30"},
    {"node27", "node28", "node31", "node32"},
]

# All valid anchored 2x2x2 sub-cubes
VALID_8NODE = [
    {"node1", "node2", "node5", "node6", "node17", "node18", "node21", "node22"},
    {"node3", "node4", "node7", "node8", "node19", "node20", "node23", "node24"},
    {"node9", "node10", "node13", "node14", "node25", "node26", "node29", "node30"},
    {"node11", "node12", "node15", "node16", "node27", "node28", "node31", "node32"},
]

# Valid 16-node allocations under --segment=8: union of any two distinct
# 8-node placements (C(4, 2) = 6 pairs)
VALID_16NODE_SEGMENT = [a | b for a, b in itertools.combinations(VALID_8NODE, 2)]


def assert_valid_placement(job_id, valid_sets):
    """Assert that the job's NodeList matches one of the valid placement sets."""
    node_list = atf.get_job_parameter(job_id, "NodeList")
    nodes = set(atf.node_range_to_list(node_list))
    assert any(
        nodes == v for v in valid_sets
    ), f"Job {job_id} should be allocated in a valid placement ({valid_sets}), but was allocated in {node_list}"


@pytest.fixture(scope="module", autouse=True)
def setup():
    atf.require_version(
        (26, 5),
        component="sbin/slurmd",
        reason="topology/torus3d was added in 26.05",
    )
    atf.require_config_parameter_includes("SchedulerParameters", ("bf_interval", 1))
    atf.require_nodes(32)
    atf.require_config_parameter("SelectType", "select/cons_tres")
    atf.require_config_parameter("SelectTypeParameters", "CR_CPU")
    atf.require_config_file("topology.yaml", topology_yaml)
    atf.require_slurm_running()


def test_basic_placement():
    """Test basic 4-node placement (2x2x1)"""

    job_id = atf.submit_job_sbatch('-N 4 --exclusive --mem=1 --wrap="hostname"')
    assert job_id != 0, "Job should be accepted"
    atf.wait_for_job_state(job_id, "DONE", fatal=True)
    assert_valid_placement(job_id, VALID_4NODE)


def test_8_node_placement():
    """Test 8-node placement (2x2x2)"""

    job_id = atf.submit_job_sbatch('-N 8 --exclusive --mem=1 --wrap="hostname"')
    assert job_id != 0, "Job should be accepted"
    atf.wait_for_job_state(job_id, "DONE", fatal=True)
    assert_valid_placement(job_id, VALID_8NODE)


def test_full_torus():
    """Test full torus allocation (4x4x2 = 32 nodes)"""

    job_id = atf.submit_job_sbatch('-N 32 --exclusive --mem=1 --wrap="hostname"')
    assert job_id != 0, "Job should be accepted"
    atf.wait_for_job_state(job_id, "DONE", fatal=True)


def test_unsupported_size():
    """Test that unsupported placement size is rejected"""

    assert (
        atf.submit_job_sbatch('-N 3 --exclusive --mem=1 --wrap="hostname"', xfail=True)
        == 0
    ), "3-node job should fail -- no 3-node placement exists"


def test_concurrent_placements():
    """Test two concurrent 4-node placements fill without conflict"""

    job_id_1 = atf.submit_job_sbatch('-N 4 --exclusive --mem=1 --wrap="sleep infinity"')
    job_id_2 = atf.submit_job_sbatch('-N 4 --exclusive --mem=1 --wrap="sleep infinity"')
    atf.wait_for_job_state(job_id_1, "RUNNING", fatal=True)
    atf.wait_for_job_state(job_id_2, "RUNNING", fatal=True)

    nodes_1 = set(atf.node_range_to_list(atf.get_job_parameter(job_id_1, "NodeList")))
    nodes_2 = set(atf.node_range_to_list(atf.get_job_parameter(job_id_2, "NodeList")))
    assert not nodes_1.intersection(nodes_2), "Concurrent jobs should not share nodes"


def test_frag_aware_placement():
    """Test that frag-aware placement fills a partially used 2x2x2 block.

    Pin job1 to z=0 anchor (0,0,0): node[3,4,7,8].
    Job2 (also 2x2x1) should land at z=1 same (x,y): node[19,20,23,24],
    completing the 2x2x2 block rather than consuming a fresh z=0 placement
    which would block another 2x2x2 anchor.
    """

    expected_nodes = {"node19", "node20", "node23", "node24"}

    job_id_1 = atf.submit_job_sbatch(
        '-N 4 -w node[3,4,7,8] --exclusive --mem=1 --wrap="sleep infinity"'
    )
    atf.wait_for_job_state(job_id_1, "RUNNING", fatal=True)

    job_id_2 = atf.submit_job_sbatch('-N 4 --exclusive --mem=1 --wrap="sleep infinity"')
    atf.wait_for_job_state(job_id_2, "RUNNING", fatal=True)

    nodes_2 = set(atf.node_range_to_list(atf.get_job_parameter(job_id_2, "NodeList")))
    assert nodes_2 == expected_nodes, (
        f"Frag-aware placement should fill z=1 below job1 (node[19,20,23,24]), "
        f"but job2 got {nodes_2}"
    )


def test_torus_constrained():
    """Test that allocating all z=0 nodes prevents 8-node (2x2x2) jobs.

    Torus 4x4x2 flat mapping (x varies fastest):
      z=0: node[1-16], z=1: node[17-32]
    Every 2x2x2 placement needs 4 nodes from z=0.
    With all z=0 nodes busy, no 2x2x2 placement can be satisfied,
    even though 16 z=1 nodes are idle.
    """

    z0_placements = [
        "node[1,2,5,6]",
        "node[3,4,7,8]",
        "node[9,10,13,14]",
        "node[11,12,15,16]",
    ]
    job_ids = []
    for nodes in z0_placements:
        job_id = atf.submit_job_sbatch(
            f'-N 4 -w {nodes} --exclusive --mem=1 --wrap="sleep infinity"'
        )
        job_ids.append(job_id)

    for job_id in job_ids:
        atf.wait_for_job_state(job_id, "RUNNING", fatal=True)

    job_id_big = atf.submit_job_sbatch(
        '-N 8 --exclusive --mem=1 --wrap="sleep infinity"'
    )
    assert job_id_big != 0, "Job should be accepted"
    atf.wait_for_job_state(
        job_id_big,
        "PENDING",
        desired_reason="No_suitable_topology_unit_found",
        fatal=True,
    )


def test_segment():
    """Test segment support -- 16 nodes as 2 segments of 8"""

    job_id = atf.submit_job_sbatch(
        '-N 16 --segment=8 --exclusive --mem=1 --wrap="hostname"'
    )
    assert job_id != 0, "Segmented job should be accepted"
    atf.wait_for_job_state(job_id, "DONE", fatal=True)
    assert_valid_placement(job_id, VALID_16NODE_SEGMENT)


def test_segment_bad_divisor():
    """Test that segment size not dividing job size is rejected"""

    assert (
        atf.submit_job_sbatch(
            '-N 8 --segment=3 --exclusive --mem=1 --wrap="hostname"',
            xfail=True,
        )
        == 0
    ), "Job should fail -- 8 % 3 != 0"