"""MFP Message Pipeline — the six-stage message lifecycle.

ACCEPT → FRAME → ENCODE → VALIDATE → DECODE → DELIVER

The pipeline is a function, not a class. It receives state, calls core
functions, and returns results. State advancement is the caller's job.

Maps to: impl/I-07_pipeline.md
"""

from __future__ import annotations

import asyncio
import inspect
from collections.abc import Awaitable, Callable
from dataclasses import dataclass, field

from mfp.core.encoding import build_encoding_context, decode, encode
from mfp.core.frame import (
    assemble_message,
    sample_frame,
    sample_frame_cross_runtime,
    validate_frame,
)
from mfp.core.primitives import random_id
from mfp.core.ratchet import advance
from mfp.core.types import (
    AgentError,
    AgentErrorCode,
    AgentId,
    BilateralState,
    Channel,
    ChannelId,
    ChannelStatus,
    DecodeError,
    DeliveredMessage,
    EnvelopeHeader,
    Frame,
    FrameValidationError,
    GlobalState,
    MessageId,
    ProtocolMessage,
    Receipt,
    StateValue,
)
from mfp.observability.logging import LogContext, TimedOperation, get_logger
from mfp.observability.timeout import TimeoutError as OperationTimeoutError, with_timeout
from mfp.storage.engine import StorageConfig

logger = get_logger(__name__)

AgentCallable = Callable[[DeliveredMessage], None]
AsyncAgentCallable = Callable[[DeliveredMessage], Awaitable[None]]


@dataclass(frozen=True)
class RuntimeConfig:
    """Runtime configuration parameters."""
    # Identity (empty = auto-generate random)
    deployment_id: bytes = b""
    instance_id: bytes = b""
    # Frame
    default_frame_depth: int = 4
    # Quarantine (secure defaults for production)
    validation_failure_threshold: int = 5
    max_message_rate: int = 1000  # messages per second
    max_payload_size: int = 1_048_576  # 1 MB
    # Encoding
    encoding_algorithm: bytes = b"aes-256-gcm"
    supported_algorithms: tuple[bytes, ...] = (b"aes-256-gcm",)
    # Resource Limits (P1.4)
    max_channels_per_agent: int = 100
    max_agents: int = 10_000
    max_bilateral_channels: int = 100
    max_storage_size_mb: int = 1024  # 1 GB
    # Timeouts (P2.3)
    agent_timeout_seconds: float = 30.0    # Agent callable timeout
    pipeline_timeout_seconds: float = 5.0  # Pipeline processing timeout
    storage_timeout_seconds: float = 10.0  # Storage operation timeout
    # Persistence (None = in-memory only)
    storage: StorageConfig | None = None


@dataclass(frozen=True)
class PipelineResult:
    """Result of a successful pipeline execution."""
    receipt: Receipt
    delivered: DeliveredMessage
    new_local_state: StateValue
    frame: Frame


# ---------------------------------------------------------------------------
# Stage 1 — ACCEPT
# ---------------------------------------------------------------------------

def accept(
    sender: AgentId,
    channel: Channel,
    payload: bytes,
    config: RuntimeConfig,
) -> MessageId:
    """Validate the send request.

    Maps to: runtime-interface.md §4.2.
    """
    # Verify sender is on this channel
    if sender != channel.agent_a and sender != channel.agent_b:
        raise AgentError(AgentErrorCode.INVALID_CHANNEL, "Sender not bound to channel")

    # Verify channel status
    if channel.status == ChannelStatus.QUARANTINED:
        raise AgentError(AgentErrorCode.CHANNEL_QUARANTINED, "Channel is quarantined")
    if channel.status == ChannelStatus.CLOSED:
        raise AgentError(AgentErrorCode.CHANNEL_CLOSED, "Channel is closed")

    # Verify payload size
    if config.max_payload_size > 0 and len(payload) > config.max_payload_size:
        raise AgentError(
            AgentErrorCode.PAYLOAD_TOO_LARGE,
            f"Payload {len(payload)} bytes exceeds limit {config.max_payload_size}",
        )

    return MessageId(random_id(16))


# ---------------------------------------------------------------------------
# Stage 2 — FRAME
# ---------------------------------------------------------------------------

