"""DGR-021: versioned activation envelope for shard traffic."""

from __future__ import annotations

import json
import hashlib

import pytest

from meshnet_node.model_backend import TensorPayload
from meshnet_node.protocol import ActivationEnvelope, NamedTensor, TensorFragment


def _payload(body: bytes = b"\x01\x02\x03\x04", *, shape=(1, 2, 1)) -> TensorPayload:
    return TensorPayload(
        body=body,
        shape=list(shape),
        attention_mask_header="1,2:AAAA",
        position_ids_header="1,2:BBBB",
        past_len=7,
    )


def test_envelope_roundtrips_deterministically():
    payload = _payload()
    env_a = payload.to_envelope(
        name="activations",
        request_id="req-1",
        work_id="work-1",
        route_session="route-1",
        route_epoch=9,
        shard_start=12,
        effective_start=18,
        phase="prefill",
        position=0,
        idempotency_step=0,
        extensions={"zeta": 3, "alpha": "x"},
    )
    env_b = ActivationEnvelope.from_bytes(env_a.to_bytes())

    assert env_a.to_bytes() == env_b.to_bytes()
    assert env_b.request_id == "req-1"
    assert env_b.work_id == "work-1"
    assert env_b.route_session == "route-1"
    assert env_b.shard_start == 12
    assert env_b.effective_start == 18
    assert env_b.phase == "prefill"
    assert env_b.tensors[0].body() == payload.body
    assert env_b.tensors[0].extensions["past_len"] == 7
    assert env_b.extensions == {"alpha": "x", "zeta": 3}


def test_fragmentation_and_checksums_are_bounded():
    body = b"0123456789abcdef"
    tensor = NamedTensor.from_bytes(
        name="activations",
        body=body,
        shape=[1, 8, 1],
        dtype="bfloat16",
        max_fragment_bytes=5,
    )

    assert len(tensor.fragments) == 4
    assert all(isinstance(fragment, TensorFragment) for fragment in tensor.fragments)
    assert all(len(fragment.body) <= 5 for fragment in tensor.fragments)
    assert tensor.body() == body
    assert tensor.checksum == hashlib.sha256(body).hexdigest()


def test_unknown_fields_are_preserved_on_roundtrip():
    env = _payload().to_envelope(
        name="activations",
        request_id="req-2",
        work_id="work-2",
        route_session="route-2",
        route_epoch=1,
        shard_start=0,
        effective_start=0,
        phase="decode",
        position=3,
        idempotency_step=2,
    )
    raw = json.loads(env.to_bytes())
    raw["future_top_level"] = {"x": 1}
    raw["tensors"][0]["future_tensor_field"] = "ok"
    roundtrip = ActivationEnvelope.from_bytes(json.dumps(raw).encode())

    assert roundtrip.extensions["future_top_level"] == {"x": 1}
    assert roundtrip.tensors[0].extensions["future_tensor_field"] == "ok"


def test_size_limit_rejects_large_envelopes():
    env = _payload(body=b"x" * 64).to_envelope(
        name="activations",
        request_id="req-3",
        work_id="work-3",
        route_session="route-3",
        route_epoch=1,
        shard_start=0,
        effective_start=0,
        phase="decode",
        position=5,
        idempotency_step=4,
    )

    with pytest.raises(ValueError, match="size limit"):
        env.to_bytes(max_bytes=32)


def test_tensor_payload_roundtrips_through_envelope():
    payload = _payload(body=b"abcde", shape=(1, 5, 1))
    env = payload.to_envelope(
        name="activations",
        request_id="req-4",
        work_id="work-4",
        route_session="route-4",
        route_epoch=2,
        shard_start=6,
        effective_start=11,
        phase="decode",
        position=8,
        idempotency_step=9,
    )

    restored = TensorPayload.from_envelope(env)
    assert restored.body == payload.body
    assert restored.shape == payload.shape
    assert restored.attention_mask_header == payload.attention_mask_header
    assert restored.position_ids_header == payload.position_ids_header
    assert restored.past_len == payload.past_len