"""Certified architecture adapters for the public TensorBundle boundary.

The adapter is intentionally small: it owns boundary names and endpoint rules,
not transformer execution. llama.cpp owns local graphs; callers select a
certified adapter before accepting an activation from another Shard.
"""

from __future__ import annotations

from dataclasses import dataclass
from enum import Enum
import struct
from typing import Callable, Mapping, Sequence

from .native_protocol import (
    HIDDEN_STATES,
    ProtocolError,
    encode_bundle,
    encode_tensor,
    pb,
    validate_tail_result,
)
from .shard_engine import BoundaryBundle, EngineTensor


# This is deliberately an execution-boundary name, not a transport name.  It
# identifies the value *before* final norm/output projection.  A future wire
# codec may rename its field, but cannot reinterpret this value as logits.
DENSE_LLAMA_ARCHITECTURE = "dense-llama"
DENSE_RESIDUAL_BOUNDARY_V1 = "dense.residual.v1"


class Architecture(str, Enum):
    DENSE = "dense"
    MOE = "moe"
    MLA = "mla"


class BoundaryStage(str, Enum):
    HEAD = "head"
    MIDDLE = "middle"
    TAIL = "tail"


@dataclass(frozen=True)
class ProtocolIdentity:
    request_id: str
    runtime_recipe_digest: str
    chat_template_id: str
    chat_template_version: str
    reasoning_mode: str
    architecture: Architecture


@dataclass(frozen=True)
class SamplingParameters:
    temperature: float
    top_p: float
    top_k: int
    seed: int


@dataclass(frozen=True)
class TailOutput:
    kind: str
    value: int | object

    @classmethod
    def sampled_token(cls, token_id: int) -> "TailOutput":
        if token_id < 0:
            raise ProtocolError("sampled token id must be non-negative")
        return cls("sampled_token", token_id)

    @classmethod
    def logits(cls, logits: object) -> "TailOutput":
        """Return raw logits under the explicit tail-only output contract."""
        return cls("logits", logits)


@dataclass(frozen=True)
class TypedTailResult:
    identity: ProtocolIdentity
    sampling: SamplingParameters
    output_kind: str
    message: pb.TailResult

    @property
    def sampled_token_id(self) -> int | None:
        return self.message.sampled_token_id if self.output_kind == "sampled_token_id" else None