def frame_stage(channel: Channel, global_state: GlobalState) -> Frame:
    """Derive and sample the frame.

    Maps to: runtime-interface.md §4.3.
    """
    return sample_frame(
        local_state=channel.state.local_state,
        step=channel.state.step,
        global_state=global_state.value,
        depth=channel.depth,
    )


# ---------------------------------------------------------------------------
# Stage 3 — ENCODE
# ---------------------------------------------------------------------------

def encode_stage(
    channel: Channel,
    payload: bytes,
    frame: Frame,
    config: RuntimeConfig,
) -> ProtocolMessage:
    """Encode payload and assemble the protocol message.

    Maps to: runtime-interface.md §4.4.
    """
    ctx = build_encoding_context(
        local_state=channel.state.local_state,
        channel_id=channel.channel_id,
        step=channel.state.step,
        algorithm_id=config.encoding_algorithm,
    )
    encoded = encode(payload, ctx)
    return assemble_message(frame, encoded)


# ---------------------------------------------------------------------------
# Stage 4 — VALIDATE
# ---------------------------------------------------------------------------

def validate_stage(msg: ProtocolMessage, expected_frame: Frame) -> None:
    """Validate the message's frame pair.

    Maps to: runtime-interface.md §4.5.
    """
    if not validate_frame(msg.frame_open, msg.frame_close, expected_frame):
        raise FrameValidationError("Frame validation failed")


# ---------------------------------------------------------------------------
# Stage 5 — DECODE
# ---------------------------------------------------------------------------

def decode_stage(
    msg: ProtocolMessage,
    channel: Channel,
    config: RuntimeConfig,
) -> bytes:
    """Decode the encrypted payload.

    Maps to: runtime-interface.md §4.6.
    """
    ctx = build_encoding_context(
        local_state=channel.state.local_state,
        channel_id=channel.channel_id,
        step=channel.state.step,
        algorithm_id=config.encoding_algorithm,
    )
    payload = decode(msg.encoded_payload, ctx)
    if payload is None:
        raise DecodeError("Payload integrity check failed")
    return payload


# ---------------------------------------------------------------------------
# Stage 6 — DELIVER
# ---------------------------------------------------------------------------

def deliver_stage(
    payload: bytes,
    sender: AgentId,
    channel_id: ChannelId,
    message_id: MessageId,
    deliver: AgentCallable,
    correlation_id: str = "",
    agent_timeout_seconds: float = 30.0,
) -> DeliveredMessage:
    """Deliver decoded message to the destination agent.

    Agent callable execution is protected by timeout to prevent blocking.

    Maps to: runtime-interface.md §4.7.

    Raises:
        OperationTimeoutError: If agent callable exceeds timeout
    """
    msg = DeliveredMessage(
        payload=payload,
        sender=sender,
        channel=channel_id,
        message_id=message_id,
        correlation_id=correlation_id,
    )

    # Wrap agent callable with timeout protection
    try:
        with_timeout(
            lambda: deliver(msg),
            timeout_seconds=agent_timeout_seconds,
            operation_name=f"agent_callable_{sender.value.hex()[:8]}",
        )
    except OperationTimeoutError as e:
        # Agent timeout should trigger quarantine
        logger.error(
            f"Agent callable timeout: {e}",
            extra={
                "agent_id": sender.value.hex()[:8],
                "channel_id": channel_id.value.hex()[:8],
                "timeout_seconds": agent_timeout_seconds,
            }
        )
        raise AgentError(
            AgentErrorCode.TIMEOUT,
            f"Agent callable exceeded timeout of {agent_timeout_seconds}s"
        )

    return msg


async def deliver_stage_async(
    payload: bytes,
    sender: AgentId,
    channel_id: ChannelId,
    message_id: MessageId,
    deliver: AsyncAgentCallable,
    correlation_id: str = "",
    agent_timeout_seconds: float = 30.0,
) -> DeliveredMessage:
    """Async variant of deliver_stage. Uses asyncio.wait_for for timeout."""
    msg = DeliveredMessage(
        payload=payload,
        sender=sender,
        channel=channel_id,
        message_id=message_id,
        correlation_id=correlation_id,
    )
    try:
        await asyncio.wait_for(
            deliver(msg),
            timeout=agent_timeout_seconds,
        )
    except asyncio.TimeoutError:
        raise AgentError(
            AgentErrorCode.TIMEOUT,
            f"Async agent callable exceeded timeout of {agent_timeout_seconds}s",
        )
    return msg


