from collections.abc import Mapping

import pytest

from exo.master.placement import (
    get_transition_events,
    place_instance,
)
from exo.master.tests.conftest import (
    create_node_memory,
    create_node_network,
    create_rdma_connection,
    create_socket_connection,
)
from exo.shared.models.model_cards import ModelCard, ModelId, ModelTask
from exo.shared.topology import Topology
from exo.shared.types.backends import Backend
from exo.shared.types.commands import PlaceInstance
from exo.shared.types.common import CommandId, NodeId
from exo.shared.types.events import (
    InstanceCreated,
    InstanceDeleted,
    TaskStatusUpdated,
)
from exo.shared.types.memory import Memory
from exo.shared.types.multiaddr import Multiaddr
from exo.shared.types.profiling import (
    NetworkInterfaceInfo,
    NodeNetworkInfo,
    NodeRdmaCtlStatus,
)
from exo.shared.types.tasks import TaskId, TaskStatus, TextGeneration
from exo.shared.types.text_generation import (
    InputMessage,
    InputMessageContent,
    TextGenerationTaskParams,
)
from exo.shared.types.topology import Connection, SocketConnection
from exo.shared.types.worker.downloads import (
    DownloadCompleted,
    DownloadFailed,
    DownloadOngoing,
    DownloadProgressData,
)
from exo.shared.types.worker.instances import (
    Instance,
    InstanceId,
    InstanceMeta,
    MlxJacclInstance,
    MlxRingInstance,
)
from exo.shared.types.worker.runners import RunnerId, ShardAssignments, ShardWithId
from exo.shared.types.worker.shards import PipelineShardMetadata, Sharding


class MockShard:
    def is_primary_output(self) -> bool:
        return True


@pytest.fixture
def instance(model_card: ModelCard) -> Instance:
    return MlxRingInstance(
        instance_id=InstanceId(),
        shard_assignments=ShardAssignments(
            model_id=ModelId("test-model"),
            shards=[
                ShardWithId(
                    NodeId(),
                    RunnerId(),
                    PipelineShardMetadata(
                        model_card=model_card,
                        device_rank=0,
                        world_size=1,
                        start_layer=0,
                        end_layer=model_card.n_layers,
                        n_layers=model_card.n_layers,
                    ),
                )
            ],
            primary_output_node=0,
        ),
        hosts_by_node={},
        ephemeral_port=50000,
    )


@pytest.fixture
def model_card() -> ModelCard:
    return ModelCard(
        model_id=ModelId("test-model"),
        storage_size=Memory.from_kb(1000),
        n_layers=10,
        hidden_size=30,
        supports_tensor=True,
        tasks=[ModelTask.TextGeneration],
        backends=[Backend.MlxMetal],
    )


def _metal_only(
    node_memory: Mapping[NodeId, object],
) -> dict[NodeId, list[Backend]]:
    return {node_id: [Backend.MlxMetal] for node_id in node_memory}


def place_instance_command(model_card: ModelCard) -> PlaceInstance:
    return PlaceInstance(
        command_id=CommandId(),
        model_card=model_card,
        sharding=Sharding.Pipeline,
        instance_meta=InstanceMeta.MlxRing,
        min_nodes=1,
    )


