"""TransportStage — wraps a Transport with journal recording and capture."""

from __future__ import annotations

import logging
import time
from typing import Any

from easycat import _observability as observability
from easycat._turn_context import TurnContext
from easycat.runtime.context import RunContext
from easycat.runtime.replay import ReplayCassette, ReplayFidelity, ReplaySpec
from easycat.stages.base import (
    ControlSignal,
    StageStateSnapshot,
    audio_capture_allowed,
    audio_format_fields,
    captures_verbose_stage_io,
    journal_append_control_signal,
    journal_append_event_async,
    journal_ctx,
    live_replay_input,
    put_artifact_async,
    record_stage_failure,
)

logger = logging.getLogger(__name__)


class TransportStage:
    """Stage wrapper around a :class:`Transport`.

    ``execute`` wraps the outbound ``send_audio`` call: the chunk's
    bytes are stored as a ``replay_critical`` artifact and attached to
    ``stage_complete`` via ``output_ref``, so a replay can reconstruct
    exactly what the transport pushed toward the client.
    """

    name = "transport"

    def __init__(self, provider: Any, *, journal: Any = None) -> None:
        self._provider = provider
        self._journal = journal

    async def execute(self, input: Any, ctx: RunContext, turn: TurnContext) -> bool:
        """Send audio through the wrapped transport.

        Returns ``True`` when the transport accepted the chunk for
        delivery and ``False`` when it was silently dropped (for
        example because the peer is disconnected). Legacy transports
        that return ``None`` after a successful send are treated as
        delivered for backward compatibility.
        """
        ctx = journal_ctx(ctx, self._journal)
        capture_enabled = captures_verbose_stage_io(ctx)
        started = time.perf_counter()
        result_attr = "pass"
        state_before = self.snapshot_state() if capture_enabled else None
        audio_bytes = getattr(input, "data", None) if not isinstance(input, bytes) else input
        start_sequence = None
        output_ref = None
        extra: dict[str, Any] | None = None
        if capture_enabled:
            output_ref = await put_artifact_async(
                ctx,
                audio_bytes,
                capture_allowed=audio_capture_allowed(ctx, input),
            )
            extra = {
                "audio_bytes": (
                    len(audio_bytes) if isinstance(audio_bytes, (bytes, bytearray)) else 0
                ),
            }
            extra.update(audio_format_fields(input))
            start_sequence = await journal_append_event_async(
                ctx,
                stage=self.name,
                name="stage_start",
                turn_id=turn.id,
                state_before=state_before,
                data_extra=extra,
            )
        metrics_available = observability.metrics_available()
        try:
            span = observability.span(
                "easycat.transport.send",
                {"easycat.stage": self.name, "easycat.surface": "tts"},
            )
            with span:
                if metrics_available and isinstance(audio_bytes, (bytes, bytearray)):
                    observability.increment_counter(
                        "easycat.audio.bytes.total",
                        value=len(audio_bytes),
                        attributes={"easycat.surface": "tts"},
                    )
                    observability.increment_counter(
                        "easycat.audio.frames.total",
                        attributes={"easycat.surface": "tts"},
                    )
                delivered = await self._provider.send_audio(input)
                result = True if delivered is None else bool(delivered)
                if not result:
                    # A falsy return is a silent drop, not a successful send;
                    # tag the latency sample so drops are distinguishable from
                    # deliveries in easycat.stage.latency. Legacy implicit
                    # ``None`` returns still mean delivered.
                    result_attr = "drop"
        except Exception as exc:
            result_attr = "fail"
            elapsed_ms = (time.perf_counter() - started) * 1000
            record_stage_failure(
                exc,
                ctx,
                stage=self.name,
                provider=type(self._provider).__name__.lower(),
                surface="tts",
                elapsed_ms=elapsed_ms,
                sequence=start_sequence,
                turn_id=turn.id,
                state_before=state_before,
            )
            raise
        finally:
            if metrics_available:
                observability.record_histogram(
                    "easycat.stage.latency",
                    time.perf_counter() - started,
                    {"easycat.stage": self.name, "easycat.result": result_attr},
                )
        if capture_enabled:
            state_after = self.snapshot_state()
            complete_extra = {
                **(extra or {}),
                "delivered": result,
                "elapsed_ms": (time.perf_counter() - started) * 1000,
            }
            await journal_append_event_async(
                ctx,
                stage=self.name,
                name="stage_complete",
                turn_id=turn.id,
                state_before=state_before,
                state_after=state_after,
                output_ref=output_ref,
                data_extra=complete_extra,
            )
        return result

    def snapshot_state(self) -> StageStateSnapshot:
        return StageStateSnapshot(
            stage_name=self.name,
            fields={"provider": type(self._provider).__name__},
        )

    def replay(
        self,
        spec: ReplaySpec,
        cassette: ReplayCassette | None = None,
    ) -> Any:
        """Replay Transport stage.

        Transport replay is a passthrough cassette. The transport stage
        only ever sends *outbound* frames, which ``execute`` captures
        under the ``stage_complete`` ``output_ref``. Both ``ARTIFACT``
        and ``LIVE`` therefore return those captured outbound frames;
        ``LIVE`` exists so the caller can push them through a fresh
        transport (for offline analysis, not a real socket).
        """
        overrides = spec.overrides
        if spec.fidelity is ReplayFidelity.LIVE:
            return live_replay_input(
                spec, cassette, record_name="stage_complete", source="output_ref"
            )

        if "data" in overrides or "result" in overrides:
            return overrides.get("data", overrides.get("result"))
        if cassette is not None:
            record = cassette.last_record("stage_complete") or cassette.last_record()
            if record is not None:
                blob = cassette.blob(record.get("output_ref"))
                if blob is not None:
                    return blob
                data = record.get("data") or {}
                if isinstance(data, dict):
                    for key in ("data", "result"):
                        if key in data:
                            return data[key]
        return None

    async def handle_upstream(
        self,
        signal: ControlSignal,
        ctx: RunContext | None = None,
    ) -> None:
        """Observe and journal an upstream control signal (no cancel here).

        Cancellation is owned by the ``CancelOrchestrator`` / turn runner;
        this method only records that the stage saw the signal.
        """
        logger.debug("TransportStage received upstream signal: %s", signal)
        if ctx is not None:
            journal_append_control_signal(
                journal_ctx(ctx, self._journal), stage=self.name, signal=signal
            )
