"""Isolated concurrent local Hot KV State for distributed Shards (DGR-007).

Hot KV State stays local to the node serving a Shard (RALPH runtime decision #7).
A concurrent server must map each ``(Route Session ID, route epoch)`` to an
isolated bounded KV context (decision #8) so that one request can never clear or
corrupt another's cache.

This module owns the *lifecycle and storage* of that state and is deliberately
backend-agnostic:

* :class:`HotKvStateManager` is the single mutation entry point. It maps
  ``(session_id, route_epoch)`` to a :class:`SessionCache`, allocates KV **only
  for the owned layer range**, and enforces a byte budget, a session cap, and a
  TTL through LRU/TTL eviction. It rejects stale route epochs and incompatible
  cache recipes, and returns an **explicit** :class:`CacheMiss` when state the
  caller expected is gone (evicted, released, desynchronised, or never held) so
  the head degrades to a from-token-zero re-prefill instead of corrupting output
  (RALPH decision #14: unverified KV is never migrated silently).
* :class:`LayerKvCache` / :class:`SessionCache` are the per-owned-layer K/V
  containers. They are plain ``numpy`` arrays so the default deterministic test
  suite needs no torch, GPU, download, or API credit; the pinned llama.cpp worker
  (DGR-008) maps a llama sequence onto the same container contract.
* :class:`KvBoundaryAdapter` wraps a KV-aware ``ShardComputation`` (the DGR-006
  duck type plus ``run_layers_cached``) so a Shard can run cached prefill/decode
  through the manager while honouring the architecture-defined boundary contract
  (head embeds tokens, middle/tail bypass embedding, non-tail emits the
  unnormalized residual, tail samples).

The manager owns *all* cache mutation: a computation reads the existing cache and
returns the new K/V for the appended positions, and the manager decides whether
that append fits the budget. That keeps eviction, accounting, and isolation in one
place instead of scattered across backends.
"""

from __future__ import annotations

import threading
import time
from collections import OrderedDict
from dataclasses import dataclass, field
from enum import Enum
from typing import Any, Callable, Mapping

import numpy as np

from meshnet_node.boundary_adapter import (
    BOUNDARY_SCHEMA_VERSION,
    BoundaryBundle,
    BoundaryContractError,
    SamplingContract,
    ShardRole,
    TailOutput,
    certified_architecture,
    role_for_range,
)
from meshnet_node.runtime_recipe import compatibility_fingerprint


class HotKvStateError(RuntimeError):
    """Base class for Hot KV State errors."""


class StaleRouteEpochError(HotKvStateError):
    """Raised when a request references a route epoch older than the current one.

    A newer route epoch means the route was re-planned; the old epoch's KV is
    unverified against the new plan and must never be silently reused.
    """


class IncompatibleCacheRecipeError(HotKvStateError):
    """Raised when a request's cache recipe does not match the loaded shard.

    A different quantization / dtype / owned range / architecture produces a KV
    layout this node cannot reuse without corrupting output.
    """


class KvBudgetExceededError(HotKvStateError):
    """Raised when a single session cannot fit the configured byte budget.

    Other sessions are evicted first (LRU); this fires only when even one session
    alone exceeds the budget, which is a misconfiguration rather than pressure.
    """


class KvCacheMissError(HotKvStateError):
    """Raised by the strict accessor when expected session state is absent.

    Prefer :meth:`HotKvStateManager.resolve`, which returns a structured
    :class:`CacheMiss` instead of raising, when the caller wants to fall back to a
    stateless re-prefill.
    """

    def __init__(self, miss: "CacheMiss") -> None:
        super().__init__(str(miss))
        self.miss = miss


class CacheMissReason(str, Enum):
    """Why a lookup produced a cache miss (all benign; retry from token zero)."""

    UNKNOWN_SESSION = "unknown-session"
    EVICTED_TTL = "evicted-ttl"
    EVICTED_LRU = "evicted-lru"
    RELEASED = "released"
    SUPERSEDED_EPOCH = "superseded-epoch"
    SEQ_LEN_MISMATCH = "seq-len-mismatch"


