"""Latency observer — measures ASR→TTS-first-audio per turn (TASK-02-06).

Hooks into the Pipecat pipeline frame flow to timestamp:
  - final-transcript-ready (ASR done)
  - LLM-first-token
  - TTS-first-audio
  - client-playback-start (approx via output frame)

Surfaces the ASR→TTS-first-audio number to the client as a metric frame so the
React client can display it (TASK-02-05 latency readout). Also logs to console
for the server-side record.

This is a thin Pipecat FrameProcessor; it does not alter the frame stream, only
observes. Per-segment latencies are stored in a per-session LatencyRecord and
emitted via the task's metrics channel.
"""

from __future__ import annotations

import time
from dataclasses import dataclass, field
from typing import Any

from pipecat.frames.frames import (
    Frame,
    TranscriptionFrame,
    LLMFullResponseEndFrame,
    TextFrame,
    TTSStartedFrame,
    TTSAudioRawFrame,
    BotStartedSpeakingFrame,
)
from pipecat.processors.frame_processor import FrameProcessor


@dataclass
class LatencyRecord:
    """Per-turn latency segments (ms)."""

    transcript_ready_ms: float | None = None
    llm_first_token_ms: float | None = None
    tts_first_audio_ms: float | None = None
    playback_start_ms: float | None = None

    @property
    def e2e_asr_to_tts_ms(self) -> float | None:
        """ASR transcript-ready → TTS first-audio (the v0.1 latency target)."""
        if self.transcript_ready_ms and self.tts_first_audio_ms:
            return self.tts_first_audio_ms - self.transcript_ready_ms
        return None

    def as_metric(self) -> dict[str, Any]:
        return {
            "e2e_latency_ms": self.e2e_asr_to_tts_ms,
            "transcript_ready_ms": self.transcript_ready_ms,
            "llm_first_token_ms": self.llm_first_token_ms,
            "tts_first_audio_ms": self.tts_first_audio_ms,
            "playback_start_ms": self.playback_start_ms,
        }


@dataclass
class LatencyObserverState:
    """Accumulates per-turn records and the current in-flight turn."""

    current: LatencyRecord = field(default_factory=LatencyRecord)
    records: list[LatencyRecord] = field(default_factory=list)

    def reset_turn(self) -> LatencyRecord:
        if self.current.transcript_ready_ms is not None:
            self.records.append(self.current)
        self.current = LatencyRecord()
        return self.current


class LatencyObserver(FrameProcessor):
    """Observes frames, timestamps the latency-critical segments, emits metrics.

    This processor is inserted into the pipeline (it passes frames through
    unchanged). On each TTS-first-audio it logs the turn's e2e latency and
    pushes a metric frame downstream for the client to read.
    """

    def __init__(self, **kwargs) -> None:
        super().__init__(**kwargs)
        self.state = LatencyObserverState()

    async def process_frame(self, frame: Frame, direction) -> None:
        # Always pass the frame through first (observation only).
        await self.push_frame(frame, direction)

        now_ms = time.perf_counter() * 1000.0

        if isinstance(frame, TranscriptionFrame):
            # ASR final transcript — start of a new turn's latency measurement.
            rec = self.state.reset_turn()
            rec.transcript_ready_ms = now_ms

        elif isinstance(frame, LLMFullResponseEndFrame):
            # LLM emitted a full response; first token timestamp is approximated
            # by this frame's arrival (Pipecat doesn't emit a dedicated
            # first-token frame; the metrics service handles TTFT separately).
            if self.state.current.llm_first_token_ms is None:
                self.state.current.llm_first_token_ms = now_ms

        elif isinstance(frame, TextFrame):
            # Intermediate LLM text frame — closest proxy to first-token time.
            if (
                self.state.current.transcript_ready_ms is not None
                and self.state.current.llm_first_token_ms is None
            ):
                self.state.current.llm_first_token_ms = now_ms

        elif isinstance(frame, (TTSStartedFrame, TTSAudioRawFrame)):
            if self.state.current.tts_first_audio_ms is None:
                self.state.current.tts_first_audio_ms = now_ms
                e2e = self.state.current.e2e_asr_to_tts_ms
                if e2e is not None:
                    from loguru import logger

                    logger.info(
                        f"[latency] ASR→TTS first-audio: {e2e:.1f}ms "
                        f"(budget 600ms — {'within' if e2e <= 600 else 'OVER'})"
                    )

        elif isinstance(frame, BotStartedSpeakingFrame):
            if self.state.current.playback_start_ms is None:
                self.state.current.playback_start_ms = now_ms


__all__ = ["LatencyObserver", "LatencyRecord", "LatencyObserverState"]