import sys
from unittest.mock import MagicMock, patch

import pytest

import ray
from ray._private.accelerators import TPUAcceleratorManager, tpu
from ray.util.tpu import SlicePlacementGroup


def test_get_current_pod_name_smoke():
    with patch(
        "ray._private.accelerators.tpu.TPUAcceleratorManager.get_current_node_tpu_name",
        return_value="my-tpu",
    ):
        name = ray.util.tpu.get_current_pod_name()
    assert name == "my-tpu"


def test_empty_get_current_pod_name_returns_none():
    with patch(
        "ray._private.accelerators.tpu.TPUAcceleratorManager.get_current_node_tpu_name",
        return_value="",
    ):
        name = ray.util.tpu.get_current_pod_name()
    assert name is None


@pytest.mark.parametrize(
    "test_case",
    [
        # (number_chips_per_host, parsed accl_type, expected_worker_count)
        (4, "v2-4", 1),
        (4, "v3-32", 4),
        (4, "v4-8", 1),
        (4, "v4-16", 2),
        (8, "v5litepod-4", 1),
        (8, "v5litepod-8", 1),
        (8, "v5litepod-16", 2),
        (8, "v5litepod-32", 4),
        (4, "v5p-4", 1),
        (4, "v5p-8", 1),
        (4, "v5p-16", 2),
        (4, "v6e-4", 1),
        (8, "v6e-8", 1),
        (4, "v6e-8", 2),
        (8, "v6e-16", 2),
        (4, "v7x-8", 1),
        (4, "v7x-16", 2),
    ],
)
@patch("glob.glob")
def test_worker_count(mock_glob, test_case):
    num_devices, accelerator_type, expected_worker_count = test_case
    mock_glob.return_value = ["/dev/accel" + str(x) for x in range(num_devices)]
    TPUAcceleratorManager.get_current_node_num_accelerators.cache_clear()

    with patch(
        "ray._private.accelerators.tpu.TPUAcceleratorManager."
        "get_current_node_tpu_pod_type",
        return_value=accelerator_type,
    ):
        worker_count = ray.util.tpu.get_current_pod_worker_count()

    assert worker_count == expected_worker_count


@patch("glob.glob")
def test_num_tpu_chips(mock_glob):
    mock_glob.return_value = [
        "/dev/accel0",
        "/dev/accel1",
        "/dev/accel2",
        "/dev/accel3",
    ]
    TPUAcceleratorManager.get_current_node_num_accelerators.cache_clear()
    num_tpu_chips = ray.util.tpu.get_num_tpu_chips_on_node()
    assert num_tpu_chips == 4


@pytest.mark.parametrize(
    "test_case",
    [
        # (accelerator_type, accelerator_topology, expected_result)
        ("v2-16", "4x4", True),
        ("v2-256", "16x16", True),
        ("v2-4", "2x2", False),
        ("v3-16", "4x4", True),
        ("v3-1024", "32x32", True),
        ("v3-4", "4x16", False),
        ("v4-4", "2x2x1", True),
        ("v4-32", "2x4x4", True),
        ("v4-2048", "8x8x16", True),
        ("v4-4", "16x16x16", False),
        ("v5p-128", "4x4x4", True),
        ("v5p-4096", "16x16x16", True),
        ("v5p-12288", "16x16x24", True),
        ("v5p-4", "24x24x24", False),
        ("v5litepod-16", "2x8", True),
        ("v5litepod-256", "16x16", True),
        ("v5litepod-4", "2x2", True),
        ("v6e-8", "2x4", True),
        ("v6e-16", "4x4", True),
        ("v6e-64", "8x8", True),
        ("v6e-4", "4x16", False),
        ("tpu7x-16", "2x2x2", True),
        ("tpu7x-64", "2x4x4", True),
        ("v7x-8", "4x4", False),
    ],
)
@patch("glob.glob")
def test_is_valid_tpu_accelerator_topology(_mock_glob, test_case):
    """Test valid TPU accelerator topologies."""
    accelerator_type, accelerator_topology, expected_result = test_case
    actual_result = TPUAcceleratorManager.is_valid_tpu_accelerator_topology(
        accelerator_type, accelerator_topology
    )

    assert actual_result == expected_result


def test_get_current_node_labels_env_only(monkeypatch):
    # Simulate GKE TPU environment variables
    monkeypatch.setenv("TPU_NAME", "tpu-worker-group-2")
    monkeypatch.setenv("TPU_WORKER_ID", "0")
    monkeypatch.setenv("TPU_ACCELERATOR_TYPE", "v6e-16")
    monkeypatch.setenv("TPU_TOPOLOGY", "4x4")

    tpu_labels = TPUAcceleratorManager.get_current_node_accelerator_labels()

    assert tpu_labels["ray.io/tpu-slice-name"] == "tpu-worker-group-2"
    assert tpu_labels["ray.io/tpu-worker-id"] == "0"
    assert tpu_labels["ray.io/tpu-topology"] == "4x4"
    assert tpu_labels["ray.io/tpu-pod-type"] == "v6e-16"


