"""MFP Runtime — the central stateful engine.

Owns all protocol state: agent table, channel registry, global ratchet
state Sg. Orchestrates the message lifecycle by composing pure core
functions with the pipeline, channel, and quarantine modules.

Maps to: impl/I-06_runtime.md
"""

from __future__ import annotations

import threading
import time
from dataclasses import dataclass, field

from mfp.agent.identity import generate_agent_id
from mfp.core.merkle import IncrementalSg
from mfp.core.primitives import random_bytes, sha256
from mfp.core.ratchet import compose_ordered
from mfp.core.types import (
    AgentError,
    AgentErrorCode,
    AgentId,
    AgentState,
    AgentStatus,
    Channel,
    ChannelId,
    ChannelInfo,
    ChannelState,
    ChannelStatus,
    FrameValidationError,
    GlobalState,
    BilateralState,
    PROTOCOL_VERSION,
    Receipt,
    RuntimeId,
    StateValue,
    DEFAULT_FRAME_DEPTH,
)

from mfp.runtime.channels import (
    ChannelRegistry,
    advance_channel,
    close_channel as ch_close,
    establish_channel as ch_establish,
    get_channel,
    get_channels_for_agent,
)
from mfp.runtime.deduplication import DeduplicationTracker
from mfp.runtime.pipeline import (
    AgentCallable,
    RuntimeConfig,
    process_message,
)
from mfp.federation.bilateral import (
    BilateralChannel,
    CrossRuntimeChannel,
    bootstrap_deterministic,
    bootstrap_ceremonial,
    advance_bilateral_state,
    derive_bilateral_id,
)
from mfp.federation.transport import (
    ConnectionPool,
    TransportConfig,
    TransportServer,
)
from mfp.federation.wire import (
    parse_wire_message,
    validate_envelope,
)
from mfp.observability.logging import AuditEventType, LogContext, get_logger, log_audit_event
from mfp.observability.metrics import get_metrics_collector
from mfp.storage.engine import RuntimeMeta, StorageEngine
from mfp.runtime.quarantine import (
    check_rate_limit,
    check_rate_limit_window,
    check_validation_failure,
    increment_failure_count,
    quarantine_agent as q_quarantine_agent,
    quarantine_channel as q_quarantine_channel,
    reset_failure_count,
    restore_agent as q_restore_agent,
    restore_channel as q_restore_channel,
)


# ---------------------------------------------------------------------------
# Internal Types
# ---------------------------------------------------------------------------

@dataclass
class AgentRecord:
    """Internal agent record. Not exposed to agents.

    Maps to: I-06 §6.1.
    """
    agent_id: AgentId
    state: AgentState
    callable: AgentCallable  # AgentCallable or AsyncAgentCallable
    is_async: bool = False
    channels: set[bytes] = field(default_factory=set)
    message_count: int = 0  # Kept for backward compat / metrics
    _message_timestamps: list[float] = field(default_factory=list)
    quarantine_reason: str = ""


# ---------------------------------------------------------------------------
# Runtime
# ---------------------------------------------------------------------------

