"""DGR-006 architecture-defined activation-boundary contract."""

from __future__ import annotations

import struct

import pytest

from meshnet_node.architecture_boundary import (
    Architecture,
    BoundaryStage,
    ProtocolIdentity,
    SamplingParameters,
    TailOutput,
    adapter_for,
)
from meshnet_node.native_protocol import ProtocolError, decode_bundle


def _f32(values: list[float]) -> bytes:
    return struct.pack("<" + "f" * len(values), *values)


def _values(payload: bytes) -> tuple[float, ...]:
    return struct.unpack("<" + "f" * (len(payload) // 4), payload)


def test_dense_whole_model_and_two_ranges_match_for_prefill_and_greedy_decode() -> None:
    adapter = adapter_for(Architecture.DENSE)
    embeddings = {3: [1.0, 2.0], 7: [3.0, 4.0]}
    head = adapter.bundle_from_token_ids([3, 7], lambda token: embeddings[token])

    def head_layers(bundle):
        return [value + 10.0 for value in _values(decode_bundle(bundle)["hidden_states"])]

    def tail_layers(residual: list[float]) -> tuple[list[float], int]:
        # The tail alone applies this fixture's final norm/head and greedy argmax.
        logits = [sum(residual) / len(residual) + 20.0, 0.0]
        return logits, max(range(len(logits)), key=logits.__getitem__)

    # Whole-model prefill retains the unnormalized residual locally. The
    # two-range lane sends that same residual before the tail-only head.
    whole_prefill = head_layers(head)
    seam = adapter.bundle_from_named_payloads(
        {"hidden_states": _f32(whole_prefill)}, shape=[1, 2, 2]
    )
    two_range_prefill = list(_values(decode_bundle(adapter.input_for(BoundaryStage.TAIL, seam))["hidden_states"]))
    assert two_range_prefill == whole_prefill

    whole_logits, whole_token = tail_layers(whole_prefill)
    two_logits, two_token = tail_layers(two_range_prefill)
    assert two_logits == whole_logits
    assert two_token == whole_token == 0

    # One greedy decode step is the same contract with [batch, token, hidden].
    decode_seam = adapter.bundle_from_named_payloads(
        {"hidden_states": _f32([5.0, 6.0])}, shape=[1, 1, 2]
    )
    assert tail_layers([5.0, 6.0]) == tail_layers(
        list(_values(decode_bundle(adapter.input_for(BoundaryStage.TAIL, decode_seam))["hidden_states"]))
    )


def test_middle_and_tail_reject_token_ids_and_require_boundary_bundle() -> None:
    adapter = adapter_for(Architecture.MOE)

    with pytest.raises(ProtocolError, match="head"):
        adapter.bundle_from_token_ids([1], lambda _: [1.0])

    with pytest.raises(ProtocolError, match="TensorBundle"):
        adapter.input_for(BoundaryStage.MIDDLE, None)


@pytest.mark.parametrize(
    ("architecture", "names"),
    [
        (Architecture.MOE, {"hidden_states", "router_logits"}),
        (Architecture.MLA, {"hidden_states", "mla_position_state"}),
    ],
)
def test_architecture_adapters_route_and_validate_their_named_sidebands(
    architecture: Architecture, names: set[str]
) -> None:
    adapter = adapter_for(architecture)
    bundle = adapter.bundle_from_named_payloads(
        {
            name: (_f32([1.0, 2.0]) if name == "hidden_states" else _f32([0.0]))
            for name in names
        }
    )

    assert set(decode_bundle(adapter.input_for(BoundaryStage.MIDDLE, bundle))) == names
    with pytest.raises(ProtocolError, match="requires"):
        adapter.bundle_from_named_payloads({"hidden_states": _f32([1.0, 2.0])})


def test_unknown_architecture_fails_closed() -> None:
    with pytest.raises(ProtocolError, match="unsupported architecture"):
        adapter_for("unchecked-name-substitution")


def test_typed_tail_result_binds_sampling_and_request_recipe_identity() -> None:
    adapter = adapter_for(Architecture.DENSE)
    identity = ProtocolIdentity(
        request_id="request-1",
        runtime_recipe_digest="sha256:recipe",
        chat_template_id="llama3",
        chat_template_version="2",
        reasoning_mode="max",
        architecture=Architecture.DENSE,
    )
    result = adapter.tail_result(
        identity=identity,
        sampling=SamplingParameters(temperature=0.0, top_p=1.0, top_k=0, seed=9),
        output=TailOutput.sampled_token(42),
    )

    assert result.identity.request_id == "request-1"
    assert result.sampled_token_id == 42
    assert result.output_kind == "sampled_token_id"
    assert result.message.WhichOneof("output") == "sampled_token_id"