"""US-002 integration test: two-node shard pipeline with activation tensors flowing A → B."""

import json
import urllib.request

from meshnet_node.server import StubNodeServer
from meshnet_gateway.server import GatewayServer


def test_two_node_activation_pipeline():
    # Node A owns layers 0-15 (first shard, not last — returns activations).
    node_a = StubNodeServer(shard_start=0, shard_end=15, is_last_shard=False)
    port_a = node_a.start()

    # Node B owns layers 16-31 (last shard — returns text after receiving activations).
    node_b = StubNodeServer(shard_start=16, shard_end=31, is_last_shard=True)
    port_b = node_b.start()

    # Gateway hardcodes the two-node inference route for this story.
    gateway = GatewayServer(
        inference_route=[
            f"http://127.0.0.1:{port_a}",
            f"http://127.0.0.1:{port_b}",
        ]
    )
    gateway_port = gateway.start()

    try:
        payload = json.dumps({
            "model": "stub-model",
            "messages": [{"role": "user", "content": "hello"}],
        }).encode()

        req = urllib.request.Request(
            f"http://127.0.0.1:{gateway_port}/v1/chat/completions",
            data=payload,
            headers={"Content-Type": "application/json"},
            method="POST",
        )

        with urllib.request.urlopen(req) as resp:
            assert resp.status == 200
            body = json.loads(resp.read())

        # Gateway must assemble a valid OpenAI-format response.
        assert body["id"] == "chatcmpl-stub"
        assert body["object"] == "chat.completion"
        assert body["model"] == "stub-model"
        assert len(body["choices"]) == 1
        choice = body["choices"][0]
        assert choice["message"]["role"] == "assistant"
        assert choice["message"]["content"] == "stub response to: hello"
        assert choice["finish_reason"] == "stop"

        # Node B must have received an activation tensor from node A — this is
        # the observable proof that the A → B activation hop occurred.
        assert node_b.received_activations, "node B did not receive activation tensors from node A"

        # Node A must NOT have received activations (it is the first shard).
        assert not node_a.received_activations, "node A should not receive activations — it is first in the route"
    finally:
        gateway.stop()
        node_a.stop()
        node_b.stop()