"""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


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


@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)


@dataclass(frozen=True)
class TypedTailResult:
    identity: ProtocolIdentity
    sampling: SamplingParameters
    output_kind: str
    sampled_token_id: int | None = None
    logits: object | None = None


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

    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)

    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 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 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":
            raise ProtocolError("uncertified tail output kind")
        return TypedTailResult(identity, sampling, output.kind, sampled_token_id=int(output.value))


_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