"""Integration tests for two-node shard pipelines with activation tensors flowing A to 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).
    "Two node activation pipeline\n\nTags: inference, integration, node, startup"
    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()


def test_binary_activation_pipeline_chunks_512_token_prefill():
    "Binary activation pipeline chunks 512 token prefill\n\nTags: inference, integration, node, startup"
    node_a = StubNodeServer(shard_start=0, shard_end=15, is_last_shard=False)
    port_a = node_a.start()
    node_b = StubNodeServer(shard_start=16, shard_end=31, is_last_shard=True)
    port_b = node_b.start()
    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:
        prompt = " ".join(f"tok{i}" for i in range(512))
        payload = json.dumps({
            "model": "stub-model",
            "messages": [{"role": "user", "content": prompt}],
        }).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())

        assert body["choices"][0]["message"]["content"] == f"stub response to: {prompt}"
        assert node_a.forward_chunk_count == 4
        assert node_b.forward_chunk_count == 4
        assert not node_a.received_activations
        assert node_b.received_activations

        chunk_responses = gateway.last_binary_chunk_responses
        assert len(chunk_responses) == 4
        for index, chunk in enumerate(chunk_responses):
            assert chunk.shape == [1, 128, 64]
            assert chunk.dtype == "bfloat16"
            assert chunk.chunk_index == index
            assert chunk.chunk_total == 4
            assert chunk.headers["x-meshnet-wire"] == "2"
            assert len(chunk.body) == 1 * 128 * 64 * 2
    finally:
        gateway.stop()
        node_a.stop()
        node_b.stop()