@pytest.mark.parametrize(
    "available_memory,total_layers,expected_layers",
    [
        ((500, 500, 1000), 12, (3, 3, 6)),
        ((500, 500, 500), 12, (4, 4, 4)),
        ((312, 468, 1092), 12, (2, 3, 7)),
    ],
)
def test_get_instance_placements_create_instance(
    available_memory: tuple[int, int, int],
    total_layers: int,
    expected_layers: tuple[int, int, int],
    model_card: ModelCard,
):
    # arrange
    model_card = model_card.model_copy(
        update={
            "n_layers": total_layers,
            "storage_size": Memory.from_bytes(
                sum(available_memory)
            ),  # make it exactly fit across all nodes
        }
    )
    topology = Topology()

    cic = place_instance_command(model_card)
    node_id_a = NodeId()
    node_id_b = NodeId()
    node_id_c = NodeId()
    node_to_layers = {
        node_id_a: expected_layers[0],
        node_id_b: expected_layers[1],
        node_id_c: expected_layers[2],
    }

    # fully connected (directed) between the 3 nodes
    conn_a_b = Connection(
        source=node_id_a, sink=node_id_b, edge=create_socket_connection(1)
    )
    conn_b_c = Connection(
        source=node_id_b, sink=node_id_c, edge=create_socket_connection(2)
    )
    conn_c_a = Connection(
        source=node_id_c, sink=node_id_a, edge=create_socket_connection(3)
    )
    conn_c_b = Connection(
        source=node_id_c, sink=node_id_b, edge=create_socket_connection(4)
    )
    conn_a_c = Connection(
        source=node_id_a, sink=node_id_c, edge=create_socket_connection(5)
    )
    conn_b_a = Connection(
        source=node_id_b, sink=node_id_a, edge=create_socket_connection(6)
    )

    node_memory = {
        node_id_a: create_node_memory(available_memory[0]),
        node_id_b: create_node_memory(available_memory[1]),
        node_id_c: create_node_memory(available_memory[2]),
    }
    node_network = {
        node_id_a: create_node_network(),
        node_id_b: create_node_network(),
        node_id_c: create_node_network(),
    }
    topology.add_node(node_id_a)
    topology.add_node(node_id_b)
    topology.add_node(node_id_c)
    topology.add_connection(conn_a_b)
    topology.add_connection(conn_b_c)
    topology.add_connection(conn_c_a)
    topology.add_connection(conn_c_b)
    topology.add_connection(conn_a_c)
    topology.add_connection(conn_b_a)

    # act
    placements = place_instance(
        cic, topology, {}, node_memory, node_network, _metal_only(node_memory)
    )

    # assert
    assert len(placements) == 1
    instance_id = list(placements.keys())[0]
    instance = placements[instance_id]
    assert instance.shard_assignments.model_id == model_card.model_id

    for nid, _, shard in (shards := instance.shard_assignments.shards):
        assert shard.end_layer - shard.start_layer == node_to_layers[nid]

    assert shards[0].shard.start_layer == 0
    assert shards[-1].shard.end_layer == total_layers


def test_get_instance_placements_one_node_exact_fit() -> None:
    topology = Topology()
    node_id = NodeId()
    topology.add_node(node_id)
    node_memory = {node_id: create_node_memory(1000 * 1024)}
    node_network = {node_id: create_node_network()}
    cic = place_instance_command(
        ModelCard(
            model_id=ModelId("test-model"),
            storage_size=Memory.from_kb(1000),
            n_layers=10,
            hidden_size=1000,
            supports_tensor=True,
            tasks=[ModelTask.TextGeneration],
            backends=[Backend.MlxMetal],
        ),
    )
    placements = place_instance(
        cic, topology, {}, node_memory, node_network, _metal_only(node_memory)
    )

    assert len(placements) == 1
    instance_id = list(placements.keys())[0]
    instance = placements[instance_id]
    assert instance.shard_assignments.model_id == "test-model"
    assert len(instance.shard_assignments.shards) == 1


def test_get_instance_placements_one_node_fits_with_extra_memory() -> None:
    topology = Topology()
    node_id = NodeId()
    topology.add_node(node_id)
    node_memory = {node_id: create_node_memory(1001 * 1024)}
    node_network = {node_id: create_node_network()}
    cic = place_instance_command(
        ModelCard(
            model_id=ModelId("test-model"),
            storage_size=Memory.from_kb(1000),
            n_layers=10,
            hidden_size=1000,
            supports_tensor=True,
            tasks=[ModelTask.TextGeneration],
            backends=[Backend.MlxMetal],
        ),
    )
    placements = place_instance(
        cic, topology, {}, node_memory, node_network, _metal_only(node_memory)
    )

    assert len(placements) == 1
    instance_id = list(placements.keys())[0]
    instance = placements[instance_id]
    assert instance.shard_assignments.model_id == "test-model"
    assert len(instance.shard_assignments.shards) == 1