def test_get_current_node_tpu_topology_from_metadata():
    tpu_env_string = "TPU_ACCELERATOR:v6e.\nTOPOLOGY: '2x2x4'\nTPU_HOST_BOUNDS:0,1,1,2"

    with patch(
        "ray._private.accelerators.tpu._get_tpu_metadata", return_value=tpu_env_string
    ):
        topology = TPUAcceleratorManager.get_current_node_tpu_topology()
        assert topology == "2x2x4"


@pytest.mark.parametrize(
    "topology, accelerator_type, expected_pod_type, should_raise",
    [
        ("2x4", "TPU-V6E", "v6e-8", False),
        ("2x2x2", "TPU-V4", "v4-16", False),
        ("4x8", "TPU-V3", "v3-64", False),
        ("2x2x1", "TPU-V5P", "v5p-8", False),
        ("4x4", "TPU-V5P", "v5p-32", False),
        ("8x16", "TPU-V6E", "v6e-128", False),
        ("", "TPU-V3", None, False),
        ("4x", "TPU-V3", None, True),
        ("2x2x2", "TPU-V7X", "v7x-16", False),
    ],
)
def test_infer_tpu_pod_type_from_topology(
    topology, accelerator_type, expected_pod_type, should_raise
):
    if should_raise:
        with pytest.raises(ValueError):
            tpu.infer_tpu_pod_type_from_topology(topology, accelerator_type)
    else:
        actual_result = tpu.infer_tpu_pod_type_from_topology(topology, accelerator_type)
        assert actual_result == expected_pod_type


@pytest.fixture
def ray_start_cpu():
    address_info = ray.init(num_cpus=1)
    yield address_info
    ray.shutdown()


@pytest.fixture
def ray_tpu_cluster(ray_start_cluster):
    """
    Simulates a Ray cluster with two multi-host TPU v4-16 slices.
    """
    pod_type = "v4-16"
    topology = "2x2x2"

    cluster = ray_start_cluster
    slice_0_env_common = {
        "TPU_NAME": "test-slice-0",
        "TPU_ACCELERATOR_TYPE": pod_type,
        "TPU_TOPOLOGY": topology,
    }
    slice_0_head_labels = {
        "ray.io/tpu-slice-name": "test-slice-0",
        "ray.io/tpu-worker-id": "0",
        "ray.io/tpu-pod-type": pod_type,
        "ray.io/tpu-topology": topology,
    }
    slice_0_worker_labels = {
        "ray.io/tpu-slice-name": "test-slice-0",
        "ray.io/tpu-worker-id": "1",
        "ray.io/tpu-pod-type": pod_type,
        "ray.io/tpu-topology": topology,
    }
    cluster.add_node(
        num_cpus=2,
        resources={"TPU": 4, f"TPU-{pod_type}-head": 1},
        env_vars={**slice_0_env_common, "TPU_WORKER_ID": "0"},
        labels=slice_0_head_labels,
    )
    cluster.add_node(
        num_cpus=2,
        resources={"TPU": 4},
        env_vars={**slice_0_env_common, "TPU_WORKER_ID": "1"},
        labels=slice_0_worker_labels,
    )

    slice_1_env_common = {
        "TPU_NAME": "test-slice-1",
        "TPU_ACCELERATOR_TYPE": pod_type,
        "TPU_TOPOLOGY": topology,
    }
    slice_1_head_labels = {
        "ray.io/tpu-slice-name": "test-slice-1",
        "ray.io/tpu-worker-id": "0",
        "ray.io/tpu-pod-type": pod_type,
        "ray.io/tpu-topology": topology,
    }
    slice_1_worker_labels = {
        "ray.io/tpu-slice-name": "test-slice-1",
        "ray.io/tpu-worker-id": "1",
        "ray.io/tpu-pod-type": pod_type,
        "ray.io/tpu-topology": topology,
    }
    cluster.add_node(
        num_cpus=2,
        resources={"TPU": 4, f"TPU-{pod_type}-head": 1},
        env_vars={**slice_1_env_common, "TPU_WORKER_ID": "0"},
        labels=slice_1_head_labels,
    )
    cluster.add_node(
        num_cpus=2,
        resources={"TPU": 4},
        env_vars={**slice_1_env_common, "TPU_WORKER_ID": "1"},
        labels=slice_1_worker_labels,
    )

    ray.init(address=cluster.address)
    yield cluster
    ray.shutdown()


