"""Architecture-defined boundary input/output and dense-Llama parity (DGR-006).

These tests prove the boundary contract with a *pure-numpy* dense-Llama reference
model: no download, no GPU, no torch, no API credit. The reference implements the
same ``ShardComputation`` duck type the real llama.cpp/PyTorch backends expose, so
whole-model execution and a two-range (or three-range) split are the exact same
arithmetic applied to the exact same float32 residual stream. Splitting the layer
stack at a seam and shipping the *unnormalized* residual bundle across a simulated
process boundary must reproduce the whole-model tokens bit-for-bit.
"""

from __future__ import annotations

import numpy as np
import pytest

from meshnet_node.boundary_adapter import (
    BOUNDARY_SCHEMA_VERSION,
    BoundaryAdapter,
    BoundaryBundle,
    BoundaryContractError,
    SamplingContract,
    ShardRole,
    TailOutput,
    UncertifiedArchitectureError,
    certified_architecture,
    is_certified_architecture,
    role_for_range,
)

# Documented parity tolerance. The split path applies the identical layer
# functions in the identical order to the identical float32 arrays, so the
# residual seam is bit-exact in practice; the tolerance is a conservative guard.
PARITY_ATOL = 1e-6


# --------------------------------------------------------------------------- #
# Pure-numpy dense-Llama reference model (test fixture, not production).
# --------------------------------------------------------------------------- #


class _ReferenceDenseLlama:
    """A tiny deterministic dense-Llama: RMSNorm, RoPE attention, SwiGLU MLP."""

    architecture_adapter = "dense-llama"

    def __init__(
        self,
        *,
        vocab: int = 48,
        hidden: int = 32,
        n_layers: int = 6,
        n_heads: int = 4,
        intermediate: int = 64,
        rms_eps: float = 1e-6,
        rope_theta: float = 10000.0,
        seed: int = 20260715,
    ) -> None:
        assert hidden % n_heads == 0
        self.vocab = vocab
        self.hidden = hidden
        self.n_layers = n_layers
        self.n_heads = n_heads
        self.head_dim = hidden // n_heads
        assert self.head_dim % 2 == 0
        self.rms_eps = rms_eps
        self.rope_theta = rope_theta

        rng = np.random.default_rng(seed)

        def w(*shape: int) -> np.ndarray:
            return (rng.standard_normal(shape) * 0.08).astype(np.float32)

        self.embed = w(vocab, hidden)
        self.layers = []
        for _ in range(n_layers):
            self.layers.append(
                {
                    "in_ln": (1.0 + rng.standard_normal(hidden) * 0.02).astype(np.float32),
                    "q": w(hidden, hidden),
                    "k": w(hidden, hidden),
                    "v": w(hidden, hidden),
                    "o": w(hidden, hidden),
                    "post_ln": (1.0 + rng.standard_normal(hidden) * 0.02).astype(np.float32),
                    "gate": w(intermediate, hidden),
                    "up": w(intermediate, hidden),
                    "down": w(hidden, intermediate),
                }
            )
        self.final_ln = (1.0 + rng.standard_normal(hidden) * 0.02).astype(np.float32)
        self.lm_head_w = w(vocab, hidden)

        inv_freq = 1.0 / (
            rope_theta ** (np.arange(0, self.head_dim, 2, dtype=np.float32) / self.head_dim)
        )
        self.inv_freq = inv_freq.astype(np.float32)

    # -- primitive ops -----------------------------------------------------
    def _rmsnorm(self, x: np.ndarray, weight: np.ndarray) -> np.ndarray:
        variance = np.mean(x.astype(np.float32) ** 2, axis=-1, keepdims=True)
        normed = x / np.sqrt(variance + self.rms_eps)
        return (normed * weight).astype(np.float32)

    def _rope(self, positions: np.ndarray):
        # positions: (batch, seq) -> cos/sin: (batch, seq, head_dim)
        angles = positions[..., None].astype(np.float32) * self.inv_freq[None, None, :]
        emb = np.concatenate([angles, angles], axis=-1)
        return np.cos(emb).astype(np.float32), np.sin(emb).astype(np.float32)

    @staticmethod
    def _rotate_half(x: np.ndarray) -> np.ndarray:
        half = x.shape[-1] // 2
        return np.concatenate([-x[..., half:], x[..., :half]], axis=-1)

    def _apply_rope(self, t: np.ndarray, cos: np.ndarray, sin: np.ndarray) -> np.ndarray:
        # t: (batch, n_heads, seq, head_dim); cos/sin: (batch, seq, head_dim)
        cos = cos[:, None, :, :]
        sin = sin[:, None, :, :]
        return t * cos + self._rotate_half(t) * sin

    def _attention(self, x: np.ndarray, layer: dict, positions: np.ndarray) -> np.ndarray:
        batch, seq, _ = x.shape
        q = (x @ layer["q"].T).reshape(batch, seq, self.n_heads, self.head_dim)
        k = (x @ layer["k"].T).reshape(batch, seq, self.n_heads, self.head_dim)
        v = (x @ layer["v"].T).reshape(batch, seq, self.n_heads, self.head_dim)
        q = q.transpose(0, 2, 1, 3)
        k = k.transpose(0, 2, 1, 3)
        v = v.transpose(0, 2, 1, 3)
        cos, sin = self._rope(positions)
        q = self._apply_rope(q, cos, sin)
        k = self._apply_rope(k, cos, sin)
        scores = (q @ k.transpose(0, 1, 3, 2)) / np.sqrt(self.head_dim)
        causal = np.triu(np.full((seq, seq), -1e30, dtype=np.float32), k=1)
        scores = scores + causal[None, None, :, :]
        scores = scores - scores.max(axis=-1, keepdims=True)
        weights = np.exp(scores)
        weights = weights / weights.sum(axis=-1, keepdims=True)
        out = weights @ v
        out = out.transpose(0, 2, 1, 3).reshape(batch, seq, self.hidden)
        return (out @ layer["o"].T).astype(np.float32)

    def _mlp(self, x: np.ndarray, layer: dict) -> np.ndarray:
        gate = x @ layer["gate"].T
        up = x @ layer["up"].T
        silu = gate * (1.0 / (1.0 + np.exp(-gate)))
        return ((silu * up) @ layer["down"].T).astype(np.float32)

    def _run_layer(self, x: np.ndarray, layer: dict, positions: np.ndarray) -> np.ndarray:
        h = x + self._attention(self._rmsnorm(x, layer["in_ln"]), layer, positions)
        h = h + self._mlp(self._rmsnorm(h, layer["post_ln"]), layer)
        return h.astype(np.float32)


