"""Architecture-defined boundary input/output for distributed Shards (DGR-006).

A public-network Shard is a contiguous range of transformer layers (RALPH runtime
decision #1). For disjoint processes to reproduce whole-model execution, every
Shard must agree on *exactly* what boundary state it consumes and emits:

* The **head** owns token embedding: it accepts token IDs and turns them into the
  residual stream. No other Shard may embed tokens.
* **Middle and tail** Shards bypass token embedding entirely; they accept the named
  boundary bundle (the residual stream handed over by the previous range).
* A **non-tail** Shard emits the *unnormalized* architecture-defined residual /
  boundary — before the final norm, before the LM head, and before any tail-only
  row pruning — so the next range sees precisely the state the whole model would
  have at that layer index.
* The **tail** owns the final norm + LM head and turns the residual into logits or
  a sampled token through an explicit sampling contract.

This module is deliberately backend-agnostic. It enforces the boundary *contract*
and defers the arithmetic to a ``ShardComputation`` (a duck-typed object exposing
``embed_tokens`` / ``run_layers`` / ``final_norm`` / ``lm_head``). The pinned
llama.cpp worker (DGR-008) and the reference PyTorch backend both satisfy that
protocol, and the numpy reference model in the tests proves whole-model versus
two-range parity without any download, GPU, or API credit.

The adapter **fails closed** for uncertified architectures: only architectures
that have passed real certification (dense Llama-family first, per RALPH runtime
decision #13) are accepted. Everything else raises rather than silently guessing a
tensor layout — Qwen3/Qwen3-MoE stays registered-but-dark until DGR-015 certifies
its own adapter.
"""

from __future__ import annotations

from dataclasses import dataclass, field
from enum import Enum
from typing import Any

import numpy as np

# The boundary bundle wire schema version. This is the ``boundary_schema_version``
# carried by ``runtime_recipe.RuntimeRecipeIdentity``; a receiver refuses a bundle
# whose schema it does not implement (forward/backward compatibility is a routing
# concern, not a silent reinterpretation).
BOUNDARY_SCHEMA_VERSION = 1


class BoundaryAdapterError(RuntimeError):
    """Base class for boundary-contract violations."""


class UncertifiedArchitectureError(BoundaryAdapterError):
    """Raised when a boundary adapter is requested for an uncertified architecture.

    Failing closed here is a safety property: an unknown architecture has an
    unknown tensor layout, so guessing where the residual boundary lives would
    silently corrupt distributed output. The architecture must pass real
    certification first.
    """


class BoundaryContractError(BoundaryAdapterError):
    """Raised when a Shard is fed the wrong boundary input for its role.

    Examples: a head handed a residual bundle instead of token IDs, a middle
    Shard handed token IDs it must not embed, or a boundary bundle whose
    architecture / schema / seam layer does not match the receiving range.
    """


@dataclass(frozen=True)
class ArchitectureBoundary:
    """The architecture-defined boundary description for one certified adapter.

    These fields are what makes the boundary *architecture-defined* rather than a
    hardcoded assumption: the residual tensor name, whether the tail normalizes
    before the LM head, and whether row pruning is a tail-only concern all come
    from here.
    """

    adapter: str
    boundary_tensor_name: str
    boundary_schema_version: int
    normalizes_before_head: bool
    prunes_rows_at_tail: bool


# Certified architectures only. Dense Llama-family is first (RALPH runtime decision
# #13 + native discipline). Aliases map the many spellings a runtime recipe /
# GGUF / HF config may use onto the single canonical adapter id. Anything not in
# this table fails closed.
_DENSE_LLAMA = ArchitectureBoundary(
    adapter="dense-llama",
    boundary_tensor_name="residual_stream",
    boundary_schema_version=BOUNDARY_SCHEMA_VERSION,
    normalizes_before_head=True,
    prunes_rows_at_tail=True,
)

_CERTIFIED_ARCHITECTURES: dict[str, ArchitectureBoundary] = {
    "dense-llama": _DENSE_LLAMA,
    "dense_llama": _DENSE_LLAMA,
    "llama": _DENSE_LLAMA,
    "llamaforcausallm": _DENSE_LLAMA,
    "llamamodel": _DENSE_LLAMA,
}