@pytest.fixture
def ray_v6e_tpu_cluster(ray_start_cluster):
    """
    Simulates a Ray cluster with two v6e-8 slices (2x4 topology).

    """
    pod_type = "v6e-8"
    topology = "2x4"
    cluster = ray_start_cluster

    for i in range(2):
        env_common = {
            "TPU_NAME": f"test-v6e-slice-{i}",
            "TPU_ACCELERATOR_TYPE": pod_type,
            "TPU_TOPOLOGY": topology,
        }
        head_labels = {
            "ray.io/tpu-slice-name": f"test-v6e-slice-{i}",
            "ray.io/tpu-worker-id": "0",
            "ray.io/tpu-pod-type": pod_type,
            "ray.io/tpu-topology": topology,
        }
        # A single-host v6e-8 has 8 chips on one node
        cluster.add_node(
            num_cpus=4,
            resources={"TPU": 8, f"TPU-{pod_type}-head": 1},
            env_vars={**env_common, "TPU_WORKER_ID": "0"},
            labels=head_labels,
        )

    ray.init(address=cluster.address)
    yield cluster
    ray.shutdown()


def test_fetch_tpu_slice_name_from_pg(ray_tpu_cluster):
    """Tests that the slice name can be fetched from a PG."""
    tpu_head_pg = ray.util.placement_group(bundles=[{"TPU-v4-16-head": 1}])
    ray.get(tpu_head_pg.ready())

    expected_unique_slice_names = {"test-slice-0", "test-slice-1"}
    slice_name = tpu.fetch_tpu_slice_name_from_pg(tpu_head_pg)
    assert slice_name in expected_unique_slice_names

    ray.util.remove_placement_group(tpu_head_pg)


def test_reserve_tpu_slice(ray_tpu_cluster):
    """Tests that a TPU slice can be successfully reserved."""
    reserved_name_0, hg_pg_0 = tpu.reserve_tpu_slice(
        topology="2x2x2", accelerator_type="TPU-V4"
    )
    reserved_name_1, hg_pg_1 = tpu.reserve_tpu_slice(
        topology="2x2x2", accelerator_type="TPU-V4"
    )

    # Ensure the placement groups reserving the TPU slice using the head worker are valid.
    assert hg_pg_0 is not None, "Expected placement group for slice 0, got None"
    assert hg_pg_1 is not None, "Expected placement group for slice 1, got None"

    assert (
        reserved_name_0 != reserved_name_1
    ), f"Expected to reserve two different slices, but got the same name: {reserved_name_0}"
    expected_unique_slice_names = {"test-slice-0", "test-slice-1"}
    actual_reserved_names = {reserved_name_0, reserved_name_1}
    assert actual_reserved_names == expected_unique_slice_names, (
        f"Got unexpected slice names. Expected {expected_unique_slice_names}, "
        f"but got {actual_reserved_names}"
    )


def test_slice_placement_group(ray_tpu_cluster):
    """Test that single TPU slice can be successfully reserved."""
    slice_placement_group = ray.util.tpu.slice_placement_group(
        topology="2x2x2",
        accelerator_version="v4",
    )
    assert slice_placement_group.chips_per_host == 4
    assert slice_placement_group.num_hosts == 2
    assert slice_placement_group.placement_group.bundle_count == 2
    assert slice_placement_group.placement_group.bundle_specs == [
        {"TPU": 4, "CPU": 1.0},
        {"TPU": 4, "CPU": 1.0},
    ]


def test_multi_slice_placement_group(ray_tpu_cluster):
    """Test that multiple whole TPU slices can be successfully reserved"""
    multi_slice_placement_group = ray.util.tpu.slice_placement_group(
        topology="2x2x2",
        accelerator_version="v4",
        num_slices=2,
    )
    assert multi_slice_placement_group.placement_group.bundle_count == 4
    assert multi_slice_placement_group.num_hosts == 4
    assert multi_slice_placement_group.placement_group.bundle_specs == [
        {"TPU": 4, "CPU": 1.0},  # slice 1, host 1
        {"TPU": 4, "CPU": 1.0},  # slice 1, host 2
        {"TPU": 4, "CPU": 1.0},  # slice 2, host 1
        {"TPU": 4, "CPU": 1.0},  # slice 2, host 2
    ]


@patch("ray.util.tpu.placement_group")
@patch("ray.util.tpu.remove_placement_group")
@patch("ray.util.tpu.reserve_tpu_slice")
def test_slice_placement_group_partial_failure_cleanup(
    mock_reserve, mock_remove_pg, mock_create_pg
):
    """
    Verifies that if a multi-slice request fails halfway through,
    the TPU head placement groups are cleaned up to prevent leaks.
    """
    fake_head_pg_1 = MagicMock(name="head_pg_1")
    mock_reserve.side_effect = [("slice_1", fake_head_pg_1), None]

    with pytest.raises(RuntimeError, match="Failed to reserve TPU slice"):
        SlicePlacementGroup(topology="2x2x2", accelerator_version="v4", num_slices=2)

    # Validate that 2 TPU util attempted to reserve two slices, failed, and
    # correctly cleaned up the hanging TPU head placement groups.
    assert mock_reserve.call_count == 2
    mock_remove_pg.assert_called_once_with(fake_head_pg_1)
    mock_create_pg.assert_not_called()