class _ReferenceShard:
    """A contiguous inclusive layer range of the reference model.

    Satisfies the ``ShardComputation`` duck type used by ``BoundaryAdapter``.
    """

    def __init__(
        self,
        model: _ReferenceDenseLlama,
        start_layer: int,
        end_layer: int,
        *,
        architecture_adapter: str | None = None,
    ) -> None:
        self._model = model
        self.start_layer = start_layer
        self.end_layer = end_layer
        self.total_layers = model.n_layers
        self.architecture_adapter = architecture_adapter or model.architecture_adapter

    def embed_tokens(self, token_ids: np.ndarray) -> np.ndarray:
        return self._model.embed[np.asarray(token_ids)]

    def run_layers(self, hidden: np.ndarray, *, positions: np.ndarray) -> np.ndarray:
        h = np.asarray(hidden, dtype=np.float32)
        for idx in range(self.start_layer, self.end_layer + 1):
            h = self._model._run_layer(h, self._model.layers[idx], positions)
        return h

    def final_norm(self, hidden: np.ndarray) -> np.ndarray:
        return self._model._rmsnorm(np.asarray(hidden, dtype=np.float32), self._model.final_ln)

    def lm_head(self, hidden: np.ndarray) -> np.ndarray:
        return np.asarray(hidden, dtype=np.float32) @ self._model.lm_head_w.T


# --------------------------------------------------------------------------- #
# Whole-model and split reference drivers.
# --------------------------------------------------------------------------- #


def _whole_model_next_token(model: _ReferenceDenseLlama, token_ids: list[int]) -> TailOutput:
    shard = _ReferenceShard(model, 0, model.n_layers - 1)
    adapter = BoundaryAdapter(shard)
    result = adapter.forward(token_ids=np.asarray(token_ids)[None, :])
    assert isinstance(result, TailOutput)
    return result