@dataclass(frozen=True)
class ArchitectureBoundaryAdapter:
    architecture: Architecture
    required_names: frozenset[str]

    @property
    def protocol_architecture(self) -> int:
        return {
            Architecture.DENSE: pb.ARCHITECTURE_TYPE_DENSE,
            Architecture.MOE: pb.ARCHITECTURE_TYPE_MOE,
            Architecture.MLA: pb.ARCHITECTURE_TYPE_MLA,
        }[self.architecture]

    def bundle_from_token_ids(
        self,
        token_ids: Sequence[int],
        token_embedding: Callable[[int], Sequence[float]],
    ):
        """Head-only embedding entry point; middle/tail never receive IDs."""
        if self.architecture is not Architecture.DENSE:
            raise ProtocolError("head token embedding is not certified for this architecture")
        if not token_ids:
            raise ProtocolError("head requires at least one token id")
        rows = [tuple(token_embedding(token)) for token in token_ids]
        if not rows or not rows[0] or any(len(row) != len(rows[0]) for row in rows):
            raise ProtocolError("token embedding returned inconsistent hidden widths")
        payload = struct.pack("<" + "f" * (len(rows) * len(rows[0])), *(x for row in rows for x in row))
        return self.bundle_from_named_payloads({HIDDEN_STATES: payload}, shape=[1, len(rows), len(rows[0])])

    def bundle_from_named_payloads(
        self, payloads: Mapping[str, bytes], *, shape: Sequence[int] | None = None
    ):
        names = set(payloads)
        if not self.required_names <= names:
            missing = sorted(self.required_names - names)
            raise ProtocolError(f"{self.architecture.value} boundary requires {missing}")
        tensors = []
        for name, payload in payloads.items():
            tensor_shape = list(shape) if name == HIDDEN_STATES and shape else [len(payload) // 4]
            if len(payload) % 4:
                raise ProtocolError(f"{name!r} F32 fixture payload is not word aligned")
            tensors.append(encode_tensor(name, payload, tensor_shape, pb.DTYPE_FLOAT32))
        return encode_bundle(
            tensors,
            architecture=self.protocol_architecture,
            boundary_point="pre_tail_residual",
        )

    def input_for(self, stage: BoundaryStage, bundle):
        """Accept architecture state only after the head embedding boundary."""
        if stage is BoundaryStage.HEAD:
            raise ProtocolError("head accepts token ids and owns token embedding")
        if bundle is None:
            raise ProtocolError(f"{stage.value} requires a TensorBundle")
        from .native_protocol import decode_bundle

        payloads = decode_bundle(bundle)
        if bundle.architecture != self.protocol_architecture:
            raise ProtocolError("boundary architecture does not match certified adapter")
        if bundle.boundary_point != "pre_tail_residual":
            raise ProtocolError("unsupported architecture boundary point")
        if not self.required_names <= set(payloads):
            raise ProtocolError(f"{self.architecture.value} boundary requires {sorted(self.required_names)}")
        return bundle

    def tail_result(
        self, *, identity: ProtocolIdentity, sampling: SamplingParameters, output: TailOutput
    ) -> TypedTailResult:
        if identity.architecture is not self.architecture:
            raise ProtocolError("tail result architecture does not match certified adapter")
        if not identity.request_id or not identity.runtime_recipe_digest:
            raise ProtocolError("tail result requires exact request and recipe identity")
        if output.kind == "sampled_token":
            if not isinstance(output.value, int):
                raise ProtocolError("sampled tail output must carry an integer token id")
            message = pb.TailResult(
                identity=pb.RequestRecipeIdentity(
                    request_id=identity.request_id,
                    runtime_recipe_digest=identity.runtime_recipe_digest,
                    chat_template_id=identity.chat_template_id,
                    chat_template_version=identity.chat_template_version,
                    reasoning_mode=identity.reasoning_mode,
                    architecture=self.protocol_architecture,
                ),
                sampling=pb.SamplingParameters(
                    temperature=sampling.temperature,
                    top_p=sampling.top_p,
                    top_k=sampling.top_k,
                    seed=sampling.seed,
                    greedy=sampling.temperature == 0.0,
                ),
                sampled_token_id=output.value,
            )
        elif output.kind == "logits":
            if not isinstance(output.value, pb.TensorBundle):
                raise ProtocolError("logits tail output must carry a TensorBundle")
            # Validate the logits bundle before putting it in the result; this
            # rejects an incompatible boundary schema rather than passing an
            # opaque tensor on to sampling.
            from .native_protocol import decode_bundle

            decode_bundle(output.value)
            message = pb.TailResult(
                identity=pb.RequestRecipeIdentity(
                    request_id=identity.request_id,
                    runtime_recipe_digest=identity.runtime_recipe_digest,
                    chat_template_id=identity.chat_template_id,
                    chat_template_version=identity.chat_template_version,
                    reasoning_mode=identity.reasoning_mode,
                    architecture=self.protocol_architecture,
                ),
                sampling=pb.SamplingParameters(
                    temperature=sampling.temperature,
                    top_p=sampling.top_p,
                    top_k=sampling.top_k,
                    seed=sampling.seed,
                    greedy=sampling.temperature == 0.0,
                ),
                logits=output.value,
            )
        else:
            raise ProtocolError("uncertified tail output kind")
        validate_tail_result(message)
        return TypedTailResult(identity, sampling, message.WhichOneof("output"), message)


@dataclass(frozen=True)
class DenseLayerRange:
    """A certified, inclusive dense-Llama range within one loaded model."""

    start_layer: int
    end_layer: int
    total_layers: int
    architecture: str = DENSE_LLAMA_ARCHITECTURE

    def __post_init__(self) -> None:
        if self.architecture != DENSE_LLAMA_ARCHITECTURE:
            raise ProtocolError("dense boundary executor only certifies dense-llama")
        if self.start_layer < 0 or self.end_layer < self.start_layer:
            raise ProtocolError("dense range is empty or inverted")
        if self.total_layers <= self.end_layer:
            raise ProtocolError("dense range lies outside the model")

    @property
    def is_head(self) -> bool:
        return self.start_layer == 0

    @property
    def is_tail(self) -> bool:
        return self.end_layer == self.total_layers - 1


class DenseRangeBoundaryExecutor:
    """Execute one dense range without leaking endpoint ownership.

    ``run_layers`` owns only the local transformer blocks and receives/returns
    the raw residual.  It never receives a final norm/head callback.  Only a
    tail range receives ``tail_output``; consequently row pruning and logits
    projection cannot accidentally happen before the final stage.
    """

    def __init__(
        self,
        layer_range: DenseLayerRange,
        *,
        embed_tokens: Callable[[tuple[int, ...]], EngineTensor],
        run_layers: Callable[[EngineTensor], EngineTensor],
        tail_output: Callable[[EngineTensor], TailOutput] | None = None,
    ) -> None:
        if layer_range.is_tail != (tail_output is not None):
            raise ProtocolError("only a dense tail range may own final norm/output")
        self._range = layer_range
        self._embed_tokens = embed_tokens
        self._run_layers = run_layers
        self._tail_output = tail_output

    def execute(
        self,
        *,
        token_ids: tuple[int, ...] | None = None,
        boundary: BoundaryBundle | None = None,
    ) -> BoundaryBundle | TailOutput:
        if self._range.is_head:
            if token_ids is None or boundary is not None or not token_ids:
                raise ProtocolError("dense head accepts non-empty token ids and no boundary bundle")
            residual = self._embed_tokens(token_ids)
        else:
            if token_ids is not None or boundary is None:
                raise ProtocolError("dense middle/tail requires a named residual boundary bundle")
            residual = self._residual_from_boundary(boundary)

        residual = self._run_layers(residual)
        if residual.name != HIDDEN_STATES:
            raise ProtocolError("dense range must return hidden_states residual")

        if self._range.is_tail:
            assert self._tail_output is not None
            output = self._tail_output(residual)
            if output.kind not in {"logits", "sampled_token"}:
                raise ProtocolError("dense tail returned an uncertified output kind")
            return output

        # Do not normalize, project, sample, or prune rows here: this exact
        # raw output becomes the next range's input.
        return BoundaryBundle(
            tensors=(residual,),
            architecture=DENSE_LLAMA_ARCHITECTURE,
            boundary_point=DENSE_RESIDUAL_BOUNDARY_V1,
        )

    @staticmethod
    def _residual_from_boundary(boundary: BoundaryBundle) -> EngineTensor:
        if boundary.architecture != DENSE_LLAMA_ARCHITECTURE:
            raise ProtocolError("boundary architecture is not certified dense-llama")
        if boundary.boundary_point != DENSE_RESIDUAL_BOUNDARY_V1:
            raise ProtocolError("incompatible dense residual boundary schema")
        if len(boundary.tensors) != 1 or boundary.tensors[0].name != HIDDEN_STATES:
            raise ProtocolError("dense residual boundary requires exactly one hidden_states tensor")
        return boundary.tensors[0]


_ADAPTERS = {
    Architecture.DENSE: ArchitectureBoundaryAdapter(Architecture.DENSE, frozenset({HIDDEN_STATES})),
    Architecture.MOE: ArchitectureBoundaryAdapter(Architecture.MOE, frozenset({HIDDEN_STATES, "router_logits"})),
    Architecture.MLA: ArchitectureBoundaryAdapter(Architecture.MLA, frozenset({HIDDEN_STATES, "mla_position_state"})),
}


def adapter_for(architecture: Architecture | str) -> ArchitectureBoundaryAdapter:
    try:
        return _ADAPTERS[Architecture(architecture)]
    except (KeyError, ValueError):
        raise ProtocolError(f"unsupported architecture {architecture!r}") from None