@pytest.mark.parametrize(
    "accelerator_type, expected_version",
    [
        # type with "TPU-" prefix
        ("TPU-V4", "v4"),
        ("TPU-v4", "v4"),
        ("TPU-V6E", "v6e"),
        ("TPU-v5p", "v5p"),
        ("TPU-V7X", "v7x"),
        # Only the TPU version - no parsing necessary.
        ("v4", "v4"),
        ("v3", "v3"),
        ("v6e", "v6e"),
        ("v5litepod", "v5litepod"),
        ("v7x", "v7x"),
    ],
)
def test_get_tpu_version_valid(accelerator_type, expected_version):
    assert ray.util.tpu.get_tpu_version_from_type(accelerator_type) == expected_version


@pytest.mark.parametrize(
    "invalid_type",
    [
        "A100",  # GPU type
        "random-invalid-type",  # Random string
        "TPU-invalid",  # TPU prefix
        "",  # Empty string
    ],
)
def test_get_tpu_version_invalid(invalid_type):
    with pytest.raises(ValueError, match="Invalid accelerator_type"):
        ray.util.tpu.get_tpu_version_from_type(invalid_type)


@pytest.mark.parametrize(
    "topology, accelerator_type, num_workers, resources_per_worker, expected_slices",
    [
        # "2x2x1" has 4 chips, for 4 workers with TPU: 1 each we expect num_slices=1.
        ("2x2x1", "TPU-V4", 4, {"TPU": 1}, 1),
        # "2x2x1" has 4 chips, for 8 workers with TPU: 1 each we expect num_slices=2.
        ("2x2x1", "v4", 8, {"TPU": 1}, 2),
        # "2x2x2" has 8 chips and 2 hosts, defaulting to 1 TPU worker per host
        # and requesting 4 workers, we expect num_slices=2.
        ("2x2x2", "TPU-V4", 4, None, 2),
        # "2x2x4" has 16 chips and 4 hosts, defaulting to 1 TPU worker per host
        # and requesting 4 workers, we expect num_slices=1.
        ("2x2x4", "TPU-V4", 4, None, 1),
        # 0 workers requested -> fallback to 1 slice.
        ("2x2x1", "v4", 0, None, 1),
        # Invalid topology -> fallback to 1 slice.
        ("", "v4", 4, {"TPU": 1}, 1),
        ("2x2x1", "", 4, {"TPU": 1}, 1),
    ],
)
def test_get_tpu_num_slices_for_workers(
    topology, accelerator_type, num_workers, resources_per_worker, expected_slices
):
    num_slices = ray.util.tpu.get_tpu_num_slices_for_workers(
        topology=topology,
        accelerator_type=accelerator_type,
        num_workers=num_workers,
        resources_per_worker=resources_per_worker,
    )
    assert num_slices == expected_slices


def _make_mock_tpu_node(
    alive, pod_type, slice_name, worker_id, tpu_chips=4, node_id=None
):
    """Helper to mock a Ray Node dictionary returned by ray.nodes()."""
    if node_id is None:
        node_id = f"node_{slice_name}_{worker_id}"
    return {
        "NodeID": node_id,
        "Alive": alive,
        "Labels": {
            "ray.io/tpu-pod-type": pod_type,
            "ray.io/tpu-slice-name": slice_name,
            "ray.io/tpu-worker-id": str(worker_id),
        },
        "Resources": {"TPU": tpu_chips},
    }