def _split_next_token(
    model: _ReferenceDenseLlama,
    token_ids: list[int],
    cut_points: list[int],
    *,
    through_wire: bool = True,
) -> TailOutput:
    """Run the model as N contiguous ranges, shipping the bundle across each seam.

    ``cut_points`` are the last (inclusive) layer of each non-final range.
    """
    bounds = _ranges_from_cuts(cut_points, model.n_layers)
    boundary: BoundaryBundle | None = None
    result: BoundaryBundle | TailOutput | None = None
    for i, (start, end) in enumerate(bounds):
        shard = _ReferenceShard(model, start, end)
        adapter = BoundaryAdapter(shard)
        if i == 0:
            result = adapter.forward(token_ids=np.asarray(token_ids)[None, :])
        else:
            assert isinstance(boundary, BoundaryBundle)
            incoming = BoundaryBundle.unpack(boundary.pack()) if through_wire else boundary
            result = adapter.forward(boundary=incoming)
        if isinstance(result, BoundaryBundle):
            boundary = result
    assert isinstance(result, TailOutput)
    return result


def _ranges_from_cuts(cut_points: list[int], n_layers: int) -> list[tuple[int, int]]:
    bounds: list[tuple[int, int]] = []
    start = 0
    for cut in cut_points:
        bounds.append((start, cut))
        start = cut + 1
    bounds.append((start, n_layers - 1))
    return bounds


def _greedy_generate(next_token_fn, prompt: list[int], n_new: int) -> list[int]:
    tokens = list(prompt)
    generated: list[int] = []
    for _ in range(n_new):
        out = next_token_fn(tokens)
        tokens.append(out.token_id)
        generated.append(out.token_id)
    return generated


# --------------------------------------------------------------------------- #
# Certification / fail-closed.
# --------------------------------------------------------------------------- #


def test_dense_llama_and_aliases_are_certified():
    "Dense Llama-family identifiers all resolve to the one certified adapter.\n\nTags: node, boundary"
    for name in ("dense-llama", "llama", "LlamaForCausalLM", "LlamaModel"):
        boundary = certified_architecture(name)
        assert boundary.adapter == "dense-llama"
        assert boundary.boundary_tensor_name == "residual_stream"
        assert is_certified_architecture(name)


@pytest.mark.parametrize("name", ["qwen3", "qwen3-moe", "mixtral", "gpt2", "", None, 123])
def test_uncertified_architectures_fail_closed(name):
    "Uncertified architectures raise instead of guessing a tensor layout.\n\nTags: node, boundary"
    assert not is_certified_architecture(name)
    with pytest.raises(UncertifiedArchitectureError):
        certified_architecture(name)


def test_adapter_construction_fails_closed_for_uncertified_backend():
    "Building the adapter over an uncertified computation fails closed.\n\nTags: node, boundary"
    model = _ReferenceDenseLlama()
    shard = _ReferenceShard(model, 0, 2, architecture_adapter="qwen3-moe")
    with pytest.raises(UncertifiedArchitectureError):
        BoundaryAdapter(shard)


# --------------------------------------------------------------------------- #
# Roles.
# --------------------------------------------------------------------------- #


def test_role_classification():
    "Range endpoints map to head/middle/tail/full roles.\n\nTags: node, boundary"
    assert role_for_range(0, 2, 6) is ShardRole.HEAD
    assert role_for_range(2, 3, 6) is ShardRole.MIDDLE
    assert role_for_range(4, 5, 6) is ShardRole.TAIL
    assert role_for_range(0, 5, 6) is ShardRole.FULL
    assert ShardRole.HEAD.owns_embedding and not ShardRole.HEAD.owns_final_head
    assert ShardRole.TAIL.owns_final_head and not ShardRole.TAIL.owns_embedding


# --------------------------------------------------------------------------- #
# Input-side contract.
# --------------------------------------------------------------------------- #


def test_head_accepts_token_ids_and_owns_embedding():
    "The head embeds token IDs and refuses an upstream boundary bundle.\n\nTags: node, boundary"
    model = _ReferenceDenseLlama()
    head = BoundaryAdapter(_ReferenceShard(model, 0, 2))
    out = head.forward(token_ids=[1, 2, 3])
    assert isinstance(out, BoundaryBundle)

    # Head owns embedding: a residual bundle from upstream is a contract error.
    bundle = out
    with pytest.raises(BoundaryContractError, match="head owns token embedding"):
        head.forward(boundary=bundle)


