"""VADStage — wraps a VADProvider with journal recording and input capture."""

from __future__ import annotations

import contextlib
import dataclasses
import logging
import time
from typing import Any

from easycat import _observability as observability
from easycat._audio_utils import AudioFrameAligner, to_mono_chunk
from easycat._turn_context import TurnContext
from easycat.audio_format import AudioChunk
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,
    set_audio_capture_allowed,
)

logger = logging.getLogger(__name__)


def _serialize_event(event: Any) -> dict[str, Any]:
    """Serialize a VAD event into a replay-reconstructable descriptor.

    Records the event's class name under ``type`` and, when the event is
    a dataclass, its fields (timestamps, correlation ids, speech/silence
    boundaries, probabilities) so ARTIFACT replay can rebuild the event
    rather than recovering only its class name.
    """
    descriptor: dict[str, Any] = {"type": type(event).__name__}
    if dataclasses.is_dataclass(event) and not isinstance(event, type):
        descriptor["fields"] = dataclasses.asdict(event)
    return descriptor


class VADStage:
    """Stage wrapper around a :class:`VADProvider`.

    ``execute`` runs the provider on a single chunk, captures the raw
    input bytes as an ``input_ref`` artifact for LIVE-fidelity replay,
    and returns the list of events emitted.  ``stage_complete`` carries
    a serialized descriptor per event in ``data["events"]`` — each a
    dict of the event's ``type`` plus its dataclass fields (timestamps,
    correlation ids, etc.) — so ARTIFACT replay can reconstruct the VAD
    output without re-running the model.
    """

    name = "vad"

    _MAX_EVENTS_PER_CHUNK: int = 64

    def __init__(self, provider: Any, *, journal: Any = None) -> None:
        self._provider = provider
        self._journal = journal
        # Most recent VAD decision (last emitted event type), surfaced via
        # ``snapshot_state`` so ``replay_decision`` returns something real.
        self._last_decision: Any = None
        self._input_frame_aligner = AudioFrameAligner()

    def _normalize_audio_input(
        self,
        input: Any,
        *,
        capture_allowed: bool,
    ) -> Any:
        """Frame-align and downmix provider-facing PCM16 audio."""
        if not isinstance(input, AudioChunk):
            return input
        normalized_input = self._input_frame_aligner.align(input)
        if normalized_input.format.channels > 1:
            normalized_input = to_mono_chunk(normalized_input)
            set_audio_capture_allowed(normalized_input, capture_allowed)
        return normalized_input

    async def execute(self, input: Any, ctx: RunContext, turn: TurnContext) -> Any:
        ctx = journal_ctx(ctx, self._journal)
        if isinstance(input, AudioChunk) and (
            input.format.encoding != "pcm" or input.format.sample_width != 2
        ):
            raise ValueError(
                "VAD requires PCM16 audio "
                f"(got encoding={input.format.encoding!r}, "
                f"sample_width={input.format.sample_width!r})"
            )
        started = time.perf_counter()
        result_attr = "pass"
        state_before = self.snapshot_state()
        capture_detail = captures_verbose_stage_io(ctx)
        data_bytes = getattr(input, "data", None) if not isinstance(input, bytes) else input
        # ``input_ref`` stores the bytes exactly as they arrived. Keep the
        # matching geometry before the provider-facing mono normalization so
        # debugger what-if runs decode the artifact with its true channel
        # layout instead of treating interleaved stereo as mono.
        captured_audio_format = audio_format_fields(input)
        capture_allowed = audio_capture_allowed(ctx, input)
        input_ref = (
            await put_artifact_async(
                ctx,
                data_bytes,
                capture_allowed=capture_allowed,
            )
            if capture_detail
            else None
        )
        # VAD backends decode the raw byte stream as flat int16 mono (frame
        # boundaries are computed as samples*2). Align source frames before
        # downmixing so a transport split in the middle of an interleaved
        # frame cannot lose or cross-pair channel samples. We do this after
        # capturing the input_ref artifact so replay still reflects the true
        # raw input.
        provider_input = self._normalize_audio_input(
            input,
            capture_allowed=capture_allowed,
        )
        start_sequence: int | None = None
        if capture_detail:
            start_extra = {
                "audio_bytes": len(data_bytes)
                if isinstance(data_bytes, (bytes, bytearray))
                else 0,
            }
            start_extra.update(captured_audio_format)
            start_sequence = await journal_append_event_async(
                ctx,
                stage=self.name,
                name="stage_start",
                turn_id=turn.id,
                state_before=state_before,
                input_ref=input_ref,
                data_extra=start_extra,
            )
        try:
            with observability.span(
                "easycat.vad.detect",
                {"easycat.stage": self.name, "easycat.surface": "stt"},
            ):
                events: list[Any] = []
                stream = self._provider.process(provider_input)
                try:
                    async for event in stream:
                        events.append(event)
                        if len(events) >= self._MAX_EVENTS_PER_CHUNK:
                            logger.warning(
                                "VAD provider yielded too many events for a single chunk; "
                                "truncating"
                            )
                            break
                finally:
                    # The event cap deliberately leaves an over-producing
                    # provider iterator early.  ``async for`` does not close
                    # a retained async generator on ``break``, so forward the
                    # close explicitly to release its stream/session promptly.
                    aclose = getattr(stream, "aclose", None)
                    if callable(aclose):
                        with contextlib.suppress(Exception):
                            await aclose()
                result = events
        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="stt",
                elapsed_ms=elapsed_ms,
                sequence=start_sequence,
                turn_id=turn.id,
                state_before=state_before,
            )
            raise
        finally:
            observability.record_histogram(
                "easycat.stage.latency",
                time.perf_counter() - started,
                {"easycat.stage": self.name, "easycat.result": result_attr},
            )
        if events:
            self._last_decision = type(events[-1]).__name__
        if capture_detail:
            state_after = self.snapshot_state()
            await journal_append_event_async(
                ctx,
                stage=self.name,
                name="stage_complete",
                turn_id=turn.id,
                state_before=state_before,
                state_after=state_after,
                data_extra={
                    "events": [_serialize_event(e) for e in events],
                    "elapsed_ms": (time.perf_counter() - started) * 1000,
                },
            )
        return result

    def snapshot_state(self) -> StageStateSnapshot:
        fields: dict[str, Any] = {"provider": type(self._provider).__name__}
        if self._last_decision is not None:
            fields["decision"] = self._last_decision
        return StageStateSnapshot(stage_name=self.name, fields=fields)

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

        ``ARTIFACT`` returns the captured frame events from the
        cassette (or ``spec.overrides["events"]``).  ``LIVE`` returns
        the captured input audio for re-execution against a fresh VAD
        backend.
        """
        overrides = spec.overrides
        if spec.fidelity is ReplayFidelity.LIVE:
            return live_replay_input(spec, cassette)

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

    def replay_decision(self, snapshot: StageStateSnapshot) -> Any:
        """Replay the last VAD decision recorded in a snapshot.

        ``snapshot_state`` carries the most recent emitted event type under
        ``fields["decision"]`` once the stage has run, so this returns that
        verdict (or ``None`` for a snapshot taken before any events).
        """
        return snapshot.fields.get("decision", 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("VADStage received upstream signal: %s", signal)
        if ctx is not None:
            journal_append_control_signal(
                journal_ctx(ctx, self._journal), stage=self.name, signal=signal
            )