@pytest.mark.parametrize(
    "topology, accelerator_type, mock_nodes, mock_avail_resources, expected_ready",
    [
        # 1 fully intact and available v4 slice (2 physical hosts).
        (
            "2x2x2",
            "v4",
            [
                _make_mock_tpu_node(True, "v4-16", "slice-1", 0, node_id="A"),
                _make_mock_tpu_node(True, "v4-16", "slice-1", 1, node_id="B"),
            ],
            {
                "A": {"TPU": 4},
                "B": {"TPU": 4},
            },
            1,
        ),
        # 1 fully intact slice, but one node is using 2 TPUs (unavailable) -> 0 ready slices.
        (
            "2x2x2",
            "v4",
            [
                _make_mock_tpu_node(True, "v4-16", "slice-1", 0, node_id="A"),
                _make_mock_tpu_node(True, "v4-16", "slice-1", 1, node_id="B"),
            ],
            {
                "A": {"TPU": 2},
                "B": {"TPU": 4},
            },
            0,
        ),
        # Fractured slice (missing a physical host) -> 0 ready slices.
        (
            "2x2x2",
            "v4",
            [
                _make_mock_tpu_node(True, "v4-16", "slice-1", 0, node_id="A"),
                # Worker 1 is missing
            ],
            {
                "A": {"TPU": 4},
            },
            0,
        ),
        # Correct number of hosts, but missing the head node (rank 0) -> 0 ready slices.
        (
            "2x2x2",
            "v4",
            [
                _make_mock_tpu_node(True, "v4-16", "slice-1", 1, node_id="A"),
                _make_mock_tpu_node(True, "v4-16", "slice-1", 2, node_id="B"),
            ],
            {
                "A": {"TPU": 4},
                "B": {"TPU": 4},
            },
            0,
        ),
        # Fractured slice (one physical host is dead) -> 0 ready slices.
        (
            "2x2x2",
            "v4",
            [
                _make_mock_tpu_node(True, "v4-16", "slice-1", 0, node_id="A"),
                _make_mock_tpu_node(False, "v4-16", "slice-1", 1, node_id="B"),
            ],
            {
                "A": {"TPU": 4},
            },
            0,
        ),
        # 2 slices total: one intact & available, one fractured -> 1 ready slice.
        (
            "2x2x2",
            "v4",
            [
                _make_mock_tpu_node(True, "v4-16", "slice-A", 0, node_id="A0"),
                _make_mock_tpu_node(True, "v4-16", "slice-A", 1, node_id="A1"),
                _make_mock_tpu_node(True, "v4-16", "slice-B", 0, node_id="B0"),
                # slice-B worker 1 is missing
            ],
            {
                "A0": {"TPU": 4},
                "A1": {"TPU": 4},
                "B0": {"TPU": 4},
            },
            1,
        ),
        # 1 fully intact and available v6e 2x4 slice (single-host).
        (
            "2x4",
            "v6e",
            [
                _make_mock_tpu_node(
                    True, "v6e-8", "slice-1", 0, tpu_chips=8, node_id="A"
                ),
            ],
            {
                "A": {"TPU": 8},
            },
            1,
        ),
        # 1 fully intact and available v6e 2x4 slice (2 physical hosts).
        (
            "2x4",
            "v6e",
            [
                _make_mock_tpu_node(
                    True, "v6e-8", "slice-1", 0, tpu_chips=4, node_id="A"
                ),
                _make_mock_tpu_node(
                    True, "v6e-8", "slice-1", 1, tpu_chips=4, node_id="B"
                ),
            ],
            {
                "A": {"TPU": 4},
                "B": {"TPU": 4},
            },
            1,
        ),
        # 2 fully intact v6e slices.
        (
            "4x4",
            "v6e",
            [
                _make_mock_tpu_node(True, "v6e-16", "slice-1", 0, node_id="S1_0"),
                _make_mock_tpu_node(True, "v6e-16", "slice-1", 1, node_id="S1_1"),
                _make_mock_tpu_node(True, "v6e-16", "slice-1", 2, node_id="S1_2"),
                _make_mock_tpu_node(True, "v6e-16", "slice-1", 3, node_id="S1_3"),
                _make_mock_tpu_node(True, "v6e-16", "slice-2", 0, node_id="S2_0"),
                _make_mock_tpu_node(True, "v6e-16", "slice-2", 1, node_id="S2_1"),
                _make_mock_tpu_node(True, "v6e-16", "slice-2", 2, node_id="S2_2"),
                _make_mock_tpu_node(True, "v6e-16", "slice-2", 3, node_id="S2_3"),
            ],
            {
                "S1_0": {"TPU": 4},
                "S1_1": {"TPU": 4},
                "S1_2": {"TPU": 4},
                "S1_3": {"TPU": 4},
                "S2_0": {"TPU": 4},
                "S2_1": {"TPU": 4},
                "S2_2": {"TPU": 4},
                "S2_3": {"TPU": 4},
            },
            2,
        ),
    ],
)
@patch("ray.is_initialized", return_value=True)
@patch("ray._private.state.available_resources_per_node")
@patch("ray.nodes")
def test_get_num_ready_tpu_slices_calculation(
    mock_nodes_call,
    mock_avail_resources_call,
    mock_is_initialized,
    topology,
    accelerator_type,
    mock_nodes,
    mock_avail_resources,
    expected_ready,
):
    """Test that the TPU slice readiness utility correctly calculates the number of ready
    slices in different mocked scenarios, including idle resource verification."""
    mock_nodes_call.return_value = mock_nodes
    mock_avail_resources_call.return_value = mock_avail_resources

    actual_ready = ray.util.tpu.get_num_ready_tpu_slices(
        topology=topology,
        accelerator_type=accelerator_type,
    )
    assert actual_ready == expected_ready