def test_middle_and_tail_bypass_embedding_and_require_the_bundle():
    "Middle/tail Shards reject token IDs and demand the named boundary bundle.\n\nTags: node, boundary"
    model = _ReferenceDenseLlama()
    tail = BoundaryAdapter(_ReferenceShard(model, 3, 5))
    with pytest.raises(BoundaryContractError, match="bypass token embedding"):
        tail.forward(token_ids=[1, 2, 3])
    with pytest.raises(BoundaryContractError, match="must receive the named boundary bundle"):
        tail.forward()


def test_boundary_seam_layer_mismatch_is_rejected():
    "A bundle handed to the wrong range (seam layer mismatch) is rejected.\n\nTags: node, boundary"
    model = _ReferenceDenseLlama()
    head = BoundaryAdapter(_ReferenceShard(model, 0, 2))
    bundle = head.forward(token_ids=[1, 2, 3])
    assert isinstance(bundle, BoundaryBundle)
    assert bundle.next_layer == 3

    # A range that starts at layer 4 must not accept a bundle cut at layer 3.
    wrong = BoundaryAdapter(_ReferenceShard(model, 4, 5))
    with pytest.raises(BoundaryContractError, match="starts at layer 4"):
        wrong.forward(boundary=bundle)


def test_normalized_bundle_is_rejected():
    "A normalized residual is not the architecture-defined boundary.\n\nTags: node, boundary"
    model = _ReferenceDenseLlama()
    head = BoundaryAdapter(_ReferenceShard(model, 0, 2))
    bundle = head.forward(token_ids=[1, 2, 3])
    assert isinstance(bundle, BoundaryBundle)
    normalized = BoundaryBundle(
        architecture_adapter=bundle.architecture_adapter,
        schema_version=bundle.schema_version,
        tensor_name=bundle.tensor_name,
        residual=bundle.residual,
        positions=bundle.positions,
        next_layer=bundle.next_layer,
        normalized=True,
    )
    tail = BoundaryAdapter(_ReferenceShard(model, 3, 5))
    with pytest.raises(BoundaryContractError, match="UNNORMALIZED"):
        tail.forward(boundary=normalized)


# --------------------------------------------------------------------------- #
# Output-side contract.
# --------------------------------------------------------------------------- #


def test_non_tail_emits_unnormalized_full_row_boundary():
    "A non-tail Shard emits the unnormalized residual with every position row.\n\nTags: node, boundary"
    model = _ReferenceDenseLlama()
    tokens = [3, 7, 1, 9, 2]
    head = BoundaryAdapter(_ReferenceShard(model, 0, 2))
    bundle = head.forward(token_ids=tokens)
    assert isinstance(bundle, BoundaryBundle)
    assert bundle.normalized is False
    assert bundle.tensor_name == "residual_stream"
    assert bundle.schema_version == BOUNDARY_SCHEMA_VERSION
    assert bundle.next_layer == 3
    # No tail-only row pruning: all sequence positions are forwarded.
    assert bundle.residual.shape == (1, len(tokens), model.hidden)
    assert bundle.positions.shape == (1, len(tokens))

    # The emitted residual must be exactly the whole model's residual after layer 2
    # (i.e. before any final norm) — prove it is NOT normalized.
    positions = np.arange(len(tokens))[None, :]
    hidden = model.embed[np.asarray(tokens)][None, :]
    for idx in range(0, 3):
        hidden = model._run_layer(hidden, model.layers[idx], positions)
    assert np.allclose(bundle.residual, hidden, atol=0)
    assert not np.allclose(bundle.residual, model._rmsnorm(hidden, model.final_ln))


def test_tail_emits_pruned_logits_through_the_sampling_contract():
    "The tail prunes to the final row and samples through an explicit contract.\n\nTags: node, boundary"
    model = _ReferenceDenseLlama()
    out = _whole_model_next_token(model, [4, 8, 15, 16, 23])
    assert isinstance(out, TailOutput)
    assert out.logits.shape == (1, model.vocab)  # tail-only row pruning to last row
    assert out.sampling.mode == "greedy"
    assert 0 <= out.token_id < model.vocab
    assert out.token_id == int(np.argmax(out.logits[0]))