def certified_architecture(name: Any) -> ArchitectureBoundary:
    """Return the certified boundary description for ``name`` or fail closed.

    ``name`` may be the canonical adapter id (``dense-llama``), an HF architecture
    class (``LlamaForCausalLM``), or a GGUF/config ``model_type`` (``llama``).
    Uncertified architectures raise ``UncertifiedArchitectureError``.
    """
    if not isinstance(name, str) or not name.strip():
        raise UncertifiedArchitectureError(
            "architecture adapter must be a non-empty string; "
            "the boundary adapter refuses to guess a tensor layout"
        )
    key = name.strip().lower()
    boundary = _CERTIFIED_ARCHITECTURES.get(key)
    if boundary is None:
        raise UncertifiedArchitectureError(
            f"architecture {name!r} is not certified for the boundary adapter; "
            f"certified adapters: {sorted(set(v.adapter for v in _CERTIFIED_ARCHITECTURES.values()))}. "
            "Uncertified architectures stay registered-but-dark until real "
            "certification passes."
        )
    return boundary


def is_certified_architecture(name: Any) -> bool:
    """Return True when ``name`` maps to a certified boundary adapter."""
    try:
        certified_architecture(name)
    except UncertifiedArchitectureError:
        return False
    return True


class ShardRole(str, Enum):
    """Where a contiguous layer range sits in the whole model."""

    HEAD = "head"
    MIDDLE = "middle"
    TAIL = "tail"
    FULL = "full"

    @property
    def owns_embedding(self) -> bool:
        return self in (ShardRole.HEAD, ShardRole.FULL)

    @property
    def owns_final_head(self) -> bool:
        return self in (ShardRole.TAIL, ShardRole.FULL)


def role_for_range(start_layer: int, end_layer: int, total_layers: int) -> ShardRole:
    """Classify a contiguous inclusive layer range within a model of ``total_layers``."""
    if total_layers <= 0:
        raise ValueError("total_layers must be positive")
    if start_layer < 0 or end_layer < start_layer:
        raise ValueError("require 0 <= start_layer <= end_layer")
    if end_layer > total_layers - 1:
        raise ValueError(
            f"end_layer {end_layer} exceeds last layer index {total_layers - 1}"
        )
    is_head = start_layer == 0
    is_tail = end_layer >= total_layers - 1
    if is_head and is_tail:
        return ShardRole.FULL
    if is_head:
        return ShardRole.HEAD
    if is_tail:
        return ShardRole.TAIL
    return ShardRole.MIDDLE


@dataclass(frozen=True)
class BoundaryBundle:
    """The versioned named-tensor bundle handed between adjacent Shard ranges.

    ``residual`` is the *unnormalized* architecture-defined residual stream with
    every position row intact (no tail-only pruning). ``next_layer`` is the layer
    index the receiving range must start at — it is the overlap-safe effective
    start of the seam, so a receiver can reject a bundle meant for a different cut.
    """

    architecture_adapter: str
    schema_version: int
    tensor_name: str
    residual: np.ndarray
    positions: np.ndarray
    next_layer: int
    normalized: bool = False

    def named_tensor_fields(self) -> dict[str, Any]:
        """Return the wire-shaped description of the residual tensor.

        These are exactly the fields the DGR-002 ``NamedTensor`` carries (name,
        shape, dtype, byte order, raw bytes), so a worker can serialize this
        bundle into the gRPC protobuf without re-deriving them.
        """
        residual = np.ascontiguousarray(self.residual)
        return {
            "name": self.tensor_name,
            "shape": list(residual.shape),
            "dtype": residual.dtype.name,
            "byte_order": _byte_order(residual.dtype),
            "data": residual.tobytes(),
        }

    def pack(self) -> dict[str, Any]:
        """Serialize the bundle to a transport-agnostic dict (proves the seam).

        The residual and positions are carried as raw little/big-endian bytes plus
        shape/dtype so that a truly disjoint process can reconstruct the exact
        array — this is what lets two OS processes reproduce whole-model math.
        """
        residual = np.ascontiguousarray(self.residual)
        positions = np.ascontiguousarray(self.positions)
        return {
            "architecture_adapter": self.architecture_adapter,
            "schema_version": self.schema_version,
            "tensor_name": self.tensor_name,
            "next_layer": self.next_layer,
            "normalized": self.normalized,
            "residual": {
                "shape": list(residual.shape),
                "dtype": residual.dtype.str,
                "data": residual.tobytes(),
            },
            "positions": {
                "shape": list(positions.shape),
                "dtype": positions.dtype.str,
                "data": positions.tobytes(),
            },
        }

    @classmethod
    def unpack(cls, payload: dict[str, Any]) -> "BoundaryBundle":
        """Reconstruct a bundle produced by :meth:`pack`."""
        residual = _array_from_wire(payload["residual"])
        positions = _array_from_wire(payload["positions"])
        return cls(
            architecture_adapter=payload["architecture_adapter"],
            schema_version=int(payload["schema_version"]),
            tensor_name=payload["tensor_name"],
            residual=residual,
            positions=positions,
            next_layer=int(payload["next_layer"]),
            normalized=bool(payload.get("normalized", False)),
        )