@pytest.mark.parametrize(
    "topology, accelerator_type, mock_nodes, expected_intact",
    [
        # 1 fully intact v4 slice (2 physical hosts).
        (
            "2x2x2",
            "v4",
            [
                _make_mock_tpu_node(True, "v4-16", "slice-1", 0, node_id="A"),
                _make_mock_tpu_node(True, "v4-16", "slice-1", 1, node_id="B"),
            ],
            1,
        ),
        # Fractured slice (missing a physical host) -> 0 intact slices.
        (
            "2x2x2",
            "v4",
            [
                _make_mock_tpu_node(True, "v4-16", "slice-1", 0, node_id="A"),
            ],
            0,
        ),
        # Missing head node (rank 0) -> 0 intact slices.
        (
            "2x2x2",
            "v4",
            [
                _make_mock_tpu_node(True, "v4-16", "slice-1", 1, node_id="A"),
                _make_mock_tpu_node(True, "v4-16", "slice-1", 2, node_id="B"),
            ],
            0,
        ),
        # One physical host is dead -> 0 intact slices.
        (
            "2x2x2",
            "v4",
            [
                _make_mock_tpu_node(True, "v4-16", "slice-1", 0, node_id="A"),
                _make_mock_tpu_node(False, "v4-16", "slice-1", 1, node_id="B"),
            ],
            0,
        ),
        # 2 slices: one intact, one fractured -> 1 intact slice.
        (
            "2x2x2",
            "v4",
            [
                _make_mock_tpu_node(True, "v4-16", "slice-A", 0, node_id="A0"),
                _make_mock_tpu_node(True, "v4-16", "slice-A", 1, node_id="A1"),
                _make_mock_tpu_node(True, "v4-16", "slice-B", 0, node_id="B0"),
            ],
            1,
        ),
        # 2 fully intact v6e slices.
        (
            "4x4",
            "v6e",
            [
                _make_mock_tpu_node(True, "v6e-16", "slice-1", 0, node_id="S1_0"),
                _make_mock_tpu_node(True, "v6e-16", "slice-1", 1, node_id="S1_1"),
                _make_mock_tpu_node(True, "v6e-16", "slice-1", 2, node_id="S1_2"),
                _make_mock_tpu_node(True, "v6e-16", "slice-1", 3, node_id="S1_3"),
                _make_mock_tpu_node(True, "v6e-16", "slice-2", 0, node_id="S2_0"),
                _make_mock_tpu_node(True, "v6e-16", "slice-2", 1, node_id="S2_1"),
                _make_mock_tpu_node(True, "v6e-16", "slice-2", 2, node_id="S2_2"),
                _make_mock_tpu_node(True, "v6e-16", "slice-2", 3, node_id="S2_3"),
            ],
            2,
        ),
    ],
)
@patch("ray.is_initialized", return_value=True)
@patch("ray.nodes")
def test_get_num_tpu_slices_calculation(
    mock_nodes_call,
    mock_is_initialized,
    topology,
    accelerator_type,
    mock_nodes,
    expected_intact,
):
    """Test that the intact TPU slice utility counts slices based purely on
    physical integrity (all hosts alive, correct chip count) without checking
    whether they are idle."""
    mock_nodes_call.return_value = mock_nodes

    actual_intact = ray.util.tpu.get_num_tpu_slices(
        topology=topology,
        accelerator_type=accelerator_type,
    )
    assert actual_intact == expected_intact


@patch("ray.is_initialized", return_value=False)
def test_get_num_tpu_slices_uninitialized(mock_is_initialized):
    """Test that the utility gracefully handles an uninitialized Ray context."""
    assert ray.util.tpu.get_num_tpu_slices("2x2x2", "v4") == 0


def test_get_num_ready_tpu_slices(ray_tpu_cluster):
    """
    Tests the get_num_ready_tpu_slices utility against a real Ray cluster.
    The ray_tpu_cluster fixture provisions two v4-16 slices (2x2x2 topology).
    """
    ready_slices = ray.util.tpu.get_num_ready_tpu_slices(
        topology="2x2x2", accelerator_type="v4"
    )
    assert ready_slices == 2


@patch("ray.is_initialized", return_value=False)
def test_get_num_ready_tpu_slices_uninitialized(mock_is_initialized):
    """Test that the utility gracefully handles an uninitialized Ray context."""
    assert ray.util.tpu.get_num_ready_tpu_slices("2x2x2", "v4") == 0