def test_sampling_contract_rejects_uncertified_modes():
    "Only the certified greedy sampling mode is accepted.\n\nTags: node, boundary"
    with pytest.raises(BoundaryContractError):
        SamplingContract(mode="top_p")


# --------------------------------------------------------------------------- #
# The core parity gate.
# --------------------------------------------------------------------------- #


def test_two_range_prefill_parity_matches_whole_model():
    "Whole-model vs two-range prefill produce the same next-token logits and token.\n\nTags: node, boundary, parity"
    model = _ReferenceDenseLlama()
    prompt = [5, 12, 3, 41, 7, 19, 2, 33]

    whole = _whole_model_next_token(model, prompt)
    split = _split_next_token(model, prompt, cut_points=[2])

    assert np.allclose(whole.logits, split.logits, atol=PARITY_ATOL)
    assert whole.token_id == split.token_id


def test_three_range_prefill_parity_exercises_the_middle_role():
    "A head/middle/tail split reproduces whole-model prefill through two seams.\n\nTags: node, boundary, parity"
    model = _ReferenceDenseLlama()
    prompt = [9, 1, 44, 6, 30, 11]

    whole = _whole_model_next_token(model, prompt)
    split = _split_next_token(model, prompt, cut_points=[1, 3])

    assert np.allclose(whole.logits, split.logits, atol=PARITY_ATOL)
    assert whole.token_id == split.token_id


def test_two_range_greedy_decode_parity_matches_whole_model():
    "Whole-model vs two-range greedy decode produce identical token sequences.\n\nTags: node, boundary, parity"
    model = _ReferenceDenseLlama()
    prompt = [2, 17, 8, 25]
    n_new = 12

    whole_tokens = _greedy_generate(
        lambda toks: _whole_model_next_token(model, toks), prompt, n_new
    )
    split_tokens = _greedy_generate(
        lambda toks: _split_next_token(model, toks, cut_points=[2]), prompt, n_new
    )

    assert whole_tokens == split_tokens
    assert len(whole_tokens) == n_new


def test_boundary_bundle_wire_round_trip_is_exact():
    "Packing and unpacking the boundary bundle reconstructs the exact arrays.\n\nTags: node, boundary"
    model = _ReferenceDenseLlama()
    head = BoundaryAdapter(_ReferenceShard(model, 0, 2))
    bundle = head.forward(token_ids=[1, 2, 3, 4])
    assert isinstance(bundle, BoundaryBundle)

    restored = BoundaryBundle.unpack(bundle.pack())
    assert np.array_equal(restored.residual, bundle.residual)
    assert np.array_equal(restored.positions, bundle.positions)
    assert restored.next_layer == bundle.next_layer
    assert restored.architecture_adapter == bundle.architecture_adapter

    fields = bundle.named_tensor_fields()
    assert fields["name"] == "residual_stream"
    assert fields["shape"] == [1, 4, model.hidden]
    assert fields["byte_order"] in ("little", "big")


def test_alias_architecture_still_parity_matches():
    "A Shard advertised as 'llama' interoperates with the canonical adapter.\n\nTags: node, boundary, parity"
    model = _ReferenceDenseLlama()
    prompt = [7, 3, 22, 5]

    whole = _whole_model_next_token(model, prompt)

    # Head advertises 'LlamaForCausalLM', tail advertises 'llama'; both certify to
    # the same canonical adapter, so the seam contract still matches.
    head = BoundaryAdapter(_ReferenceShard(model, 0, 2, architecture_adapter="LlamaForCausalLM"))
    bundle = head.forward(token_ids=np.asarray(prompt)[None, :])
    assert isinstance(bundle, BoundaryBundle)
    tail = BoundaryAdapter(_ReferenceShard(model, 3, 5, architecture_adapter="llama"))
    split = tail.forward(boundary=BoundaryBundle.unpack(bundle.pack()))
    assert isinstance(split, TailOutput)

    assert np.allclose(whole.logits, split.logits, atol=PARITY_ATOL)
    assert whole.token_id == split.token_id