def test_get_instance_placements_one_node_not_fit() -> None:
    topology = Topology()
    node_id = NodeId()
    topology.add_node(node_id)
    node_memory = {node_id: create_node_memory(1000 * 1024)}
    node_network = {node_id: create_node_network()}
    cic = place_instance_command(
        model_card=ModelCard(
            model_id=ModelId("test-model"),
            storage_size=Memory.from_kb(1001),
            n_layers=10,
            hidden_size=1000,
            supports_tensor=True,
            tasks=[ModelTask.TextGeneration],
            backends=[Backend.MlxMetal],
        ),
    )

    with pytest.raises(ValueError, match="No cycles found with sufficient memory"):
        place_instance(
            cic, topology, {}, node_memory, node_network, _metal_only(node_memory)
        )


def test_get_transition_events_no_change(instance: Instance):
    # arrange
    instance_id = InstanceId()
    current_instances = {instance_id: instance}
    target_instances = {instance_id: instance}

    # act
    events = get_transition_events(current_instances, target_instances, {})

    # assert
    assert len(events) == 0


def test_get_transition_events_create_instance(instance: Instance):
    # arrange
    instance_id = InstanceId()
    current_instances: dict[InstanceId, Instance] = {}
    target_instances: dict[InstanceId, Instance] = {instance_id: instance}

    # act
    events = get_transition_events(current_instances, target_instances, {})

    # assert
    assert len(events) == 1
    assert isinstance(events[0], InstanceCreated)


def test_get_transition_events_delete_instance(instance: Instance):
    # arrange
    instance_id = InstanceId()
    current_instances: dict[InstanceId, Instance] = {instance_id: instance}
    target_instances: dict[InstanceId, Instance] = {}

    # act
    events = get_transition_events(current_instances, target_instances, {})

    # assert
    assert len(events) == 1
    assert isinstance(events[0], InstanceDeleted)
    assert events[0].instance_id == instance_id


def test_placement_selects_leaf_nodes(
    model_card: ModelCard,
):
    # arrange
    topology = Topology()

    model_card = model_card.model_copy(update={"storage_size": Memory.from_bytes(1000)})

    node_id_a = NodeId()
    node_id_b = NodeId()
    node_id_c = NodeId()
    node_id_d = NodeId()

    node_memory = {
        node_id_a: create_node_memory(500),
        node_id_b: create_node_memory(600),
        node_id_c: create_node_memory(600),
        node_id_d: create_node_memory(500),
    }
    node_network = {
        node_id_a: create_node_network(),
        node_id_b: create_node_network(),
        node_id_c: create_node_network(),
        node_id_d: create_node_network(),
    }

    topology.add_node(node_id_a)
    topology.add_node(node_id_b)
    topology.add_node(node_id_c)
    topology.add_node(node_id_d)

    # Daisy chain topology (directed)
    topology.add_connection(
        Connection(source=node_id_a, sink=node_id_b, edge=create_socket_connection(1))
    )
    topology.add_connection(
        Connection(source=node_id_b, sink=node_id_a, edge=create_socket_connection(1))
    )
    topology.add_connection(
        Connection(source=node_id_b, sink=node_id_c, edge=create_socket_connection(1))
    )
    topology.add_connection(
        Connection(source=node_id_c, sink=node_id_b, edge=create_socket_connection(1))
    )
    topology.add_connection(
        Connection(source=node_id_c, sink=node_id_d, edge=create_socket_connection(1))
    )
    topology.add_connection(
        Connection(source=node_id_d, sink=node_id_c, edge=create_socket_connection(1))
    )

    cic = place_instance_command(model_card=model_card)

    # act
    placements = place_instance(
        cic, topology, {}, node_memory, node_network, _metal_only(node_memory)
    )

    # assert
    assert len(placements) == 1
    instance = list(placements.values())[0]

    assigned_nodes = set(map(lambda it: it.node_id, instance.shard_assignments.shards))
    assert assigned_nodes == set((node_id_a, node_id_b)) or assigned_nodes == set(
        (
            node_id_c,
            node_id_d,
        )
    )


