"""DGR-035 dense range boundary execution contract."""

from __future__ import annotations

import struct

import pytest

from meshnet_node.architecture_boundary import (
    DENSE_LLAMA_ARCHITECTURE,
    DENSE_RESIDUAL_BOUNDARY_V1,
    DenseLayerRange,
    DenseRangeBoundaryExecutor,
    TailOutput,
)
from meshnet_node.native_protocol import HIDDEN_STATES, ProtocolError
from meshnet_node.shard_engine import BoundaryBundle, EngineTensor


def _tensor(values: tuple[float, ...]) -> EngineTensor:
    return EngineTensor(HIDDEN_STATES, (1, len(values)), "f32", struct.pack("<" + "f" * len(values), *values))


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


def _embed(token_ids: tuple[int, ...]) -> EngineTensor:
    return _tensor(tuple(float(token) for token in token_ids))


def _layers(residual: EngineTensor) -> EngineTensor:
    return _tensor(tuple(value + 10.0 for value in _values(residual)))


def test_head_and_middle_handoff_the_same_unnormalized_named_residual() -> None:
    head = DenseRangeBoundaryExecutor(DenseLayerRange(0, 1, 4), embed_tokens=_embed, run_layers=_layers)
    middle = DenseRangeBoundaryExecutor(DenseLayerRange(2, 2, 4), embed_tokens=_embed, run_layers=_layers)

    head_out = head.execute(token_ids=(1, 2))
    assert isinstance(head_out, BoundaryBundle)
    assert head_out.architecture == DENSE_LLAMA_ARCHITECTURE
    assert head_out.boundary_point == DENSE_RESIDUAL_BOUNDARY_V1
    assert _values(head_out.tensors[0]) == (11.0, 12.0)

    middle_out = middle.execute(boundary=head_out)
    assert isinstance(middle_out, BoundaryBundle)
    # The raw residual is carried through. No tail norm/output or row pruning
    # can run because this executor has no tail callback.
    assert _values(middle_out.tensors[0]) == (21.0, 22.0)


def test_tail_bypasses_embedding_and_has_an_explicit_sampled_output_contract() -> None:
    tail = DenseRangeBoundaryExecutor(
        DenseLayerRange(3, 3, 4),
        embed_tokens=_embed,
        run_layers=_layers,
        tail_output=lambda residual: TailOutput.sampled_token(int(sum(_values(residual)))),
    )
    boundary = BoundaryBundle((_tensor((3.0, 4.0)),), DENSE_LLAMA_ARCHITECTURE, DENSE_RESIDUAL_BOUNDARY_V1)

    result = tail.execute(boundary=boundary)
    assert result == TailOutput.sampled_token(27)
    with pytest.raises(ProtocolError, match="requires"):
        tail.execute(token_ids=(3,))


def test_uncertified_architecture_and_incompatible_schema_fail_closed() -> None:
    with pytest.raises(ProtocolError, match="only certifies"):
        DenseLayerRange(0, 0, 1, architecture="unchecked")

    middle = DenseRangeBoundaryExecutor(DenseLayerRange(1, 1, 3), embed_tokens=_embed, run_layers=_layers)
    bad_architecture = BoundaryBundle((_tensor((1.0,)),), "moe", DENSE_RESIDUAL_BOUNDARY_V1)
    with pytest.raises(ProtocolError, match="not certified"):
        middle.execute(boundary=bad_architecture)
    bad_schema = BoundaryBundle((_tensor((1.0,)),), DENSE_LLAMA_ARCHITECTURE, "post_middle_residual")
    with pytest.raises(ProtocolError, match="incompatible"):
        middle.execute(boundary=bad_schema)


def test_only_tail_can_be_given_final_norm_and_output_ownership() -> None:
    with pytest.raises(ProtocolError, match="only a dense tail"):
        DenseRangeBoundaryExecutor(
            DenseLayerRange(0, 1, 4), embed_tokens=_embed, run_layers=_layers, tail_output=TailOutput.sampled_token
        )
    with pytest.raises(ProtocolError, match="only a dense tail"):
        DenseRangeBoundaryExecutor(DenseLayerRange(3, 3, 4), embed_tokens=_embed, run_layers=_layers)