@dataclass(frozen=True)
class CacheMiss:
    """Explicit cache-miss response the head can act on (re-prefill).

    This is a value, not an exception: the native protocol carries a cache
    expectation/result, and a miss is a normal, expected outcome under eviction.
    """

    session_id: str
    route_epoch: int
    reason: CacheMissReason
    detail: str = ""

    def __str__(self) -> str:
        suffix = f": {self.detail}" if self.detail else ""
        return (
            f"cache miss for session {self.session_id[:8]} epoch "
            f"{self.route_epoch} ({self.reason.value}){suffix}"
        )


@dataclass(frozen=True)
class KvCacheRecipe:
    """The identity of a Shard's KV layout, used to reject incompatible reuse.

    Two recipes are compatible iff their fingerprints match — same certified
    architecture, KV dtype, head geometry, and owned layer range within the same
    whole-model layer count.
    """

    architecture_adapter: str
    kv_dtype: str
    n_kv_heads: int
    head_dim: int
    total_layers: int
    start_layer: int
    end_layer: int
    boundary_schema_version: int = BOUNDARY_SCHEMA_VERSION

    def __post_init__(self) -> None:
        # Fail closed on architecture identity (shared with the boundary adapter).
        certified_architecture(self.architecture_adapter)
        if self.n_kv_heads <= 0:
            raise ValueError("n_kv_heads must be positive")
        if self.head_dim <= 0:
            raise ValueError("head_dim must be positive")
        try:
            np.dtype(self.kv_dtype)
        except TypeError as exc:  # pragma: no cover - defensive
            raise ValueError(f"invalid kv_dtype {self.kv_dtype!r}") from exc
        # role_for_range validates 0 <= start <= end <= total_layers - 1.
        role_for_range(self.start_layer, self.end_layer, self.total_layers)
        if self.boundary_schema_version < 1:
            raise ValueError("boundary_schema_version must be >= 1")

    @property
    def owned_layers(self) -> tuple[int, ...]:
        return tuple(range(self.start_layer, self.end_layer + 1))

    @property
    def role(self) -> ShardRole:
        return role_for_range(self.start_layer, self.end_layer, self.total_layers)

    def bytes_per_token(self) -> int:
        """Bytes of KV one token adds across *owned* layers (keys + values)."""
        itemsize = np.dtype(self.kv_dtype).itemsize
        per_layer = 2 * self.n_kv_heads * self.head_dim * itemsize
        return per_layer * len(self.owned_layers)

    def fingerprint(self) -> str:
        return compatibility_fingerprint(
            {
                "kind": "hot-kv-recipe",
                # Canonicalize the architecture so 'llama' / 'LlamaForCausalLM'
                # map to the same fingerprint (they are the same layout).
                "architecture_adapter": certified_architecture(
                    self.architecture_adapter
                ).adapter,
                "kv_dtype": np.dtype(self.kv_dtype).name,
                "n_kv_heads": self.n_kv_heads,
                "head_dim": self.head_dim,
                "total_layers": self.total_layers,
                "start_layer": self.start_layer,
                "end_layer": self.end_layer,
                "boundary_schema_version": self.boundary_schema_version,
            }
        )

    def is_compatible(self, other: "KvCacheRecipe") -> bool:
        return self.fingerprint() == other.fingerprint()