def test_tensor_rdma_backend_connectivity_matrix(
    model_card: ModelCard,
):
    # arrange
    topology = Topology()
    model_card = model_card.model_copy(
        update={
            "n_layers": 12,
            "storage_size": Memory.from_bytes(1500),
        }
    )

    node_a = NodeId()
    node_b = NodeId()
    node_c = NodeId()

    node_memory = {
        node_a: create_node_memory(500),
        node_b: create_node_memory(500),
        node_c: create_node_memory(500),
    }

    ethernet_interface = NetworkInterfaceInfo(
        name="en0",
        ip_address="10.0.0.1",
    )
    ethernet_conn = SocketConnection(
        sink_multiaddr=Multiaddr(address="/ip4/10.0.0.1/tcp/8000")
    )

    node_network = {
        node_a: NodeNetworkInfo(interfaces=[ethernet_interface]),
        node_b: NodeNetworkInfo(interfaces=[ethernet_interface]),
        node_c: NodeNetworkInfo(interfaces=[ethernet_interface]),
    }

    topology.add_node(node_a)
    topology.add_node(node_b)
    topology.add_node(node_c)

    # RDMA connections (directed)
    topology.add_connection(
        Connection(source=node_a, sink=node_b, edge=create_rdma_connection(3))
    )
    topology.add_connection(
        Connection(source=node_b, sink=node_a, edge=create_rdma_connection(3))
    )
    topology.add_connection(
        Connection(source=node_b, sink=node_c, edge=create_rdma_connection(4))
    )
    topology.add_connection(
        Connection(source=node_c, sink=node_b, edge=create_rdma_connection(4))
    )
    topology.add_connection(
        Connection(source=node_a, sink=node_c, edge=create_rdma_connection(5))
    )
    topology.add_connection(
        Connection(source=node_c, sink=node_a, edge=create_rdma_connection(5))
    )

    # Ethernet connections (directed)
    topology.add_connection(Connection(source=node_a, sink=node_b, edge=ethernet_conn))
    topology.add_connection(Connection(source=node_b, sink=node_c, edge=ethernet_conn))
    topology.add_connection(Connection(source=node_c, sink=node_a, edge=ethernet_conn))
    topology.add_connection(Connection(source=node_a, sink=node_c, edge=ethernet_conn))
    topology.add_connection(Connection(source=node_b, sink=node_a, edge=ethernet_conn))
    topology.add_connection(Connection(source=node_c, sink=node_b, edge=ethernet_conn))

    cic = PlaceInstance(
        sharding=Sharding.Tensor,
        instance_meta=InstanceMeta.MlxJaccl,
        command_id=CommandId(),
        model_card=model_card,
        min_nodes=1,
    )

    node_rdma_ctl = {
        node_a: NodeRdmaCtlStatus(enabled=True),
        node_b: NodeRdmaCtlStatus(enabled=True),
        node_c: NodeRdmaCtlStatus(enabled=True),
    }

    # act
    placements = place_instance(
        cic,
        topology,
        {},
        node_memory,
        node_network,
        _metal_only(node_memory),
        node_rdma_ctl=node_rdma_ctl,
    )

    # assert
    assert len(placements) == 1
    instance_id = list(placements.keys())[0]
    instance = placements[instance_id]

    assert isinstance(instance, MlxJacclInstance)

    assert instance.jaccl_devices is not None
    assert instance.jaccl_coordinators is not None

    matrix = instance.jaccl_devices
    assert len(matrix) == 3
    for i in range(3):
        assert matrix[i][i] is None

    assigned_nodes = list(instance.shard_assignments.shards)
    node_to_idx = {node_id: idx for idx, (node_id, _, _) in enumerate(assigned_nodes)}

    idx_a = node_to_idx[node_a]
    idx_b = node_to_idx[node_b]
    idx_c = node_to_idx[node_c]

    assert matrix[idx_a][idx_b] == "rdma_en3"
    assert matrix[idx_b][idx_c] == "rdma_en4"
    assert matrix[idx_c][idx_a] == "rdma_en5"

    # Verify coordinators are set for all nodes
    assert len(instance.jaccl_coordinators) == 3
    for node_id, _, _ in assigned_nodes:
        assert node_id in instance.jaccl_coordinators
        coordinator = instance.jaccl_coordinators[node_id]
        assert ":" in coordinator
        # Rank 0 node should use 0.0.0.0, others should use connection-specific IPs
        if node_id == assigned_nodes[0]:
            assert coordinator.startswith("0.0.0.0:")
        else:
            ip_part = coordinator.split(":")[0]
            assert len(ip_part.split(".")) == 4