# ---------------------------------------------------------------------------
# Full Pipeline
# ---------------------------------------------------------------------------

def process_message(
    sender: AgentId,
    channel: Channel,
    payload: bytes,
    global_state: GlobalState,
    config: RuntimeConfig,
    deliver: AgentCallable,
    correlation_id: str = "",
) -> PipelineResult:
    """Execute the six-stage message pipeline.

    Returns PipelineResult with new state. The caller (Runtime) applies
    the state mutation and Sg recomputation.

    Maps to: runtime-interface.md §4.
    """
    # Create logging context
    context = LogContext(
        correlation_id=correlation_id,
        runtime_id="",  # Set by caller
        agent_id=sender.value.hex()[:8],
        channel_id=channel.channel_id.value.hex()[:8],
        operation="pipeline",
    )

    logger.debug("Pipeline started", context=context, step=channel.state.step)

    # Stage 1: ACCEPT
    context.stage = "ACCEPT"
    with TimedOperation("accept", context):
        message_id = accept(sender, channel, payload, config)

    # Stage 2: FRAME
    context.stage = "FRAME"
    with TimedOperation("frame", context):
        frame = frame_stage(channel, global_state)

    # Stage 3: ENCODE
    context.stage = "ENCODE"
    with TimedOperation("encode", context):
        protocol_msg = encode_stage(channel, payload, frame, config)

    # Stage 4: VALIDATE (intra-runtime self-check)
    context.stage = "VALIDATE"
    with TimedOperation("validate", context):
        validate_stage(protocol_msg, frame)

    # Stage 5: DECODE
    context.stage = "DECODE"
    with TimedOperation("decode", context):
        decoded = decode_stage(protocol_msg, channel, config)

    # Stage 6: DELIVER
    context.stage = "DELIVER"
    with TimedOperation("deliver", context):
        delivered = deliver_stage(
            decoded,
            sender,
            channel.channel_id,
            message_id,
            deliver,
            correlation_id,
            agent_timeout_seconds=config.agent_timeout_seconds,
        )

    # Compute new state (not yet applied)
    new_local = advance(channel.state.local_state, frame)

    logger.debug("Pipeline completed", context=context)

    return PipelineResult(
        receipt=Receipt(
            message_id=message_id,
            channel=channel.channel_id,
            step=channel.state.step,
            correlation_id=correlation_id,
        ),
        delivered=delivered,
        new_local_state=new_local,
        frame=frame,
    )


async def process_message_async(
    sender: AgentId,
    channel: Channel,
    payload: bytes,
    global_state: GlobalState,
    config: RuntimeConfig,
    deliver: AsyncAgentCallable,
    correlation_id: str = "",
) -> PipelineResult:
    """Async variant of process_message.

    Stages 1-5 are synchronous (CPU-bound crypto).
    Stage 6 (DELIVER) uses async agent callable with asyncio timeout.
    """
    context = LogContext(
        correlation_id=correlation_id,
        runtime_id="",
        agent_id=sender.value.hex()[:8],
        channel_id=channel.channel_id.value.hex()[:8],
        operation="pipeline_async",
    )

    # Stages 1-5: identical to sync pipeline
    message_id = accept(sender, channel, payload, config)
    frame = frame_stage(channel, global_state)
    protocol_msg = encode_stage(channel, payload, frame, config)
    validate_stage(protocol_msg, frame)
    decoded = decode_stage(protocol_msg, channel, config)

    # Stage 6: DELIVER (async)
    delivered = await deliver_stage_async(
        decoded, sender, channel.channel_id, message_id,
        deliver, correlation_id,
        agent_timeout_seconds=config.agent_timeout_seconds,
    )

    new_local = advance(channel.state.local_state, frame)

    return PipelineResult(
        receipt=Receipt(
            message_id=message_id,
            channel=channel.channel_id,
            step=channel.state.step,
            correlation_id=correlation_id,
        ),
        delivered=delivered,
        new_local_state=new_local,
        frame=frame,
    )