class LayerKvCache:
    """K/V storage for a single owned layer; sequence axis is 0.

    Keys and values are ``(seq, n_kv_heads, head_dim)``. Backends store the
    position-encoded (post-RoPE) keys so a decode step only appends the new rows.
    """

    __slots__ = ("layer_index", "n_kv_heads", "head_dim", "dtype", "keys", "values")

    def __init__(
        self, layer_index: int, n_kv_heads: int, head_dim: int, dtype: Any
    ) -> None:
        self.layer_index = int(layer_index)
        self.n_kv_heads = int(n_kv_heads)
        self.head_dim = int(head_dim)
        self.dtype = np.dtype(dtype)
        self.keys = np.empty((0, self.n_kv_heads, self.head_dim), dtype=self.dtype)
        self.values = np.empty((0, self.n_kv_heads, self.head_dim), dtype=self.dtype)

    @property
    def length(self) -> int:
        return int(self.keys.shape[0])

    def _validate(self, array: np.ndarray, name: str) -> np.ndarray:
        arr = np.asarray(array, dtype=self.dtype)
        if arr.ndim != 3 or arr.shape[1:] != (self.n_kv_heads, self.head_dim):
            raise ValueError(
                f"layer {self.layer_index} {name} must be "
                f"(seq, {self.n_kv_heads}, {self.head_dim}), got {arr.shape}"
            )
        return arr

    def append(self, keys: np.ndarray, values: np.ndarray) -> int:
        k = self._validate(keys, "keys")
        v = self._validate(values, "values")
        if k.shape[0] != v.shape[0]:
            raise ValueError(
                f"layer {self.layer_index} keys/values disagree on token count "
                f"({k.shape[0]} vs {v.shape[0]})"
            )
        self.keys = np.concatenate([self.keys, k], axis=0)
        self.values = np.concatenate([self.values, v], axis=0)
        return self.length

    def truncate(self, length: int) -> None:
        length = max(0, int(length))
        self.keys = self.keys[:length]
        self.values = self.values[:length]

    @property
    def nbytes(self) -> int:
        return int(self.keys.nbytes + self.values.nbytes)


@dataclass
class SessionCache:
    """Isolated per-``(session_id, epoch)`` KV context over the owned layers only."""

    session_id: str
    route_epoch: int
    recipe: KvCacheRecipe
    layers: "OrderedDict[int, LayerKvCache]"
    created_tick: float
    last_tick: float
    released: bool = False

    @property
    def seq_len(self) -> int:
        if not self.layers:
            return 0
        # All owned layers advance in lockstep; report the first owned layer.
        return next(iter(self.layers.values())).length

    @property
    def owned_layers(self) -> tuple[int, ...]:
        return tuple(self.layers.keys())

    def layer(self, index: int) -> LayerKvCache:
        try:
            return self.layers[index]
        except KeyError:
            raise KeyError(
                f"layer {index} is not owned by this shard "
                f"(owned {list(self.layers)})"
            ) from None

    def read_only_layers(self) -> Mapping[int, LayerKvCache]:
        """The current per-layer caches a computation reads to attend over."""
        return dict(self.layers)

    def _append(self, kv_by_layer: Mapping[int, Any]) -> int:
        provided = set(kv_by_layer)
        owned = set(self.layers)
        if provided != owned:
            raise ValueError(
                f"append must cover exactly the owned layers {sorted(owned)}, "
                f"got {sorted(provided)}"
            )
        # Pre-validate token counts so a partial append never desynchronises the
        # owned layers (append is all-or-nothing).
        new_counts = set()
        for idx, (keys, _values) in kv_by_layer.items():
            new_counts.add(int(np.asarray(keys).shape[0]))
        if len(new_counts) != 1:
            raise ValueError(
                f"append token counts disagree across layers: {sorted(new_counts)}"
            )
        for idx, (keys, values) in kv_by_layer.items():
            self.layers[idx].append(keys, values)
        return self.seq_len

    def _truncate(self, length: int) -> None:
        for cache in self.layers.values():
            cache.truncate(length)

    @property
    def nbytes(self) -> int:
        return sum(cache.nbytes for cache in self.layers.values())


@dataclass(frozen=True)
class HotKvStateConfig:
    """Bounds for the manager: memory budget, session cap, and idle TTL."""

    budget_bytes: int = 64 * 1024 * 1024
    max_sessions: int = 8
    ttl_seconds: float = 600.0
    miss_history: int = 256

    def __post_init__(self) -> None:
        if self.budget_bytes <= 0:
            raise ValueError("budget_bytes must be positive")
        if self.max_sessions < 1:
            raise ValueError("max_sessions must be >= 1")
        if self.ttl_seconds < 0:
            raise ValueError("ttl_seconds must be >= 0")
        if self.miss_history < 0:
            raise ValueError("miss_history must be >= 0")