def _build_three_node_rdma_topology() -> tuple[
    Topology, NodeId, NodeId, NodeId, dict[NodeId, NodeNetworkInfo]
]:
    topology = Topology()
    node_a = NodeId()
    node_b = NodeId()
    node_c = NodeId()

    ethernet_interface = NetworkInterfaceInfo(name="en0", ip_address="10.0.0.1")
    ethernet_conn = SocketConnection(
        sink_multiaddr=Multiaddr(address="/ip4/10.0.0.1/tcp/8000")
    )
    node_network = {
        node_a: NodeNetworkInfo(interfaces=[ethernet_interface]),
        node_b: NodeNetworkInfo(interfaces=[ethernet_interface]),
        node_c: NodeNetworkInfo(interfaces=[ethernet_interface]),
    }

    for n in (node_a, node_b, node_c):
        topology.add_node(n)

    rdma_pairs = [
        (node_a, node_b, 3),
        (node_b, node_a, 3),
        (node_b, node_c, 4),
        (node_c, node_b, 4),
        (node_a, node_c, 5),
        (node_c, node_a, 5),
    ]
    for src, sink, iface in rdma_pairs:
        topology.add_connection(
            Connection(source=src, sink=sink, edge=create_rdma_connection(iface))
        )

    socket_pairs = [
        (node_a, node_b),
        (node_b, node_c),
        (node_c, node_a),
        (node_a, node_c),
        (node_b, node_a),
        (node_c, node_b),
    ]
    for src, sink in socket_pairs:
        topology.add_connection(Connection(source=src, sink=sink, edge=ethernet_conn))

    return topology, node_a, node_b, node_c, node_network


def test_place_mlx_jaccl_rejects_when_a_node_has_rdma_ctl_disabled(
    model_card: ModelCard,
):
    # arrange
    model_card = model_card.model_copy(
        update={"n_layers": 12, "storage_size": Memory.from_bytes(1500)}
    )
    topology, node_a, node_b, node_c, node_network = _build_three_node_rdma_topology()
    node_memory = {
        node_a: create_node_memory(500),
        node_b: create_node_memory(500),
        node_c: create_node_memory(500),
    }
    node_rdma_ctl = {
        node_a: NodeRdmaCtlStatus(enabled=True),
        node_b: NodeRdmaCtlStatus(enabled=True),
        node_c: NodeRdmaCtlStatus(enabled=False),
    }
    cic = PlaceInstance(
        sharding=Sharding.Tensor,
        instance_meta=InstanceMeta.MlxJaccl,
        command_id=CommandId(),
        model_card=model_card,
        min_nodes=3,
    )

    # act / assert
    with pytest.raises(
        ValueError, match="Requested RDMA \\(MlxJaccl\\) but no RDMA-connected cycles"
    ):
        place_instance(
            cic,
            topology,
            {},
            node_memory,
            node_network,
            _metal_only(node_memory),
            node_rdma_ctl=node_rdma_ctl,
        )


def test_place_mlx_jaccl_rejects_when_node_rdma_ctl_missing(model_card: ModelCard):
    """A node with no observed rdma_ctl status must not participate in RDMA placement."""
    # arrange
    model_card = model_card.model_copy(
        update={"n_layers": 12, "storage_size": Memory.from_bytes(1500)}
    )
    topology, node_a, node_b, node_c, node_network = _build_three_node_rdma_topology()
    node_memory = {
        node_a: create_node_memory(500),
        node_b: create_node_memory(500),
        node_c: create_node_memory(500),
    }
    # node_c has no rdma_ctl entry at all
    node_rdma_ctl = {
        node_a: NodeRdmaCtlStatus(enabled=True),
        node_b: NodeRdmaCtlStatus(enabled=True),
    }
    cic = PlaceInstance(
        sharding=Sharding.Tensor,
        instance_meta=InstanceMeta.MlxJaccl,
        command_id=CommandId(),
        model_card=model_card,
        min_nodes=3,
    )

    # act / assert
    with pytest.raises(ValueError):
        place_instance(
            cic,
            topology,
            {},
            node_memory,
            node_network,
            _metal_only(node_memory),
            node_rdma_ctl=node_rdma_ctl,
        )