@pytest.mark.parametrize(
    "node_dict, expected_slice_name",
    [
        (_make_mock_tpu_node(True, "v4-16", "slice-1", 0), "slice-1"),
        (_make_mock_tpu_node(True, "v6e-8", "slice-A", 1), "slice-A"),
        (_make_mock_tpu_node(True, "v4-16", "", 0), ""),
        ({"Alive": True, "Labels": {}}, None),  # Missing TPU slice name
        ({"Alive": True}, None),  # Missing Node labels dict
    ],
)
def test_get_tpu_slice_name_from_node(node_dict, expected_slice_name):
    """Tests that the utility correctly extracts the TPU slice name from a node dictionary."""
    assert ray.util.tpu.get_tpu_slice_name_from_node(node_dict) == expected_slice_name


@patch("ray.is_initialized", return_value=True)
@patch("ray.nodes")
def test_get_tpu_nodes_for_slice(mock_nodes_call, mock_is_initialized):
    """Tests that the utility correctly filters alive nodes for a specific slice."""
    mock_nodes = [
        _make_mock_tpu_node(True, "v4-16", "slice-A", 0),
        _make_mock_tpu_node(True, "v4-16", "slice-A", 1),
        _make_mock_tpu_node(False, "v4-16", "slice-A", 2),  # Dead node
        _make_mock_tpu_node(True, "v4-16", "slice-B", 0),  # Wrong slice
    ]
    mock_nodes_call.return_value = mock_nodes

    # Call ray.nodes() to fetch from GCS
    nodes_a = ray.util.tpu.get_tpu_nodes_for_slice("slice-A")
    assert len(nodes_a) == 2
    assert nodes_a[0]["Labels"][ray._raylet.RAY_NODE_TPU_WORKER_ID_KEY] == "0"
    assert nodes_a[1]["Labels"][ray._raylet.RAY_NODE_TPU_WORKER_ID_KEY] == "1"
    assert mock_nodes_call.call_count == 1

    # Pass cached nodes directly
    nodes_b = ray.util.tpu.get_tpu_nodes_for_slice("slice-B", nodes=mock_nodes)
    assert len(nodes_b) == 1
    assert nodes_b[0]["Labels"][ray._raylet.RAY_NODE_TPU_SLICE_NAME_KEY] == "slice-B"
    assert mock_nodes_call.call_count == 1  # Call count remains 1

    # Use non-existent slice
    nodes_c = ray.util.tpu.get_tpu_nodes_for_slice("slice-C", nodes=mock_nodes)
    assert len(nodes_c) == 0


@patch("ray.is_initialized", return_value=False)
def test_get_tpu_nodes_for_slice_uninitialized(mock_is_initialized):
    """Test that the utility gracefully handles an uninitialized Ray context."""
    assert ray.util.tpu.get_tpu_nodes_for_slice("slice-A") == []


def test_get_tpu_worker_resources_chips_per_vm_override():
    """Test that chips_per_vm correctly overrides the default resource calculations."""

    # Default behavior: v6e 2x4 defaults to a single 8-chip host
    num_workers, resources = ray.util.tpu.get_tpu_worker_resources(
        topology="2x4", accelerator_type="v6e"
    )
    assert num_workers == 1
    assert resources["TPU"] == 8

    # Override behavior: v6e 2x4 forced to 4 chips per VM (2 hosts)
    num_workers_override, resources_override = ray.util.tpu.get_tpu_worker_resources(
        topology="2x4", accelerator_type="v6e", chips_per_vm=4
    )
    assert num_workers_override == 2
    assert resources_override["TPU"] == 4


def test_slice_placement_group_chips_per_vm_override(ray_v6e_tpu_cluster):
    """Test that SlicePlacementGroup respects chips_per_vm for host calculation."""

    # Default behavior (1 VM with 8 chips)
    default_pg = SlicePlacementGroup(topology="2x4", accelerator_version="v6e")
    assert default_pg.chips_per_host == 8
    assert default_pg.num_hosts == 1
    assert default_pg.num_bundles == 1
    assert default_pg.bundle_resources["TPU"] == 8

    # User-specified override behavior (2 VMs with 4 chips each)
    override_pg = SlicePlacementGroup(
        topology="2x4", accelerator_version="v6e", chips_per_vm=4
    )
    assert override_pg.chips_per_host == 4
    assert override_pg.num_hosts == 2
    assert override_pg.num_bundles == 2
    assert override_pg.bundle_resources["TPU"] == 4