class Runtime:
    """MFP Runtime — the central stateful engine.

    Owns all protocol state. Orchestrates the message lifecycle.
    Composes pure core functions into a working engine.

    Maps to: runtime-interface.md §1, §9.
    """

    def __init__(self, config: RuntimeConfig | None = None) -> None:
        self._config = config or RuntimeConfig()
        self._identity = self._derive_identity()
        self._agents: dict[bytes, AgentRecord] = {}
        self._channels: ChannelRegistry = {}
        self._sg: GlobalState | None = None
        self._incremental_sg: IncrementalSg | None = None  # Merkle tree for O(log N) Sg
        self._agent_counter: int = 0
        self._deduplication = DeduplicationTracker()
        self._logger = get_logger(__name__)

        # Per-channel ordering locks (spec: runtime-interface.md §4.9)
        self._channel_locks: dict[bytes, threading.Lock] = {}
        self._sg_lock = threading.Lock()

        # Federation
        self._bilateral_channels: dict[bytes, BilateralChannel] = {}
        self._cross_runtime_channels: dict[bytes, CrossRuntimeChannel] = {}
        self._transport_server: TransportServer | None = None
        self._connection_pool: ConnectionPool | None = None
        self._pending_wire_messages: list[tuple[bytes, bytes]] = []

        # Persistence (None = in-memory only)
        self._storage: StorageEngine | None = None
        if self._config.storage is not None:
            self._storage = StorageEngine(self._config.storage)
            self._persist_meta()
        self._storage_degraded: bool = False

        # Log runtime initialization
        context = LogContext(
            correlation_id="init",
            runtime_id=self._identity.data.hex()[:8],
            operation="runtime_init",
        )
        self._logger.info("Runtime initialized", context=context)

    def _get_channel_lock(self, channel_id_bytes: bytes) -> threading.Lock:
        """Get or create a per-channel lock for message ordering."""
        if channel_id_bytes not in self._channel_locks:
            self._channel_locks[channel_id_bytes] = threading.Lock()
        return self._channel_locks[channel_id_bytes]

    @classmethod
    def recover(
        cls,
        config: RuntimeConfig,
        agent_callables: dict[bytes, AgentCallable] | None = None,
    ) -> Runtime:
        """Recover a Runtime from persistent storage.

        Restores agents, channels, Merkle tree, and Sg from the database.
        Agent callables must be supplied externally (they cannot be serialized).
        Agents whose callables are not provided remain in BOUND state.

        Args:
            config: RuntimeConfig with storage configured.
            agent_callables: Mapping of agent_id bytes -> callable.
                             Agents not in this dict stay BOUND.

        Returns:
            A fully reconstructed Runtime.

        Raises:
            ValueError: If config.storage is None or DB is empty.
        """
        if config.storage is None:
            raise ValueError("Cannot recover without storage configuration")

        callables = agent_callables or {}
        storage = StorageEngine(config.storage)
        result = storage.recover()
        if result is None:
            storage.close()
            raise ValueError("No state to recover — database is empty")

        # Build runtime with the same identity
        rt = cls.__new__(cls)
        rt._config = config
        rt._identity = StateValue(result.meta.runtime_id)
        rt._agents = {}
        rt._channels = {}
        rt._sg = None
        rt._incremental_sg = None
        rt._agent_counter = result.meta.agent_counter
        rt._deduplication = DeduplicationTracker()
        rt._logger = get_logger(__name__)
        rt._channel_locks = {}
        rt._sg_lock = threading.Lock()
        rt._bilateral_channels = {}
        rt._cross_runtime_channels = {}
        rt._transport_server = None
        rt._connection_pool = None
        rt._pending_wire_messages = []
        rt._storage = storage
        rt._storage_degraded = False

        # Reconstruct channels
        status_map = {
            "active": ChannelStatus.ACTIVE,
            "quarantined": ChannelStatus.QUARANTINED,
            "closed": ChannelStatus.CLOSED,
        }
        for ch_row in result.channels:
            if ch_row.status == "closed":
                continue
            channel = Channel(
                channel_id=ChannelId(ch_row.channel_id),
                agent_a=AgentId(ch_row.agent_a),
                agent_b=AgentId(ch_row.agent_b),
                state=ChannelState(
                    local_state=StateValue(ch_row.local_state),
                    step=ch_row.step,
                ),
                depth=ch_row.depth,
                status=status_map.get(ch_row.status, ChannelStatus.ACTIVE),
                validation_failure_count=ch_row.validation_failure_count,
            )
            rt._channels[ch_row.channel_id] = channel

        # Reconstruct agents
        state_map = {
            "bound": AgentState.BOUND,
            "active": AgentState.ACTIVE,
            "quarantined": AgentState.QUARANTINED,
        }
        def _noop(msg): pass

        for ag_row in result.agents:
            agent_callable = callables.get(ag_row.agent_id, _noop)
            has_callable = ag_row.agent_id in callables
            state = state_map.get(ag_row.state, AgentState.BOUND)

            # Downgrade to BOUND if callable not provided and was ACTIVE
            if state == AgentState.ACTIVE and not has_callable:
                state = AgentState.BOUND
                rt._logger.warning(
                    f"Agent {ag_row.agent_id.hex()[:8]} recovered without callable — set to BOUND"
                )

            record = AgentRecord(
                agent_id=AgentId(ag_row.agent_id),
                state=state,
                callable=agent_callable,
                message_count=ag_row.message_count,
                quarantine_reason=ag_row.quarantine_reason,
            )

            # Rebuild agent's channel set
            for ch_id_bytes, ch in rt._channels.items():
                if ch.agent_a.value == ag_row.agent_id or ch.agent_b.value == ag_row.agent_id:
                    record.channels.add(ch_id_bytes)

            rt._agents[ag_row.agent_id] = record

        # Rebuild Merkle tree and Sg
        if rt._channels:
            channel_states = [
                (ch.channel_id, ch.state.local_state)
                for ch in rt._channels.values()
            ]
            rt._incremental_sg = IncrementalSg.from_channel_states(channel_states)
            rt._recompute_sg()

        rt._logger.info(
            f"Recovery complete: {len(rt._agents)} agents, "
            f"{len(rt._channels)} channels, "
            f"Sg={'present' if rt._sg else 'none'}",
        )

        for w in result.warnings:
            rt._logger.warning(f"Recovery warning: {w}")

        return rt

    # ------------------------------------------------------------------
    # Identity
    # ------------------------------------------------------------------

    def _derive_identity(self) -> StateValue:
        """Derive runtime identity from deployment and instance IDs.

        SHA-256(deployment_id || instance_id). Random values used if
        config fields are empty (sufficient for single-runtime operation).

        Maps to: I-06 §5.
        """
        deployment = self._config.deployment_id or random_bytes(32)
        instance = self._config.instance_id or random_bytes(32)
        return sha256(deployment + instance)

    def _generate_agent_id(self) -> AgentId:
        """Generate a unique agent ID via I-11 identity scheme.

        SHA-256(runtime_identity || counter || random_suffix).
        Counter persistence is handled by save_agent() in StorageEngine.
        """
        self._agent_counter += 1
        return generate_agent_id(self._identity, self._agent_counter)

    def _persist_meta(self) -> None:
        """Save runtime metadata to storage."""
        import time
        if not self._storage:
            return
        meta = RuntimeMeta(
            runtime_id=self._identity.data,
            deployment_id=self._config.deployment_id or b"",
            instance_id=self._config.instance_id or b"",
            agent_counter=self._agent_counter,
            schema_version=1,
            created_at=int(time.time()),
        )
        self._storage.save_runtime_meta(meta)

    # ------------------------------------------------------------------
    # Global State
    # ------------------------------------------------------------------

    def _recompute_sg(self) -> None:
        """Update Sg using incremental Merkle tree (O(1) - already updated).

        The Merkle tree is updated incrementally in establish_channel,
        close_channel, and after channel state advances. This method
        just syncs _sg with the tree root.

        Quarantined channels contribute (frozen state included).
        Closed channels do not (removed from registry).

        Maps to: spec.md §4.3, runtime-interface.md §8.3.
        """
        if not self._channels:
            self._sg = None
            self._incremental_sg = None
            return

        # Sg is already up-to-date in the Merkle tree
        if self._incremental_sg:
            self._sg = self._incremental_sg.get_root_hash()

    def _update_channel_in_tree(self, channel_id: ChannelId, new_state: StateValue) -> None:
        """Update a single channel state in the Merkle tree (O(log N))."""
        if self._incremental_sg:
            self._incremental_sg.update_channel(channel_id, new_state)

    # ------------------------------------------------------------------
    # Internal Lookups
    # ------------------------------------------------------------------

    def _lookup_agent(self, agent_id: AgentId) -> AgentRecord:
        """Look up an agent by ID. Raises AgentError if not found."""
        record = self._agents.get(agent_id.value)
        if record is None:
            raise AgentError(AgentErrorCode.UNBOUND, "Agent not bound")
        return record

    def _lookup_channel(self, channel_id: ChannelId) -> "Channel":
        """Look up a channel by ID. Raises AgentError if not found."""
        channel = get_channel(self._channels, channel_id)
        if channel is None:
            raise AgentError(AgentErrorCode.INVALID_CHANNEL, "Channel not found")
        return channel

    # ------------------------------------------------------------------
    # Properties
    # ------------------------------------------------------------------

    @property
    def identity(self) -> StateValue:
        """Runtime identity (immutable after init)."""
        return self._identity

    @property
    def global_state(self) -> GlobalState | None:
        """Current Sg. None if no channels exist."""
        return self._sg

    @property
    def protocol_version(self) -> int:
        """Wire protocol version supported by this runtime."""
        return PROTOCOL_VERSION

    @property
    def storage_degraded(self) -> bool:
        """True if the last storage operation failed (circuit breaker may be open)."""
        return self._storage_degraded

    # ------------------------------------------------------------------
    # Layer 2: Administration
    # ------------------------------------------------------------------

    def bind_agent(self, agent_callable: AgentCallable) -> AgentId:
        """Bind a new agent to the runtime.

        Creates an AgentRecord in BOUND state. The agent becomes ACTIVE
        when its first channel is established. Accepts both sync and async callables.

        Maps to: I-06 §7.1.
        """
        import inspect

        # Check max_agents limit
        if len(self._agents) >= self._config.max_agents:
            raise AgentError(
                AgentErrorCode.RESOURCE_LIMIT_EXCEEDED,
                f"Maximum agent limit reached ({self._config.max_agents})"
            )

        agent_id = self._generate_agent_id()
        record = AgentRecord(
            agent_id=agent_id,
            state=AgentState.BOUND,
            callable=agent_callable,
            is_async=inspect.iscoroutinefunction(agent_callable),
        )
        self._agents[agent_id.value] = record

        # Persist
        if self._storage:
            self._storage.save_agent(agent_id.value, "bound", self._identity.data)

        # Log agent binding
        context = LogContext(
            correlation_id=f"bind-{agent_id.value.hex()[:8]}",
            runtime_id=self._identity.data.hex()[:8],
            agent_id=agent_id.value.hex()[:8],
            operation="bind_agent",
        )
        log_audit_event(AuditEventType.AGENT_BOUND, context, agent_id=agent_id.value.hex()[:8])

        return agent_id

    def rebind_agent(self, agent_id: AgentId, agent_callable: AgentCallable) -> None:
        """Re-attach a callable to an existing BOUND agent.

        Used after recovery when an agent was restored without its callable.
        Transitions the agent to ACTIVE if it has channels.

        Maps to: I-06 §7.1 (rebind variant).
        """
        record = self._lookup_agent(agent_id)
        if record.state not in (AgentState.BOUND, AgentState.ACTIVE):
            raise AgentError(
                AgentErrorCode.UNBOUND,
                f"Cannot rebind agent in {record.state.name} state",
            )
        record.callable = agent_callable
        if record.channels and record.state == AgentState.BOUND:
            record.state = AgentState.ACTIVE
            if self._storage:
                self._storage.update_agent_state(agent_id.value, "active")

    def unbind_agent(self, agent_id: AgentId) -> None:
        """Unbind an agent from the runtime.

        Closes all channels, sets state to TERMINATED, removes from registry.

        Maps to: I-06 §7.2.
        """
        record = self._lookup_agent(agent_id)

        # Close all channels the agent participates in
        for ch_id_bytes in list(record.channels):
            self.close_channel(ChannelId(ch_id_bytes))

        record.state = AgentState.TERMINATED
        del self._agents[agent_id.value]

        # Persist
        if self._storage:
            self._storage.delete_agent(agent_id.value)

        # Audit
        context = LogContext(
            correlation_id=f"unbind-{agent_id.value.hex()[:8]}",
            runtime_id=self._identity.data.hex()[:8],
            operation="unbind_agent",
        )
        log_audit_event(AuditEventType.AGENT_UNBOUND, context, agent_id=agent_id.value.hex()[:8], reason="unbound")

    def establish_channel(
        self,
        agent_a: AgentId,
        agent_b: AgentId,
        depth: int = DEFAULT_FRAME_DEPTH,
    ) -> ChannelId:
        """Establish a channel between two bound agents.

        Derives Sl0, registers channel, adds to agents' channel sets,
        transitions BOUND agents to ACTIVE, recomputes Sg.

        Maps to: I-06 §8.1.
        """
        rec_a = self._lookup_agent(agent_a)
        rec_b = self._lookup_agent(agent_b)

        if rec_a.state == AgentState.QUARANTINED:
            raise AgentError(AgentErrorCode.QUARANTINED, "Agent A is quarantined")
        if rec_b.state == AgentState.QUARANTINED:
            raise AgentError(AgentErrorCode.QUARANTINED, "Agent B is quarantined")

        # Check max_channels_per_agent limit
        if len(rec_a.channels) >= self._config.max_channels_per_agent:
            raise AgentError(
                AgentErrorCode.RESOURCE_LIMIT_EXCEEDED,
                f"Agent A channel limit reached ({self._config.max_channels_per_agent})"
            )
        if len(rec_b.channels) >= self._config.max_channels_per_agent:
            raise AgentError(
                AgentErrorCode.RESOURCE_LIMIT_EXCEEDED,
                f"Agent B channel limit reached ({self._config.max_channels_per_agent})"
            )

        channel = ch_establish(
            registry=self._channels,
            runtime_identity=self._identity,
            agent_a=agent_a,
            agent_b=agent_b,
            depth=depth,
        )

        # Register channel with both agents
        rec_a.channels.add(channel.channel_id.value)
        rec_b.channels.add(channel.channel_id.value)

        # Transition BOUND → ACTIVE
        if rec_a.state == AgentState.BOUND:
            rec_a.state = AgentState.ACTIVE
        if rec_b.state == AgentState.BOUND:
            rec_b.state = AgentState.ACTIVE

        # Add channel to Merkle tree or rebuild if first channel
        if self._incremental_sg is None:
            # First channel - build initial tree
            channel_states = [
                (ch.channel_id, ch.state.local_state)
                for ch in self._channels.values()
            ]
            self._incremental_sg = IncrementalSg.from_channel_states(channel_states)
        else:
            # Add new channel to existing tree
            self._incremental_sg.add_channel(channel.channel_id, channel.state.local_state)

        self._recompute_sg()  # Sync _sg with tree root

        # Persist
        if self._storage:
            self._storage.save_channel(channel, self._sg)
            if rec_a.state == AgentState.ACTIVE:
                self._storage.update_agent_state(agent_a.value, "active")
            if rec_b.state == AgentState.ACTIVE:
                self._storage.update_agent_state(agent_b.value, "active")

        # Audit
        context = LogContext(
            correlation_id=f"establish-{channel.channel_id.value.hex()[:8]}",
            runtime_id=self._identity.data.hex()[:8],
            operation="establish_channel",
        )
        log_audit_event(
            AuditEventType.CHANNEL_ESTABLISHED, context,
            channel_id=channel.channel_id.value.hex()[:8],
            agent_a=agent_a.value.hex()[:8],
            agent_b=agent_b.value.hex()[:8],
        )

        return channel.channel_id

    def close_channel(self, channel_id: ChannelId) -> None:
        """Close a channel and zero its state.

        Removes from agents' channel sets, zeros Sl, recomputes Sg.

        Maps to: I-06 §8.2.
        """
        channel = self._lookup_channel(channel_id)

        # Remove from agents' channel sets
        rec_a = self._agents.get(channel.agent_a.value)
        if rec_a:
            rec_a.channels.discard(channel_id.value)
        rec_b = self._agents.get(channel.agent_b.value)
        if rec_b:
            rec_b.channels.discard(channel_id.value)

        # Remove from Merkle tree before closing
        if self._incremental_sg:
            self._incremental_sg.remove_channel(channel_id)

        ch_close(self._channels, channel_id)
        self._deduplication.clear_channel(channel_id)
        self._channel_locks.pop(channel_id.value, None)  # Clean up lock
        self._recompute_sg()  # Sync _sg with tree root

        # Deactivate agents that have no remaining channels (ACTIVE → BOUND)
        for rec in (rec_a, rec_b):
            if rec and rec.state == AgentState.ACTIVE and not rec.channels:
                rec.state = AgentState.BOUND

        # Persist
        if self._storage:
            self._storage.close_channel(channel_id.value, self._sg)
            for rec in (rec_a, rec_b):
                if rec and rec.state == AgentState.BOUND:
                    self._storage.update_agent_state(rec.agent_id.value, "bound")

        # Audit
        context = LogContext(
            correlation_id=f"close-{channel_id.value.hex()[:8]}",
            runtime_id=self._identity.data.hex()[:8],
            operation="close_channel",
        )
        log_audit_event(AuditEventType.CHANNEL_CLOSED, context, channel_id=channel_id.value.hex()[:8], reason="closed")

    def quarantine_agent(self, agent_id: AgentId, reason: str = "") -> None:
        """Quarantine an agent and all its active channels.

        Maps to: runtime-interface.md §8.3.
        """
        record = self._lookup_agent(agent_id)
        q_quarantine_agent(record, self._channels, reason)

        # Persist
        if self._storage:
            self._storage.update_agent_state(agent_id.value, "quarantined", reason)
            for ch_id_bytes in record.channels:
                ch = self._channels.get(ch_id_bytes)
                if ch and ch.status == ChannelStatus.QUARANTINED:
                    self._storage.quarantine_channel(ch_id_bytes)

        # Record metrics
        metrics = get_metrics_collector()
        metrics.increment_quarantine_events(reason=reason or "unknown")

    def quarantine_channel(self, channel_id: ChannelId, reason: str = "") -> None:
        """Quarantine a single channel.

        Maps to: runtime-interface.md §8.3.
        """
        channel = self._lookup_channel(channel_id)
        q_quarantine_channel(channel, reason)

        # Persist
        if self._storage:
            self._storage.quarantine_channel(channel_id.value)

        # Record metrics
        metrics = get_metrics_collector()
        metrics.increment_quarantine_events(reason=reason or "unknown")

    def restore_agent(self, agent_id: AgentId) -> None:
        """Restore a quarantined agent and its channels.

        Maps to: runtime-interface.md §8.4.
        """
        record = self._lookup_agent(agent_id)
        q_restore_agent(record, self._channels)

        # Persist
        if self._storage:
            self._storage.update_agent_state(agent_id.value, "active")
            for ch_id_bytes in record.channels:
                ch = self._channels.get(ch_id_bytes)
                if ch and ch.status == ChannelStatus.ACTIVE:
                    self._storage.restore_channel(ch_id_bytes)

        # Audit
        context = LogContext(
            correlation_id=f"restore-{agent_id.value.hex()[:8]}",
            runtime_id=self._identity.data.hex()[:8],
            operation="restore_agent",
        )
        log_audit_event(AuditEventType.RESTORE_AGENT, context, agent_id=agent_id.value.hex()[:8])

    def restore_channel(self, channel_id: ChannelId) -> None:
        """Restore a quarantined channel.

        Maps to: runtime-interface.md §8.4.
        """
        channel = self._lookup_channel(channel_id)
        q_restore_channel(channel)

        # Persist
        if self._storage:
            self._storage.restore_channel(channel_id.value)

        # Audit
        context = LogContext(
            correlation_id=f"restore-{channel_id.value.hex()[:8]}",
            runtime_id=self._identity.data.hex()[:8],
            operation="restore_channel",
        )
        log_audit_event(AuditEventType.RESTORE_CHANNEL, context, channel_id=channel_id.value.hex()[:8])

    # ------------------------------------------------------------------
    # Federation (Layer 2)
    # ------------------------------------------------------------------

    @property
    def runtime_id(self) -> RuntimeId:
        """This runtime's identity as a RuntimeId for federation."""
        return RuntimeId(self._identity)

    def bootstrap_bilateral_deterministic(
        self, peer_runtime_id: RuntimeId,
    ) -> BilateralChannel:
        """Bootstrap a bilateral channel using deterministic derivation.

        Both runtimes independently compute identical initial state
        from their sorted identities. No out-of-band exchange needed.

        Maps to: federation.md §3 (deterministic bootstrap).
        """
        bid = derive_bilateral_id(self.runtime_id, peer_runtime_id)
        if bid in self._bilateral_channels:
            return self._bilateral_channels[bid]

        if len(self._bilateral_channels) >= self._config.max_bilateral_channels:
            raise AgentError(
                AgentErrorCode.RESOURCE_LIMIT_EXCEEDED,
                "Maximum bilateral channel limit reached",
            )

        state = bootstrap_deterministic(self.runtime_id, peer_runtime_id)
        bc = BilateralChannel(
            bilateral_id=bid,
            local_runtime=self.runtime_id,
            peer_runtime=peer_runtime_id,
            state=state,
        )
        self._bilateral_channels[bid] = bc

        if self._storage:
            from mfp.storage.engine import BilateralRow
            import time
            self._storage.save_bilateral(BilateralRow(
                bilateral_id=bid,
                runtime_id_local=self._identity.data,
                runtime_id_peer=peer_runtime_id.value.data,
                ratchet_state=state.ratchet_state.data,
                shared_prng_seed=state.shared_prng_seed.data,
                step=state.step,
                status="active",
                created_at=int(time.time()),
                updated_at=int(time.time()),
            ))

        self._logger.info(
            f"Bilateral channel bootstrapped (deterministic) with peer "
            f"{peer_runtime_id.value.data.hex()[:8]}",
        )

        # Audit
        context = LogContext(
            correlation_id=f"bilateral-{bid.hex()[:8]}",
            runtime_id=self._identity.data.hex()[:8],
            operation="bootstrap_bilateral_deterministic",
        )
        log_audit_event(
            AuditEventType.BILATERAL_BOOTSTRAP, context,
            bilateral_id=bid.hex()[:8],
            peer_runtime=peer_runtime_id.value.data.hex()[:8],
        )

        return bc

    def bootstrap_bilateral_ceremonial(
        self,
        peer_runtime_id: RuntimeId,
        shared_secret: bytes,
    ) -> BilateralChannel:
        """Bootstrap a bilateral channel using DH shared secret.

        Requires out-of-band key exchange. Provides initial confidentiality.

        Maps to: federation.md §4 (trust ceremony).
        """
        bid = derive_bilateral_id(self.runtime_id, peer_runtime_id)
        if bid in self._bilateral_channels:
            return self._bilateral_channels[bid]

        if len(self._bilateral_channels) >= self._config.max_bilateral_channels:
            raise AgentError(
                AgentErrorCode.RESOURCE_LIMIT_EXCEEDED,
                "Maximum bilateral channel limit reached",
            )

        state = bootstrap_ceremonial(self.runtime_id, peer_runtime_id, shared_secret)
        bc = BilateralChannel(
            bilateral_id=bid,
            local_runtime=self.runtime_id,
            peer_runtime=peer_runtime_id,
            state=state,
        )
        self._bilateral_channels[bid] = bc

        if self._storage:
            from mfp.storage.engine import BilateralRow
            import time
            self._storage.save_bilateral(BilateralRow(
                bilateral_id=bid,
                runtime_id_local=self._identity.data,
                runtime_id_peer=peer_runtime_id.value.data,
                ratchet_state=state.ratchet_state.data,
                shared_prng_seed=state.shared_prng_seed.data,
                step=state.step,
                status="active",
                created_at=int(time.time()),
                updated_at=int(time.time()),
            ))

        self._logger.info(
            f"Bilateral channel bootstrapped (ceremonial) with peer "
            f"{peer_runtime_id.value.data.hex()[:8]}",
        )
        return bc

    def establish_cross_runtime_channel(
        self,
        local_agent: AgentId,
        remote_agent_id: bytes,
        bilateral_id: bytes,
        depth: int = DEFAULT_FRAME_DEPTH,
    ) -> ChannelId:
        """Establish a channel that spans two runtimes.

        The local agent must be ACTIVE. The remote agent is identified
        by raw bytes (its AgentId lives on the peer runtime).

        Maps to: federation.md §6.
        """
        rec = self._lookup_agent(local_agent)
        if rec.state != AgentState.ACTIVE and rec.state != AgentState.BOUND:
            raise AgentError(AgentErrorCode.UNBOUND, "Local agent not available")

        bc = self._bilateral_channels.get(bilateral_id)
        if bc is None:
            raise AgentError(
                AgentErrorCode.INVALID_CHANNEL,
                "Bilateral channel not found",
            )

        from mfp.core.primitives import random_id
        ch_id = ChannelId(random_id(16))
        crc = CrossRuntimeChannel(
            channel_id=ch_id,
            local_agent_id=local_agent.value,
            remote_agent_id=remote_agent_id,
            bilateral_id=bilateral_id,
            state=ChannelState(
                local_state=bc.state.ratchet_state,
                step=bc.state.step,
            ),
            depth=depth,
        )
        self._cross_runtime_channels[ch_id.value] = crc
        bc.cross_runtime_channels.append(crc)

        # Activate agent if BOUND
        if rec.state == AgentState.BOUND:
            rec.state = AgentState.ACTIVE
            if self._storage:
                self._storage.update_agent_state(local_agent.value, "active")

        self._logger.info(
            f"Cross-runtime channel established: {ch_id.value.hex()[:8]}",
        )
        return ch_id

    async def start_federation(self, transport_config: TransportConfig) -> None:
        """Start federation transport: TCP server + connection pool.

        The server listens for incoming messages from peer runtimes.
        The connection pool manages outbound connections.
        """
        self._connection_pool = ConnectionPool(transport_config)
        self._transport_server = TransportServer(
            transport_config, self._handle_incoming_message,
        )
        await self._transport_server.start()
        self._logger.info(
            f"Federation started on {transport_config.host}:{transport_config.port}",
        )

    async def stop_federation(self) -> None:
        """Stop federation transport."""
        if self._transport_server:
            await self._transport_server.stop()
            self._transport_server = None
        if self._connection_pool:
            await self._connection_pool.close_all()
            self._connection_pool = None

    async def _handle_incoming_message(
        self, header: "EnvelopeHeader", msg: "ProtocolMessage",
    ) -> None:
        """Handle a message received from a peer runtime.

        Validates the envelope, finds the cross-runtime channel,
        decodes the payload, and delivers to the local agent.
        """
        from mfp.core.encoding import build_encoding_context, decode
        from mfp.core.frame import validate_frame, sample_frame_cross_runtime

        # Validate envelope
        known_channels = set(self._cross_runtime_channels.keys())
        errors = validate_envelope(header, known_channels=known_channels)
        if errors:
            self._logger.warning(f"Invalid envelope: {errors}")
            return

        # Find cross-runtime channel
        crc = self._cross_runtime_channels.get(header.channel_id.value)
        if crc is None:
            self._logger.warning(f"Unknown channel in envelope: {header.channel_id.value.hex()[:8]}")
            return

        bc = self._bilateral_channels.get(crc.bilateral_id)
        if bc is None:
            self._logger.warning(f"No bilateral channel for CRC")
            return

        # Validate frame using bilateral state
        expected_frame = sample_frame_cross_runtime(
            local_state=bc.state.ratchet_state,
            step=bc.state.step,
            bilateral_ratchet_state=bc.state.ratchet_state,
            shared_prng_seed=bc.state.shared_prng_seed,
            depth=header.frame_depth,
        )
        if not validate_frame(msg.frame_open, msg.frame_close, expected_frame):
            self._logger.error("Cross-runtime frame validation failed")
            return

        # Decode payload
        encoding_ctx = build_encoding_context(
            bc.state.ratchet_state,
            ChannelId(header.channel_id.value if isinstance(header.channel_id, ChannelId) else header.channel_id),
            bc.state.step,
        )
        plaintext = decode(msg.encoded_payload, encoding_ctx)
        if plaintext is None:
            self._logger.error("Cross-runtime payload decode failed")
            return

        # Deliver to local agent
        local_agent_rec = self._agents.get(crc.local_agent_id)
        if local_agent_rec and local_agent_rec.state == AgentState.ACTIVE:
            from mfp.core.types import DeliveredMessage, MessageId
            from mfp.core.primitives import random_id
            delivered = DeliveredMessage(
                payload=plaintext,
                sender=AgentId(crc.remote_agent_id),
                channel=ChannelId(crc.channel_id.value),
                message_id=MessageId(random_id(16)),
            )
            try:
                local_agent_rec.callable(delivered)
            except Exception as e:
                self._logger.error(f"Agent delivery failed: {e}")

        # Advance bilateral state
        bc.state = advance_bilateral_state(bc.state, expected_frame)

        self._logger.info("Cross-runtime message received and delivered")

    async def flush_pending_messages(self) -> None:
        """Send all pending wire messages via the connection pool.

        Called by the transport layer or an event loop after send().
        """
        if not self._connection_pool:
            return

        while self._pending_wire_messages:
            peer_id, wire_bytes = self._pending_wire_messages.pop(0)
            # Find peer endpoint from bilateral channels
            for bc in self._bilateral_channels.values():
                if bc.peer_runtime.value.data == peer_id:
                    # Would need host:port mapping — for now, log
                    self._logger.debug(
                        f"Would send {len(wire_bytes)} bytes to peer {peer_id.hex()[:8]}"
                    )
                    break

    # ------------------------------------------------------------------
    # Persistence migration
    # ------------------------------------------------------------------

    def enable_persistence(self, storage_config: "StorageConfig") -> None:
        """Upgrade a running in-memory Runtime to persistent mode.

        Snapshots current state (agents, channels, Sg) into a new
        StorageEngine. From this point forward, all mutations persist.

        Raises ValueError if storage is already active.
        """
        from mfp.storage.engine import StorageConfig as _SC  # noqa: F811
        if self._storage is not None:
            raise ValueError("Storage is already active")

        storage = StorageEngine(storage_config)

        # Persist runtime metadata
        import time
        meta = RuntimeMeta(
            runtime_id=self._identity.data,
            deployment_id=self._config.deployment_id or b"",
            instance_id=self._config.instance_id or b"",
            agent_counter=self._agent_counter,
            schema_version=1,
            created_at=int(time.time()),
        )
        storage.save_runtime_meta(meta)

        # Persist all agents
        for record in self._agents.values():
            storage.save_agent(
                record.agent_id.value,
                record.state.name.lower(),
                self._identity.data,
            )

        # Persist all channels and Sg
        for channel in self._channels.values():
            storage.save_channel(channel, self._sg)

        if self._sg:
            storage.save_sg_cache(self._sg)

        self._storage = storage
        self._logger.info("Persistence enabled — state migrated to storage")

    # ------------------------------------------------------------------
    # Layer 1: Agent-Facing
    # ------------------------------------------------------------------

    def send(
        self,
        sender: AgentId,
        channel_id: ChannelId,
        payload: bytes,
    ) -> Receipt:
        """Process a message through the six-stage pipeline.

        On success: advance ratchet, recompute Sg, return receipt.
        On validation failure: increment failure count, auto-quarantine
        if threshold reached, re-raise.
        On any failure: no state change (atomic).

        Maps to: runtime-interface.md §4.
        """
        # Generate correlation ID for this send operation
        from mfp.core.primitives import random_bytes
        correlation_id = random_bytes(16).hex()[:16]

        context = LogContext(
            correlation_id=correlation_id,
            runtime_id=self._identity.data.hex()[:8],
            agent_id=sender.value.hex()[:8],
            channel_id=channel_id.value.hex()[:8],
            operation="send",
        )

        self._logger.debug("Message send initiated", context=context, payload_size=len(payload))

        sender_rec = self._lookup_agent(sender)

        # Verify sender state
        if sender_rec.state == AgentState.QUARANTINED:
            raise AgentError(AgentErrorCode.QUARANTINED, "Sender is quarantined")
        if sender_rec.state != AgentState.ACTIVE:
            raise AgentError(AgentErrorCode.UNBOUND, "Sender is not active")

        # Rate limit check (sliding window — messages per second)
        if check_rate_limit_window(sender_rec._message_timestamps, self._config.max_message_rate):
            self.quarantine_agent(sender, "Rate limit exceeded")
            raise AgentError(AgentErrorCode.QUARANTINED, "Rate limit exceeded")

        # Cross-runtime path
        crc = self._cross_runtime_channels.get(channel_id.value)
        if crc is not None:
            return self._send_cross_runtime(sender, channel_id, payload, crc, correlation_id, context)

        channel = self._lookup_channel(channel_id)

        # Determine destination agent
        if sender.value == channel.agent_a.value:
            dest_id = channel.agent_b
        elif sender.value == channel.agent_b.value:
            dest_id = channel.agent_a
        else:
            raise AgentError(AgentErrorCode.INVALID_CHANNEL, "Sender not on channel")

        dest_rec = self._lookup_agent(dest_id)

        # Global state must exist (channels exist)
        if self._sg is None:
            raise AgentError(AgentErrorCode.INVALID_CHANNEL, "No global state")

        # Per-channel lock ensures message ordering (spec: runtime-interface.md §4.9)
        channel_lock = self._get_channel_lock(channel_id.value)
        with channel_lock:
            try:
                result = process_message(
                    sender=sender,
                    channel=channel,
                    payload=payload,
                    global_state=self._sg,
                    config=self._config,
                    deliver=dest_rec.callable,
                    correlation_id=correlation_id,
                )
            except AgentError as e:
                # Agent timeout: quarantine the destination agent
                if e.code == AgentErrorCode.TIMEOUT:
                    self._logger.warning(
                        f"Agent timeout, quarantining destination agent: {str(e)}",
                        context=context,
                        dest_agent=dest_id.value.hex()[:8],
                    )
                    self.quarantine_agent(dest_id, f"Agent callable timeout: {str(e)}")
                raise
            except FrameValidationError as e:
                # Validation failure: track and potentially quarantine
                increment_failure_count(channel)

                # Record metrics
                metrics = get_metrics_collector()
                metrics.increment_validation_failures(error_type=type(e).__name__)

                if check_validation_failure(
                    channel, self._config.validation_failure_threshold
                ):
                    q_quarantine_channel(channel)
                raise

            # Replay protection: record message ID
            self._deduplication.is_duplicate(channel_id, result.receipt.message_id)

            # Success: advance state
            advance_channel(channel, result.new_local_state)
            reset_failure_count(channel)
            sender_rec.message_count += 1
            sender_rec._message_timestamps.append(time.time())

            # Update Merkle tree incrementally (O(log N) instead of O(N))
            self._update_channel_in_tree(channel_id, result.new_local_state)
            with self._sg_lock:
                self._recompute_sg()  # Sync _sg with tree root

            # Persist ratchet advancement and Sg.
            if self._storage:
                try:
                    self._storage.advance_channel(
                        channel_id.value, result.new_local_state, self._sg,
                    )
                    self._storage_degraded = False
                except Exception as e:
                    self._storage_degraded = True
                    self._logger.warning(
                        f"Storage write failed after send — state diverged: {e}",
                        context=context,
                    )

        # Add correlation_id to receipt
        receipt = Receipt(
            message_id=result.receipt.message_id,
            channel=result.receipt.channel,
            step=result.receipt.step,
            correlation_id=correlation_id,
        )

        # Audit
        log_audit_event(
            AuditEventType.MESSAGE_SENT, context,
            channel_id=channel_id.value.hex()[:8],
            sender=sender.value.hex()[:8],
            step=receipt.step,
        )

        self._logger.info("Message sent successfully", context=context, step=result.receipt.step)

        # Record metrics
        metrics = get_metrics_collector()
        metrics.increment_messages_sent(
            agent_id=sender.value.hex()[:8],
            channel_id=channel_id.value.hex()[:8],
        )
        metrics.increment_messages_received(
            agent_id=dest_id.value.hex()[:8],
        )
        metrics.observe_message_size(len(payload))

        return receipt

    async def send_async(
        self,
        sender: AgentId,
        channel_id: ChannelId,
        payload: bytes,
    ) -> Receipt:
        """Async variant of send(). Uses async pipeline for async agent callables.

        For sync callables, delegates to the sync pipeline.
        For async callables, uses process_message_async with asyncio timeout.
        """
        from mfp.core.primitives import random_bytes
        from mfp.runtime.pipeline import process_message_async

        correlation_id = random_bytes(16).hex()[:16]
        context = LogContext(
            correlation_id=correlation_id,
            runtime_id=self._identity.data.hex()[:8],
            agent_id=sender.value.hex()[:8],
            channel_id=channel_id.value.hex()[:8],
            operation="send_async",
        )

        sender_rec = self._lookup_agent(sender)
        if sender_rec.state == AgentState.QUARANTINED:
            raise AgentError(AgentErrorCode.QUARANTINED, "Sender is quarantined")
        if sender_rec.state != AgentState.ACTIVE:
            raise AgentError(AgentErrorCode.UNBOUND, "Sender is not active")

        if check_rate_limit_window(sender_rec._message_timestamps, self._config.max_message_rate):
            self.quarantine_agent(sender, "Rate limit exceeded")
            raise AgentError(AgentErrorCode.QUARANTINED, "Rate limit exceeded")

        # Cross-runtime path (same as sync — outbound is not async-dependent)
        crc = self._cross_runtime_channels.get(channel_id.value)
        if crc is not None:
            return self._send_cross_runtime(sender, channel_id, payload, crc, correlation_id, context)

        channel = self._lookup_channel(channel_id)

        if sender.value == channel.agent_a.value:
            dest_id = channel.agent_b
        elif sender.value == channel.agent_b.value:
            dest_id = channel.agent_a
        else:
            raise AgentError(AgentErrorCode.INVALID_CHANNEL, "Sender not on channel")

        dest_rec = self._lookup_agent(dest_id)

        if self._sg is None:
            raise AgentError(AgentErrorCode.INVALID_CHANNEL, "No global state")

        # Per-channel lock ensures message ordering (spec: runtime-interface.md §4.9)
        channel_lock = self._get_channel_lock(channel_id.value)
        with channel_lock:
            # Use async pipeline if destination agent is async, else sync
            if dest_rec.is_async:
                result = await process_message_async(
                    sender=sender,
                    channel=channel,
                    payload=payload,
                    global_state=self._sg,
                    config=self._config,
                    deliver=dest_rec.callable,
                    correlation_id=correlation_id,
                )
            else:
                result = process_message(
                    sender=sender,
                    channel=channel,
                    payload=payload,
                    global_state=self._sg,
                    config=self._config,
                    deliver=dest_rec.callable,
                    correlation_id=correlation_id,
                )

            # Same post-processing as sync send
            self._deduplication.is_duplicate(channel_id, result.receipt.message_id)
            advance_channel(channel, result.new_local_state)
            reset_failure_count(channel)
            sender_rec.message_count += 1
            sender_rec._message_timestamps.append(time.time())
            self._update_channel_in_tree(channel_id, result.new_local_state)
            with self._sg_lock:
                self._recompute_sg()

            if self._storage:
                try:
                    self._storage.advance_channel(
                        channel_id.value, result.new_local_state, self._sg,
                    )
                    self._storage_degraded = False
                except Exception as e:
                    self._storage_degraded = True
                    self._logger.warning(f"Storage write failed after async send: {e}", context=context)

        return Receipt(
            message_id=result.receipt.message_id,
            channel=result.receipt.channel,
            step=result.receipt.step,
            correlation_id=correlation_id,
        )

    def _send_cross_runtime(
        self,
        sender: AgentId,
        channel_id: ChannelId,
        payload: bytes,
        crc: CrossRuntimeChannel,
        correlation_id: str,
        context: LogContext,
    ) -> Receipt:
        """Process a cross-runtime message: ACCEPT→FRAME→ENCODE→VALIDATE→TRANSIT.

        Produces wire bytes for transport. Does NOT deliver locally.
        The wire bytes are stored in self._pending_wire_messages for the
        transport layer to pick up, or returned via receipt metadata.
        """
        from mfp.runtime.pipeline import process_cross_runtime_outbound

        bc = self._bilateral_channels.get(crc.bilateral_id)
        if bc is None:
            raise AgentError(AgentErrorCode.INVALID_CHANNEL, "Bilateral channel not found")

        # Bilateral-level lock: serializes all agent channels on same bilateral
        bilateral_lock = self._get_channel_lock(crc.bilateral_id)
        with bilateral_lock:
            return self._send_cross_runtime_locked(
                sender, channel_id, payload, crc, bc, correlation_id, context,
            )

    def _send_cross_runtime_locked(
        self,
        sender: AgentId,
        channel_id: ChannelId,
        payload: bytes,
        crc: CrossRuntimeChannel,
        bc: BilateralChannel,
        correlation_id: str,
        context: LogContext,
    ) -> Receipt:
        """Cross-runtime send while holding bilateral lock."""
        from mfp.runtime.pipeline import process_cross_runtime_outbound

        result = process_cross_runtime_outbound(
            sender=sender,
            channel_id=channel_id,
            payload=payload,
            bilateral_state=bc.state,
            depth=crc.depth,
            config=self._config,
            sender_runtime=self._identity.data[:16],
            correlation_id=correlation_id,
            encoding_algorithm=bc.encoding_algorithm,
        )

        # Advance bilateral state
        bc.state = result.new_bilateral_state
        bc.increment_message_count()

        # Update cross-runtime channel state to track step
        crc.state = ChannelState(
            local_state=result.new_bilateral_state.ratchet_state,
            step=result.new_bilateral_state.step,
        )

        # Persist
        if self._storage:
            try:
                from mfp.storage.engine import BilateralRow
                import time
                self._storage.save_bilateral(BilateralRow(
                    bilateral_id=bc.bilateral_id,
                    runtime_id_local=self._identity.data,
                    runtime_id_peer=bc.peer_runtime.value.data,
                    ratchet_state=bc.state.ratchet_state.data,
                    shared_prng_seed=bc.state.shared_prng_seed.data,
                    step=bc.state.step,
                    status="active",
                    created_at=int(time.time()),
                    updated_at=int(time.time()),
                ))
                self._storage_degraded = False
            except Exception as e:
                self._storage_degraded = True
                self._logger.warning(
                    f"Storage write failed for bilateral state: {e}",
                    context=context,
                )

        # Store wire bytes for transport pickup
        if not hasattr(self, '_pending_wire_messages'):
            self._pending_wire_messages: list[tuple[bytes, bytes]] = []
        self._pending_wire_messages.append(
            (bc.peer_runtime.value.data, result.wire_bytes)
        )

        self._logger.info(
            f"Cross-runtime message prepared for peer "
            f"{bc.peer_runtime.value.data.hex()[:8]}",
            context=context,
        )

        sender_rec = self._lookup_agent(sender)
        sender_rec.message_count += 1

        return result.receipt

    def check_agent_state(self, agent_id: AgentId) -> AgentState:
        """Return the current state of an agent.

        Layer 1 method — safe for AgentHandle to call without
        reaching into Runtime internals.

        Maps to: runtime-interface.md §3.
        """
        record = self._agents.get(agent_id.value)
        if record is None:
            raise AgentError(AgentErrorCode.UNBOUND, "Agent not bound")
        return record.state

    def get_channels(self, agent_id: AgentId) -> list[ChannelInfo]:
        """Return agent-visible channel information.

        Maps to: runtime-interface.md §3.2.
        """
        self._lookup_agent(agent_id)  # verify bound
        return get_channels_for_agent(self._channels, agent_id)

    def get_status(self, agent_id: AgentId) -> AgentStatus:
        """Return agent-visible status.

        Maps to: runtime-interface.md §3.3.
        """
        record = self._lookup_agent(agent_id)
        return AgentStatus(
            agent_id=agent_id,
            state=record.state,
            channel_count=len(record.channels),
        )

    # ------------------------------------------------------------------
    # Lifecycle
    # ------------------------------------------------------------------

    def shutdown(self) -> None:
        """Graceful shutdown. Close all channels, terminate agents, zero state.

        Maps to: I-06 §11.
        """
        # Close all channels (zeros Sl for each)
        for ch_id_bytes in list(self._channels.keys()):
            ch_close(self._channels, ChannelId(ch_id_bytes))

        # Terminate all agents
        for record in self._agents.values():
            record.state = AgentState.TERMINATED

        self._agents.clear()
        self._channels.clear()
        self._sg = None
        self._agent_counter = 0

        # Close storage
        if self._storage:
            self._storage.close()
            self._storage = None