def _make_task(
    instance_id: InstanceId,
    status: TaskStatus = TaskStatus.Running,
) -> TextGeneration:
    return TextGeneration(
        task_id=TaskId(),
        task_status=status,
        instance_id=instance_id,
        command_id=CommandId(),
        task_params=TextGenerationTaskParams(
            model=ModelId("test-model"),
            input=[InputMessage(role="user", content=InputMessageContent("hello"))],
        ),
    )


def test_get_transition_events_delete_instance_cancels_running_tasks(
    instance: Instance,
):
    # arrange
    instance_id = InstanceId()
    current_instances: dict[InstanceId, Instance] = {instance_id: instance}
    target_instances: dict[InstanceId, Instance] = {}
    task = _make_task(instance_id, TaskStatus.Running)
    tasks = {task.task_id: task}

    # act
    events = get_transition_events(current_instances, target_instances, tasks)

    # assert – cancellation event should come before the deletion event
    assert len(events) == 2
    assert isinstance(events[0], TaskStatusUpdated)
    assert events[0].task_id == task.task_id
    assert events[0].task_status == TaskStatus.Cancelled
    assert isinstance(events[1], InstanceDeleted)
    assert events[1].instance_id == instance_id


def test_get_transition_events_delete_instance_cancels_pending_tasks(
    instance: Instance,
):
    # arrange
    instance_id = InstanceId()
    current_instances: dict[InstanceId, Instance] = {instance_id: instance}
    target_instances: dict[InstanceId, Instance] = {}
    task = _make_task(instance_id, TaskStatus.Pending)
    tasks = {task.task_id: task}

    # act
    events = get_transition_events(current_instances, target_instances, tasks)

    # assert
    assert len(events) == 2
    assert isinstance(events[0], TaskStatusUpdated)
    assert events[0].task_id == task.task_id
    assert events[0].task_status == TaskStatus.Cancelled
    assert isinstance(events[1], InstanceDeleted)


def test_get_transition_events_delete_instance_ignores_completed_tasks(
    instance: Instance,
):
    # arrange
    instance_id = InstanceId()
    current_instances: dict[InstanceId, Instance] = {instance_id: instance}
    target_instances: dict[InstanceId, Instance] = {}
    tasks = {
        t.task_id: t
        for t in [
            _make_task(instance_id, TaskStatus.Complete),
            _make_task(instance_id, TaskStatus.Failed),
            _make_task(instance_id, TaskStatus.TimedOut),
            _make_task(instance_id, TaskStatus.Cancelled),
        ]
    }

    # act
    events = get_transition_events(current_instances, target_instances, tasks)

    # assert – only the InstanceDeleted event, no cancellations
    assert len(events) == 1
    assert isinstance(events[0], InstanceDeleted)


def test_get_transition_events_delete_instance_cancels_only_matching_tasks(
    instance: Instance,
):
    # arrange
    instance_id_a = InstanceId()
    instance_id_b = InstanceId()
    current_instances: dict[InstanceId, Instance] = {
        instance_id_a: instance,
        instance_id_b: instance,
    }
    # only delete instance A, keep instance B
    target_instances: dict[InstanceId, Instance] = {instance_id_b: instance}

    task_a = _make_task(instance_id_a, TaskStatus.Running)
    task_b = _make_task(instance_id_b, TaskStatus.Running)
    tasks = {task_a.task_id: task_a, task_b.task_id: task_b}

    # act
    events = get_transition_events(current_instances, target_instances, tasks)

    # assert – only task_a should be cancelled
    cancel_events = [e for e in events if isinstance(e, TaskStatusUpdated)]
    delete_events = [e for e in events if isinstance(e, InstanceDeleted)]
    assert len(cancel_events) == 1
    assert cancel_events[0].task_id == task_a.task_id
    assert cancel_events[0].task_status == TaskStatus.Cancelled
    assert len(delete_events) == 1
    assert delete_events[0].instance_id == instance_id_a