@dataclass(frozen=True)
class SamplingContract:
    """Explicit contract for turning tail logits into a token.

    The tail never hides the sampling decision inside the adapter: the contract is
    a first-class value so the head/route can reproduce it and so greedy decoding
    is deterministic by construction. Only greedy is certified here; temperature /
    top-p are declared but must be requested explicitly and are out of scope for
    the deterministic parity gate.
    """

    mode: str = "greedy"
    temperature: float = 1.0
    top_p: float = 1.0

    def __post_init__(self) -> None:
        if self.mode not in ("greedy",):
            raise BoundaryContractError(
                f"sampling mode {self.mode!r} is not certified; only 'greedy' is "
                "deterministic and supported by the boundary adapter today"
            )

    @classmethod
    def greedy(cls) -> "SamplingContract":
        return cls(mode="greedy")

    def sample(self, last_logits: np.ndarray) -> int:
        """Return the next token id from the final-position logits row."""
        logits = np.asarray(last_logits)
        if logits.ndim == 2:
            # (batch, vocab) — parity harness uses batch size 1.
            logits = logits[0]
        if logits.ndim != 1:
            raise BoundaryContractError(
                "sampling expects the pruned final-position logits row"
            )
        return int(np.argmax(logits))


@dataclass(frozen=True)
class TailOutput:
    """What a tail Shard emits: the sampled token plus the pruned logits row."""

    token_id: int
    logits: np.ndarray
    sampling: SamplingContract