class HotKvStateManager:
    """Concurrent, bounded map of ``(session_id, epoch)`` to an isolated KV context."""

    def __init__(
        self,
        recipe: KvCacheRecipe,
        config: HotKvStateConfig | None = None,
        *,
        clock: Callable[[], float] | None = None,
    ) -> None:
        self.recipe = recipe
        self.config = config or HotKvStateConfig()
        self._clock = clock or time.monotonic
        self._sessions: "OrderedDict[tuple[str, int], SessionCache]" = OrderedDict()
        self._latest_epoch: dict[str, int] = {}
        self._misses: "OrderedDict[tuple[str, int], CacheMiss]" = OrderedDict()
        self._lock = threading.RLock()

    # -- introspection --------------------------------------------------------

    @property
    def total_bytes(self) -> int:
        with self._lock:
            return sum(s.nbytes for s in self._sessions.values())

    @property
    def session_count(self) -> int:
        with self._lock:
            self._evict_expired_locked(self._clock())
            return len(self._sessions)

    def session_keys(self) -> list[tuple[str, int]]:
        with self._lock:
            return list(self._sessions.keys())

    # -- lifecycle ------------------------------------------------------------

    def open(
        self,
        session_id: str,
        route_epoch: int,
        *,
        recipe: KvCacheRecipe | None = None,
    ) -> SessionCache:
        """Create (or replace) a fresh, empty isolated context for the session.

        A higher route epoch supersedes and frees any earlier epoch for the same
        session id; an older epoch is rejected as stale.
        """
        self._require_text(session_id, "session_id")
        route_epoch = self._require_epoch(route_epoch)
        with self._lock:
            self._check_recipe(recipe)
            self._validate_epoch_locked(session_id, route_epoch)
            now = self._clock()
            self._evict_expired_locked(now)
            self._supersede_older_epochs_locked(session_id, route_epoch)
            key = (session_id, route_epoch)
            # A re-open at the same epoch replaces the prior context entirely.
            self._sessions.pop(key, None)
            layers: "OrderedDict[int, LayerKvCache]" = OrderedDict(
                (
                    idx,
                    LayerKvCache(
                        idx,
                        self.recipe.n_kv_heads,
                        self.recipe.head_dim,
                        self.recipe.kv_dtype,
                    ),
                )
                for idx in self.recipe.owned_layers
            )
            session = SessionCache(
                session_id=session_id,
                route_epoch=route_epoch,
                recipe=self.recipe,
                layers=layers,
                created_tick=now,
                last_tick=now,
            )
            self._sessions[key] = session
            self._latest_epoch[session_id] = route_epoch
            self._misses.pop(key, None)
            self._enforce_capacity_locked(protect=key, incoming_bytes=0)
            return session

    def append(
        self,
        session_id: str,
        route_epoch: int,
        kv_by_layer: Mapping[int, Any],
        *,
        recipe: KvCacheRecipe | None = None,
        expected_seq_len: int | None = None,
    ) -> SessionCache:
        """Append new K/V (prefill or decode) to an existing isolated context.

        The computation supplies exactly the owned layers' new keys/values. The
        manager evicts other sessions (LRU) to fit the byte budget before growing
        this one, and raises :class:`KvBudgetExceededError` only if this session
        alone cannot fit.
        """
        route_epoch = self._require_epoch(route_epoch)
        with self._lock:
            self._check_recipe(recipe)
            self._validate_epoch_locked(session_id, route_epoch)
            session = self._require_live_locked(session_id, route_epoch)
            if expected_seq_len is not None and session.seq_len != expected_seq_len:
                miss = self._drop_and_record_locked(
                    (session_id, route_epoch),
                    CacheMissReason.SEQ_LEN_MISMATCH,
                    detail=f"cache holds {session.seq_len}, caller expected "
                    f"{expected_seq_len}",
                )
                raise KvCacheMissError(miss)
            n_new = self._new_token_count(kv_by_layer)
            incoming = n_new * self.recipe.bytes_per_token()
            self._enforce_capacity_locked(
                protect=(session_id, route_epoch), incoming_bytes=incoming
            )
            session._append(kv_by_layer)
            session.last_tick = self._clock()
            self._sessions.move_to_end((session_id, route_epoch))
            return session

    def truncate(
        self, session_id: str, route_epoch: int, length: int
    ) -> SessionCache:
        """Drop cached positions beyond ``length`` (rollback) for one session."""
        route_epoch = self._require_epoch(route_epoch)
        with self._lock:
            self._validate_epoch_locked(session_id, route_epoch)
            session = self._require_live_locked(session_id, route_epoch)
            if length < 0:
                raise ValueError("truncate length must be >= 0")
            session._truncate(length)
            session.last_tick = self._clock()
            self._sessions.move_to_end((session_id, route_epoch))
            return session

    def release(self, session_id: str, route_epoch: int) -> bool:
        """Free one session's context; other sessions are untouched.

        Returns True if a live context was freed. A later lookup for the released
        key yields an explicit :class:`CacheMiss`.
        """
        route_epoch = self._require_epoch(route_epoch)
        with self._lock:
            key = (session_id, route_epoch)
            existed = key in self._sessions
            self._drop_and_record_locked(key, CacheMissReason.RELEASED)
            return existed

    # -- lookup ---------------------------------------------------------------

    def resolve(
        self,
        session_id: str,
        route_epoch: int,
        *,
        recipe: KvCacheRecipe | None = None,
        expected_seq_len: int | None = None,
    ) -> SessionCache | CacheMiss:
        """Return the live context or an explicit :class:`CacheMiss`.

        Rejects stale epochs and incompatible recipes (both are protocol
        violations, not benign misses).
        """
        route_epoch = self._require_epoch(route_epoch)
        with self._lock:
            self._check_recipe(recipe)
            self._validate_epoch_locked(session_id, route_epoch)
            now = self._clock()
            self._evict_expired_locked(now)
            key = (session_id, route_epoch)
            session = self._sessions.get(key)
            if session is None:
                return self._recorded_miss_locked(key)
            if expected_seq_len is not None and session.seq_len != expected_seq_len:
                return self._drop_and_record_locked(
                    key,
                    CacheMissReason.SEQ_LEN_MISMATCH,
                    detail=f"cache holds {session.seq_len}, caller expected "
                    f"{expected_seq_len}",
                )
            session.last_tick = now
            self._sessions.move_to_end(key)
            return session

    def get(
        self,
        session_id: str,
        route_epoch: int,
        *,
        recipe: KvCacheRecipe | None = None,
        expected_seq_len: int | None = None,
    ) -> SessionCache:
        """Strict accessor: raises :class:`KvCacheMissError` on a miss."""
        result = self.resolve(
            session_id,
            route_epoch,
            recipe=recipe,
            expected_seq_len=expected_seq_len,
        )
        if isinstance(result, CacheMiss):
            raise KvCacheMissError(result)
        return result

    # -- internals ------------------------------------------------------------

    def _check_recipe(self, recipe: KvCacheRecipe | None) -> None:
        if recipe is not None and not self.recipe.is_compatible(recipe):
            raise IncompatibleCacheRecipeError(
                "request cache recipe does not match this shard's loaded recipe "
                f"(request {recipe.fingerprint()} vs shard {self.recipe.fingerprint()})"
            )

    def _validate_epoch_locked(self, session_id: str, route_epoch: int) -> None:
        latest = self._latest_epoch.get(session_id)
        if latest is not None and route_epoch < latest:
            raise StaleRouteEpochError(
                f"session {session_id[:8]} route epoch {route_epoch} is stale; "
                f"current epoch is {latest}"
            )

    def _supersede_older_epochs_locked(
        self, session_id: str, route_epoch: int
    ) -> None:
        stale_keys = [
            key
            for key in self._sessions
            if key[0] == session_id and key[1] < route_epoch
        ]
        for key in stale_keys:
            self._drop_and_record_locked(key, CacheMissReason.SUPERSEDED_EPOCH)

    def _require_live_locked(
        self, session_id: str, route_epoch: int
    ) -> SessionCache:
        now = self._clock()
        self._evict_expired_locked(now)
        key = (session_id, route_epoch)
        session = self._sessions.get(key)
        if session is None:
            raise KvCacheMissError(self._recorded_miss_locked(key))
        return session

    def _new_token_count(self, kv_by_layer: Mapping[int, Any]) -> int:
        owned = set(self.recipe.owned_layers)
        if set(kv_by_layer) != owned:
            raise ValueError(
                f"append must cover exactly the owned layers {sorted(owned)}, "
                f"got {sorted(kv_by_layer)}"
            )
        counts = {int(np.asarray(k).shape[0]) for k, _ in kv_by_layer.values()}
        if len(counts) != 1:
            raise ValueError(
                f"append token counts disagree across layers: {sorted(counts)}"
            )
        return counts.pop()

    def _enforce_capacity_locked(
        self, *, protect: tuple[str, int], incoming_bytes: int
    ) -> None:
        # Session cap: evict LRU sessions other than the protected one.
        while len(self._sessions) > self.config.max_sessions:
            victim = self._lru_victim_locked(protect)
            if victim is None:
                break
            self._drop_and_record_locked(victim, CacheMissReason.EVICTED_LRU)

        # Byte budget: the protected session's own footprint after the append.
        protected = self._sessions.get(protect)
        protected_bytes = (protected.nbytes if protected is not None else 0) + incoming_bytes
        if protected_bytes > self.config.budget_bytes:
            raise KvBudgetExceededError(
                f"session {protect[0][:8]} needs {protected_bytes} bytes which "
                f"exceeds the KV budget {self.config.budget_bytes}"
            )
        # Evict other LRU sessions until the whole store fits with the append.
        while self._total_bytes_locked() + incoming_bytes > self.config.budget_bytes:
            victim = self._lru_victim_locked(protect)
            if victim is None:
                break
            self._drop_and_record_locked(victim, CacheMissReason.EVICTED_LRU)

    def _lru_victim_locked(self, protect: tuple[str, int]) -> tuple[str, int] | None:
        for key in self._sessions:  # OrderedDict iterates oldest-first.
            if key != protect:
                return key
        return None

    def _total_bytes_locked(self) -> int:
        return sum(s.nbytes for s in self._sessions.values())

    def _evict_expired_locked(self, now: float) -> None:
        ttl = self.config.ttl_seconds
        if ttl <= 0:
            return
        expired = [
            key
            for key, session in self._sessions.items()
            if now - session.last_tick > ttl
        ]
        for key in expired:
            self._drop_and_record_locked(key, CacheMissReason.EVICTED_TTL)

    def _drop_and_record_locked(
        self,
        key: tuple[str, int],
        reason: CacheMissReason,
        *,
        detail: str = "",
    ) -> CacheMiss:
        session = self._sessions.pop(key, None)
        if session is not None:
            session.released = True
        miss = CacheMiss(
            session_id=key[0], route_epoch=key[1], reason=reason, detail=detail
        )
        self._record_miss_locked(key, miss)
        return miss

    def _record_miss_locked(self, key: tuple[str, int], miss: CacheMiss) -> None:
        if self.config.miss_history <= 0:
            return
        self._misses.pop(key, None)
        self._misses[key] = miss
        while len(self._misses) > self.config.miss_history:
            self._misses.popitem(last=False)

    def _recorded_miss_locked(self, key: tuple[str, int]) -> CacheMiss:
        recorded = self._misses.get(key)
        if recorded is not None:
            return recorded
        return CacheMiss(
            session_id=key[0],
            route_epoch=key[1],
            reason=CacheMissReason.UNKNOWN_SESSION,
        )

    @staticmethod
    def _require_text(value: Any, name: str) -> str:
        if not isinstance(value, str) or not value.strip():
            raise ValueError(f"{name} must be a non-empty string")
        return value

    @staticmethod
    def _require_epoch(value: Any) -> int:
        if isinstance(value, bool) or not isinstance(value, int):
            raise ValueError("route_epoch must be an integer")
        if value < 0:
            raise ValueError("route_epoch must be >= 0")
        return value