def _make_shard_metadata(model_card: ModelCard) -> PipelineShardMetadata:
    return PipelineShardMetadata(
        model_card=model_card,
        device_rank=0,
        world_size=1,
        start_layer=0,
        end_layer=model_card.n_layers,
        n_layers=model_card.n_layers,
    )


def test_placement_prefers_cycle_with_downloaded_model(
    model_card: ModelCard,
) -> None:
    """When two cycles are otherwise equal, prefer the one with the model already downloaded."""
    topology = Topology()

    model_card = model_card.model_copy(update={"storage_size": Memory.from_bytes(500)})

    node_a = NodeId()
    node_b = NodeId()

    node_memory = {
        node_a: create_node_memory(1000),
        node_b: create_node_memory(1000),
    }
    node_network = {
        node_a: create_node_network(),
        node_b: create_node_network(),
    }

    topology.add_node(node_a)
    topology.add_node(node_b)
    # No connections between them — two single-node cycles

    shard_meta = _make_shard_metadata(model_card)

    # node_b has the model fully downloaded, node_a does not
    download_status = {
        node_b: [
            DownloadCompleted(
                node_id=node_b,
                shard_metadata=shard_meta,
                total=model_card.storage_size,
            ),
        ],
    }

    cic = place_instance_command(model_card)
    placements = place_instance(
        cic,
        topology,
        {},
        node_memory,
        node_network,
        _metal_only(node_memory),
        download_status=download_status,
    )

    assert len(placements) == 1
    instance = list(placements.values())[0]
    assigned_nodes = set(map(lambda it: it.node_id, instance.shard_assignments.shards))
    assert assigned_nodes == {node_b}


def test_placement_prefers_cycle_with_higher_download_progress(
    model_card: ModelCard,
) -> None:
    """When two cycles are otherwise equal, prefer the one with more download progress."""
    topology = Topology()

    model_card = model_card.model_copy(update={"storage_size": Memory.from_bytes(1000)})

    node_a = NodeId()
    node_b = NodeId()

    node_memory = {
        node_a: create_node_memory(1000),
        node_b: create_node_memory(1000),
    }
    node_network = {
        node_a: create_node_network(),
        node_b: create_node_network(),
    }

    topology.add_node(node_a)
    topology.add_node(node_b)

    shard_meta = _make_shard_metadata(model_card)

    # node_a: 30% downloaded, node_b: 80% downloaded
    download_status = {
        node_a: [
            DownloadOngoing(
                node_id=node_a,
                shard_metadata=shard_meta,
                download_progress=DownloadProgressData(
                    total=Memory.from_bytes(1000),
                    downloaded=Memory.from_bytes(300),
                    downloaded_this_session=Memory.from_bytes(300),
                    completed_files=0,
                    total_files=1,
                    speed=0.0,
                    eta_ms=0,
                    files={},
                ),
            ),
        ],
        node_b: [
            DownloadOngoing(
                node_id=node_b,
                shard_metadata=shard_meta,
                download_progress=DownloadProgressData(
                    total=Memory.from_bytes(1000),
                    downloaded=Memory.from_bytes(800),
                    downloaded_this_session=Memory.from_bytes(800),
                    completed_files=0,
                    total_files=1,
                    speed=0.0,
                    eta_ms=0,
                    files={},
                ),
            ),
        ],
    }

    cic = place_instance_command(model_card)
    placements = place_instance(
        cic,
        topology,
        {},
        node_memory,
        node_network,
        _metal_only(node_memory),
        download_status=download_status,
    )

    assert len(placements) == 1
    instance = list(placements.values())[0]
    assigned_nodes = set(map(lambda it: it.node_id, instance.shard_assignments.shards))
    assert assigned_nodes == {node_b}