@dataclass
class BoundaryAdapter:
    """Enforces the architecture-defined boundary contract for one Shard range.

    Construction fails closed for uncertified architectures. The adapter derives
    the Shard's role from its range and drives a duck-typed ``ShardComputation``.
    """

    computation: Any
    sampling: SamplingContract = field(default_factory=SamplingContract.greedy)
    architecture: ArchitectureBoundary = field(init=False)
    role: ShardRole = field(init=False)
    start_layer: int = field(init=False)
    end_layer: int = field(init=False)
    total_layers: int = field(init=False)

    def __post_init__(self) -> None:
        arch_name = getattr(self.computation, "architecture_adapter", None)
        self.architecture = certified_architecture(arch_name)
        self.start_layer = int(getattr(self.computation, "start_layer"))
        self.end_layer = int(getattr(self.computation, "end_layer"))
        self.total_layers = int(getattr(self.computation, "total_layers"))
        self.role = role_for_range(
            self.start_layer, self.end_layer, self.total_layers
        )

    @property
    def is_head(self) -> bool:
        return self.role.owns_embedding

    @property
    def is_tail(self) -> bool:
        return self.role.owns_final_head

    def forward(
        self,
        *,
        token_ids: Any | None = None,
        boundary: BoundaryBundle | None = None,
    ) -> BoundaryBundle | TailOutput:
        """Run one prefill/decode pass for this range and emit its boundary output.

        Head/full ranges require ``token_ids``; middle/tail ranges require the
        ``boundary`` bundle. Non-tail ranges return a :class:`BoundaryBundle`;
        tail/full ranges return a :class:`TailOutput` through the sampling
        contract.
        """
        hidden, positions = self._ingest(token_ids, boundary)
        hidden = self.computation.run_layers(hidden, positions=positions)
        if self.is_tail:
            return self._emit_tail(hidden)
        return self._emit_boundary(hidden, positions)

    # -- input side -----------------------------------------------------------

    def _ingest(
        self, token_ids: Any | None, boundary: BoundaryBundle | None
    ) -> tuple[np.ndarray, np.ndarray]:
        if self.role.owns_embedding:
            return self._ingest_tokens(token_ids, boundary)
        return self._ingest_boundary(token_ids, boundary)

    def _ingest_tokens(
        self, token_ids: Any | None, boundary: BoundaryBundle | None
    ) -> tuple[np.ndarray, np.ndarray]:
        if token_ids is None:
            raise BoundaryContractError(
                "the head owns token embedding and must receive token IDs"
            )
        if boundary is not None:
            raise BoundaryContractError(
                "the head owns token embedding; it must not receive a boundary "
                "bundle from an upstream range"
            )
        ids = np.asarray(token_ids)
        if ids.ndim == 1:
            ids = ids[None, :]
        if ids.ndim != 2:
            raise BoundaryContractError("token IDs must be (seq,) or (batch, seq)")
        hidden = np.asarray(self.computation.embed_tokens(ids))
        positions = np.broadcast_to(
            np.arange(ids.shape[1], dtype=np.int64), ids.shape
        ).copy()
        return hidden, positions

    def _ingest_boundary(
        self, token_ids: Any | None, boundary: BoundaryBundle | None
    ) -> tuple[np.ndarray, np.ndarray]:
        if token_ids is not None:
            raise BoundaryContractError(
                "middle/tail Shards bypass token embedding; they must not receive "
                "token IDs"
            )
        if boundary is None:
            raise BoundaryContractError(
                "middle/tail Shards must receive the named boundary bundle"
            )
        self._check_boundary(boundary)
        return np.asarray(boundary.residual), np.asarray(boundary.positions)

    def _check_boundary(self, boundary: BoundaryBundle) -> None:
        if certified_architecture(boundary.architecture_adapter) is not self.architecture:
            raise BoundaryContractError(
                f"boundary bundle architecture {boundary.architecture_adapter!r} "
                f"does not match this Shard's adapter {self.architecture.adapter!r}"
            )
        if boundary.schema_version != self.architecture.boundary_schema_version:
            raise BoundaryContractError(
                f"boundary schema v{boundary.schema_version} is not supported by "
                f"this Shard (expects v{self.architecture.boundary_schema_version})"
            )
        if boundary.tensor_name != self.architecture.boundary_tensor_name:
            raise BoundaryContractError(
                f"boundary tensor {boundary.tensor_name!r} is not the "
                f"architecture-defined {self.architecture.boundary_tensor_name!r}"
            )
        if boundary.normalized:
            raise BoundaryContractError(
                "boundary bundle is normalized; a Shard range must receive the "
                "UNNORMALIZED architecture-defined residual"
            )
        if boundary.next_layer != self.start_layer:
            raise BoundaryContractError(
                f"boundary hands over at layer {boundary.next_layer} but this "
                f"Shard starts at layer {self.start_layer}"
            )

    # -- output side ----------------------------------------------------------

    def _emit_boundary(
        self, hidden: np.ndarray, positions: np.ndarray
    ) -> BoundaryBundle:
        # A non-tail Shard emits the unnormalized residual with every position row
        # intact: no final norm, no LM head, no tail-only row pruning. next_layer
        # is the receiver's overlap-safe effective start.
        return BoundaryBundle(
            architecture_adapter=self.architecture.adapter,
            schema_version=self.architecture.boundary_schema_version,
            tensor_name=self.architecture.boundary_tensor_name,
            residual=np.asarray(hidden),
            positions=np.asarray(positions),
            next_layer=self.end_layer + 1,
            normalized=False,
        )

    def _emit_tail(self, hidden: np.ndarray) -> TailOutput:
        hidden = np.asarray(hidden)
        # Tail-only row pruning: only the final position is needed to sample the
        # next token, so the LM head runs on the pruned row. A non-tail Shard is
        # forbidden from doing this (it must forward every row).
        if self.architecture.prunes_rows_at_tail:
            last_hidden = hidden[:, -1:, :]
        else:  # pragma: no cover - no certified architecture takes this path yet
            last_hidden = hidden
        if self.architecture.normalizes_before_head:
            last_hidden = np.asarray(self.computation.final_norm(last_hidden))
        logits = np.asarray(self.computation.lm_head(last_hidden))
        last_logits = logits[:, -1, :]
        token_id = self.sampling.sample(last_logits)
        return TailOutput(
            token_id=token_id, logits=last_logits, sampling=self.sampling
        )


def _byte_order(dtype: np.dtype) -> str:
    order = dtype.byteorder
    if order == "<":
        return "little"
    if order == ">":
        return "big"
    # '=' native, '|' not applicable (single byte)
    import sys

    return sys.byteorder if order in ("=", "|") else "little"


def _array_from_wire(field_payload: dict[str, Any]) -> np.ndarray:
    array = np.frombuffer(
        field_payload["data"], dtype=np.dtype(field_payload["dtype"])
    )
    return array.reshape(field_payload["shape"]).copy()