def kv_recipe_for(computation: Any) -> KvCacheRecipe:
    """Build a :class:`KvCacheRecipe` from a KV-aware ``ShardComputation``.

    The computation exposes the DGR-006 duck type plus KV geometry
    (``n_kv_heads``, ``head_dim``, ``kv_dtype``).
    """
    return KvCacheRecipe(
        architecture_adapter=str(getattr(computation, "architecture_adapter")),
        kv_dtype=str(getattr(computation, "kv_dtype", "float32")),
        n_kv_heads=int(getattr(computation, "n_kv_heads")),
        head_dim=int(getattr(computation, "head_dim")),
        total_layers=int(getattr(computation, "total_layers")),
        start_layer=int(getattr(computation, "start_layer")),
        end_layer=int(getattr(computation, "end_layer")),
    )


@dataclass
class KvBoundaryAdapter:
    """KV-aware boundary driver: cached prefill/decode through the manager.

    Mirrors the DGR-006 :class:`~meshnet_node.boundary_adapter.BoundaryAdapter`
    contract (head embeds tokens, middle/tail bypass embedding and consume the
    unnormalized residual bundle, non-tail emits the unnormalized residual, tail
    normalizes + heads + prunes + samples) but threads a per-session KV context.

    The wrapped computation must additionally expose::

        run_layers_cached(hidden, *, positions, past_kv)
            -> (hidden_out, {layer_index: (new_keys, new_values)})

    reading ``past_kv`` (the current per-owned-layer caches) and returning the new
    position-encoded K/V for the appended positions only. The manager, not the
    computation, commits those K/V so eviction and budget stay centralized.
    """

    computation: Any
    manager: HotKvStateManager
    sampling: SamplingContract = field(default_factory=SamplingContract.greedy)
    architecture: Any = 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)
    recipe: KvCacheRecipe = 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)
        self.recipe = kv_recipe_for(self.computation)
        if not self.manager.recipe.is_compatible(self.recipe):
            raise IncompatibleCacheRecipeError(
                "manager recipe does not match this computation's KV recipe"
            )

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

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

    def prefill(
        self,
        session_id: str,
        route_epoch: int,
        *,
        token_ids: Any | None = None,
        boundary: BoundaryBundle | None = None,
    ) -> BoundaryBundle | TailOutput:
        """Open a fresh isolated context and run the prompt through this range."""
        session = self.manager.open(session_id, route_epoch, recipe=self.recipe)
        return self._run_step(session, token_ids, boundary)

    def decode(
        self,
        session_id: str,
        route_epoch: int,
        *,
        token_ids: Any | None = None,
        boundary: BoundaryBundle | None = None,
        expected_seq_len: int | None = None,
    ) -> BoundaryBundle | TailOutput | CacheMiss:
        """Append one (or more) decode positions to an existing context.

        Returns an explicit :class:`CacheMiss` if the context is gone so the head
        can re-prefill from token zero instead of corrupting output.
        """
        resolved = self.manager.resolve(
            session_id,
            route_epoch,
            recipe=self.recipe,
            expected_seq_len=expected_seq_len,
        )
        if isinstance(resolved, CacheMiss):
            return resolved
        return self._run_step(resolved, token_ids, boundary)

    # -- internals ------------------------------------------------------------

    def _run_step(
        self,
        session: SessionCache,
        token_ids: Any | None,
        boundary: BoundaryBundle | None,
    ) -> BoundaryBundle | TailOutput:
        prev_len = session.seq_len
        hidden, positions = self._ingest(prev_len, token_ids, boundary)
        hidden_out, new_kv = self.computation.run_layers_cached(
            hidden, positions=positions, past_kv=session.read_only_layers()
        )
        self.manager.append(
            session.session_id,
            session.route_epoch,
            new_kv,
            recipe=self.recipe,
            expected_seq_len=prev_len,
        )
        if self.is_tail:
            return self._emit_tail(hidden_out)
        return self._emit_boundary(hidden_out, positions)

    def _ingest(
        self,
        prev_len: int,
        token_ids: Any | None,
        boundary: BoundaryBundle | None,
    ) -> tuple[np.ndarray, np.ndarray]:
        if self.role.owns_embedding:
            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))
            n_new = ids.shape[1]
            positions = np.broadcast_to(
                np.arange(prev_len, prev_len + n_new, dtype=np.int64),
                ids.shape,
            ).copy()
            return hidden, positions
        # Middle / tail: consume the boundary bundle (the unnormalized residual).
        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}"
            )

    def _emit_boundary(
        self, hidden: np.ndarray, positions: np.ndarray
    ) -> BoundaryBundle:
        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)
        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)