def test_placement_does_not_prefer_cycle_with_failed_download(
    model_card: ModelCard,
) -> None:
    """A failed download should count as 0% — not preferred over a node with no download history."""
    topology = Topology()

    model_card = model_card.model_copy(update={"storage_size": Memory.from_bytes(500)})

    node_a = NodeId()
    node_b = NodeId()

    # node_a has slightly more RAM so it would win on the RAM tiebreaker
    node_memory = {
        node_a: create_node_memory(1001),
        node_b: create_node_memory(1000),
    }
    node_network = {
        node_a: create_node_network(),
        node_b: create_node_network(),
    }

    topology.add_node(node_a)
    topology.add_node(node_b)

    shard_meta = _make_shard_metadata(model_card)

    # node_b has a failed download — should not be preferred
    download_status = {
        node_b: [
            DownloadFailed(
                node_id=node_b,
                shard_metadata=shard_meta,
                error_message="connection reset",
            ),
        ],
    }

    cic = place_instance_command(model_card)
    placements = place_instance(
        cic,
        topology,
        {},
        node_memory,
        node_network,
        _metal_only(node_memory),
        download_status=download_status,
    )

    assert len(placements) == 1
    instance = list(placements.values())[0]
    assigned_nodes = set(map(lambda it: it.node_id, instance.shard_assignments.shards))
    # node_a should win on RAM tiebreaker since failed download scores 0.0
    assert assigned_nodes == {node_a}


def test_placement_rejects_when_model_backends_disjoint_from_engine(
    model_card: ModelCard,
):
    topology = Topology()
    node_id = NodeId()
    topology.add_node(node_id)
    node_memory = {node_id: create_node_memory(1000 * 1024)}
    node_network = {node_id: create_node_network()}

    cic = place_instance_command(
        model_card.model_copy(update={"backends": [Backend.Vllm]})
    )

    with pytest.raises(ValueError, match="cannot satisfy engine MlxRing"):
        place_instance(
            cic, topology, {}, node_memory, node_network, _metal_only(node_memory)
        )


def test_placement_rejects_when_only_some_nodes_support_backend(
    model_card: ModelCard,
):
    topology = Topology()
    node_a = NodeId()
    node_b = NodeId()
    node_c = NodeId()
    for n in (node_a, node_b, node_c):
        topology.add_node(n)

    eth = create_socket_connection(1)
    for src, dst in [
        (node_a, node_b),
        (node_b, node_c),
        (node_c, node_a),
        (node_a, node_c),
        (node_c, node_b),
        (node_b, node_a),
    ]:
        topology.add_connection(Connection(source=src, sink=dst, edge=eth))

    node_memory = {n: create_node_memory(500 * 1024) for n in (node_a, node_b, node_c)}
    node_network = {n: create_node_network() for n in (node_a, node_b, node_c)}
    node_backends = {
        node_a: [Backend.MlxMetal],
        node_b: [Backend.MlxMetal],
        node_c: [Backend.MlxCuda],  # the lone CUDA-only node breaks the cycle
    }

    cic = place_instance_command(
        model_card.model_copy(
            update={
                "backends": [Backend.MlxMetal],
                "n_layers": 12,
                "storage_size": Memory.from_kb(1500),
            }
        )
    )

    with pytest.raises(ValueError, match="No cycle where every node supports"):
        place_instance(cic, topology, {}, node_memory, node_network, node_backends)


def test_mlx_jaccl_rejects_cuda_only_cycle(model_card: ModelCard):
    topology, node_a, node_b, node_c, node_network = _build_three_node_rdma_topology()
    node_memory = {n: create_node_memory(500) for n in (node_a, node_b, node_c)}
    node_rdma_ctl = {
        n: NodeRdmaCtlStatus(enabled=True) for n in (node_a, node_b, node_c)
    }
    node_backends = {n: [Backend.MlxCuda] for n in (node_a, node_b, node_c)}

    cic = PlaceInstance(
        sharding=Sharding.Tensor,
        instance_meta=InstanceMeta.MlxJaccl,
        command_id=CommandId(),
        model_card=model_card.model_copy(
            update={
                "backends": [Backend.MlxMetal, Backend.MlxCuda],
                "n_layers": 12,
                "storage_size": Memory.from_bytes(1500),
            }
        ),
        min_nodes=3,
    )

    with pytest.raises(ValueError, match="No cycle where every node supports"):
        place_instance(
            cic,
            topology,
            {},
            node_memory,
            node_network,
            node_backends,
            node_rdma_ctl=node_rdma_ctl,
        )