# ---------------------------------------------------------------------------
# Cross-Runtime Pipeline
# ---------------------------------------------------------------------------

@dataclass(frozen=True)
class CrossRuntimeResult:
    """Result of cross-runtime outbound pipeline (ACCEPT→FRAME→ENCODE→TRANSIT)."""
    receipt: Receipt
    wire_bytes: bytes
    header: EnvelopeHeader
    frame: Frame
    new_bilateral_state: BilateralState


def transit_stage(
    protocol_msg: ProtocolMessage,
    channel_id: ChannelId,
    step: int,
    frame_depth: int,
    sender_runtime: bytes,
) -> tuple[EnvelopeHeader, bytes]:
    """Stage TRANSIT: wrap a ProtocolMessage in a wire envelope for transport.

    Maps to: federation.md §6 (TRANSIT stage).
    """
    from mfp.federation.wire import build_envelope_header, assemble_wire_message

    header = build_envelope_header(
        channel_id=channel_id,
        step=step,
        frame_depth=frame_depth,
        payload_len=len(protocol_msg.encoded_payload),
        sender_runtime=sender_runtime,
    )
    wire_bytes = assemble_wire_message(header, protocol_msg)
    return header, wire_bytes


def process_cross_runtime_outbound(
    sender: AgentId,
    channel_id: ChannelId,
    payload: bytes,
    bilateral_state: BilateralState,
    depth: int,
    config: RuntimeConfig,
    sender_runtime: bytes,
    correlation_id: str = "",
    encoding_algorithm: bytes | None = None,
) -> CrossRuntimeResult:
    """Outbound cross-runtime pipeline: ACCEPT→FRAME→ENCODE→VALIDATE→TRANSIT.

    No DECODE or DELIVER — the peer runtime handles those.
    Uses deterministic cross-runtime frame sampling.
    """
    from mfp.federation.bilateral import advance_bilateral_state

    context = LogContext(
        correlation_id=correlation_id,
        runtime_id="",
        operation="cross_runtime_outbound",
    )

    # Stage 1: ACCEPT (validate payload size)
    if config.max_payload_size > 0 and len(payload) > config.max_payload_size:
        raise AgentError(
            AgentErrorCode.PAYLOAD_TOO_LARGE,
            f"Payload {len(payload)} bytes exceeds limit {config.max_payload_size}",
        )
    message_id = MessageId(random_id(16))

    # Stage 2: FRAME (cross-runtime deterministic sampling)
    frame = sample_frame_cross_runtime(
        local_state=bilateral_state.ratchet_state,
        step=bilateral_state.step,
        bilateral_ratchet_state=bilateral_state.ratchet_state,
        shared_prng_seed=bilateral_state.shared_prng_seed,
        depth=depth,
    )

    # Stage 3: ENCODE (use bilateral algorithm if negotiated, else config default)
    alg = encoding_algorithm if encoding_algorithm is not None else config.encoding_algorithm
    encoding_ctx = build_encoding_context(
        bilateral_state.ratchet_state,
        channel_id,
        bilateral_state.step,
        algorithm_id=alg,
    )
    encoded = encode(payload, encoding_ctx)
    protocol_msg = assemble_message(frame, encoded)

    # Stage 4: VALIDATE (self-check)
    if not validate_frame(protocol_msg.frame_open, protocol_msg.frame_close, frame):
        raise FrameValidationError("Cross-runtime frame self-check failed")

    # Stage TRANSIT: wrap in wire envelope
    header, wire_bytes = transit_stage(
        protocol_msg, channel_id, bilateral_state.step, depth, sender_runtime,
    )

    # Compute new bilateral state
    new_state = advance_bilateral_state(bilateral_state, frame)

    return CrossRuntimeResult(
        receipt=Receipt(
            message_id=message_id,
            channel=channel_id,
            step=bilateral_state.step,
            correlation_id=correlation_id,
        ),
        wire_bytes=wire_bytes,
        header=header,
        frame=frame,
        new_bilateral_state=new_state,
    )