def test_user_bundle_label_selector_merged(ray_tpu_cluster):
    """Verifies that user-passed bundle_label_selector is merged with dynamic TPU labels."""
    user_selectors = [{"env": "prod"}, {"env": "test"}]

    # 2x2x2 v4 = 2 hosts = 2 bundles
    slice_pg = SlicePlacementGroup(
        topology="2x2x2", accelerator_version="v4", bundle_label_selector=user_selectors
    )

    assert len(slice_pg._bundle_label_selector) == 2

    # Verify slice 0
    assert slice_pg._bundle_label_selector[0]["env"] == "prod"
    assert ray._raylet.RAY_NODE_TPU_SLICE_NAME_KEY in slice_pg._bundle_label_selector[0]

    # Verify slice 1
    assert slice_pg._bundle_label_selector[1]["env"] == "test"
    assert ray._raylet.RAY_NODE_TPU_SLICE_NAME_KEY in slice_pg._bundle_label_selector[1]


def test_user_bundle_label_selector_collision_dynamic_wins(ray_v6e_tpu_cluster):
    """Verifies that dynamic TPU labels take precedence on collision."""
    user_selectors = [{ray._raylet.RAY_NODE_TPU_SLICE_NAME_KEY: "user-requested-slice"}]

    # v6e-8 is single host (1 bundle)
    slice_pg = SlicePlacementGroup(
        topology="2x4", accelerator_version="v6e", bundle_label_selector=user_selectors
    )

    assert len(slice_pg._bundle_label_selector) == 1
    # The dynamic value should win (it generates test-v6e-slice-N)
    actual_val = slice_pg._bundle_label_selector[0][
        ray._raylet.RAY_NODE_TPU_SLICE_NAME_KEY
    ]
    assert actual_val != "user-requested-slice"
    assert "test-v6e-slice-" in actual_val


def test_user_bundle_label_selector_length_mismatch_raises():
    """Verifies that providing wrong length of selector list raises ValueError."""
    user_selectors = [{"env": "prod"}]  # Only 1 provided but 2x2x2 v4 has 2 hosts

    with pytest.raises(ValueError, match="bundle_label_selector length"):
        SlicePlacementGroup(
            topology="2x2x2",
            accelerator_version="v4",
            bundle_label_selector=user_selectors,
        )


def test_release_head_pgs_idempotent(ray_tpu_cluster):
    """Verifies that release_head_pgs() is idempotent."""
    slice_pg = SlicePlacementGroup(topology="2x2x2", accelerator_version="v4")

    assert len(slice_pg.head_placement_groups) == 1

    slice_pg.release_head_pgs()
    assert len(slice_pg.head_placement_groups) == 0

    # Call again, should not raise
    slice_pg.release_head_pgs()
    assert len(slice_pg.head_placement_groups) == 0


def test_shutdown_idempotent(ray_tpu_cluster):
    """Verifies that shutdown() is idempotent."""
    slice_pg = SlicePlacementGroup(topology="2x2x2", accelerator_version="v4")

    slice_pg.shutdown()
    assert slice_pg.placement_group is None
    assert len(slice_pg.head_placement_groups) == 0

    # Call again, should not raise
    slice_pg.shutdown()


def test_shutdown_safe_after_construction_failure():
    """Verifies that shutdown() is safe to call on a partially-constructed instance."""
    with patch(
        "ray.util.tpu.SlicePlacementGroup._reserve_slice",
        side_effect=RuntimeError("Test failure"),
    ):
        with pytest.raises(RuntimeError, match="Test failure"):
            SlicePlacementGroup(topology="2x2x2", accelerator_version="v4")

    # If the above didn't crash or leak resources, we are good.
    # We can also manually construct a partial instance and call shutdown.
    partial_pg = SlicePlacementGroup.__new__(SlicePlacementGroup)
    partial_pg._head_pgs = []
    partial_pg._placement_group = None

    # Should not raise even though it's missing attributes
    partial_pg.shutdown()


def test_release_head_pgs_after_ready_then_shutdown(ray_tpu_cluster):
    """Validates Slice PG lifecycle: wait until ready, release head PGs, then shutdown."""
    slice_pg = SlicePlacementGroup(topology="2x2x2", accelerator_version="v4")

    # Wait for ready
    ray.get(slice_pg.placement_group.ready())

    slice_pg.release_head_pgs()
    assert len(slice_pg.head_placement_groups) == 0

    slice_pg.shutdown()
    assert slice_pg.placement_group is None


def test_chips_per_vm_zero_raises_value_error():
    """Verifies that passing chips_per_vm=0 explicitly raises a ValueError instead of silently using the topology default."""
    with pytest.raises(ValueError):
        SlicePlacementGroup(
            topology="2x2x2",
            accelerator_version="v4",
            chips_per_vm=0,
        )

    # Also verify when custom resources already include a "TPU" key
    with pytest.raises(ValueError):
        SlicePlacementGroup(
            topology="2x2x2",
            accelerator_version="v4",
            resources_per_bundle={"TPU": 4},
            chips_per_vm=0,
        )


if __name__ == "__main__":
    sys.exit(pytest.main(["-sv", __file__]))