"""Chapter 4 — VAD + pre-roll.

Replace chapter 3's fixed silence timeout with a real voice-activity
detector plus a pre-roll ring buffer. The same parrot loop, now gated
on VAD turn boundaries instead of "500 ms since the last STT event."

Run with ``--no-preroll`` to compare a stream that omits cached audio
received before VAD-on.

Dependencies:
    uv sync --extra quickstart --extra deepgram --group dev
    export OPENAI_API_KEY=...      # OpenAI TTS
    export DEEPGRAM_API_KEY=...    # Streaming STT
    uv run easycat doctor
    uv run easycat doctor --env-file .env         # if keys live in .env
    uv run easycat doctor --env-file .env --json  # for parseable checks
    Add `--env-file .env` after `uv run` on script commands if keys live in `.env`.
"""

from __future__ import annotations

import argparse
import asyncio
import collections
import os
import time
import types
from contextlib import AsyncExitStack
from pathlib import Path

from easycat import LocalTransportConfig
from easycat.audio_format import PCM16_MONO_24K, AudioChunk
from easycat.debug.export import export_debug_bundle
from easycat.events import EventBus, STTEventType, VADStartSpeaking, VADStopSpeaking
from easycat.recipes import speak
from easycat.runtime import InMemoryRingBuffer, JournalRecordKind
from easycat.runtime.capabilities import close_if_supported
from easycat.stt.factory import STTProviderConfig, create_stt_provider
from easycat.transports.local import LocalTransport
from easycat.vad import VADConfig
from easycat.vad.factory import create_vad

PREROLL_FRAMES = 15  # 15 × 20 ms = 300 ms of audio *before* VAD fires
RUNS_DIR = Path(__file__).parent / "runs"


class MiniTurnDetector:
    """Tiny turn detector: VAD + pre-roll buffer.

    Consumes raw audio chunks, yields tagged events:

        ("speech_started", None)   - once per turn, at VAD-on.
        ("frame",          chunk)  - cached pre-roll first, then live speech.
        ("speech_ended",   None)   - once per turn, at VAD-off.

    About 40 lines of real logic. EasyCat's production ``TurnManager``
    (``src/easycat/turn_manager.py``) is a 5-state FSM with far more
    responsibilities (bot-speech overlap, cancellation, actions); read
    it once you understand why each extra state is there.
    """

    def __init__(self, vad, preroll_frames: int = PREROLL_FRAMES) -> None:
        self._vad = vad
        self._preroll: collections.deque[AudioChunk] = collections.deque(maxlen=preroll_frames)
        self._speaking = False

    async def frames(self, audio_iter):
        async for chunk in audio_iter:
            vad_events = [ev async for ev in self._vad.process(chunk)]

            for ev in vad_events:
                if isinstance(ev, VADStartSpeaking):
                    self._speaking = True
                    # Starting the turn is a state event, independent of
                    # whether any pre-roll frames exist. This matters when
                    # ``preroll_frames=0``: STT must still start before the
                    # current live frame is emitted below.
                    yield "speech_started", None

                    # Flush cached frames in their original order so STT
                    # sees the sounds that arrived before VAD fired.
                    while self._preroll:
                        yield "frame", self._preroll.popleft()
                elif isinstance(ev, VADStopSpeaking):
                    self._speaking = False
                    yield "speech_ended", None

            if self._speaking:
                yield "frame", chunk
            else:
                self._preroll.append(chunk)


def record_delivery(
    journal: InMemoryRingBuffer,
    *,
    session_id: str,
    text: str,
    accepted_chunks: int,
    rejected_chunks: int,
) -> None:
    """Preserve transport acceptance without claiming speaker playback."""
    journal.append(
        kind=JournalRecordKind.EVENT,
        name="parrot.delivery",
        session_id=session_id,
        data={
            "stage": "parrot",
            "committed_text": text,
            "accepted_chunks": accepted_chunks,
            "rejected_chunks": rejected_chunks,
            "t_ms": time.monotonic() * 1000,
        },
    )
    if rejected_chunks:
        print(
            "  transport rejected "
            f"{rejected_chunks}/{accepted_chunks + rejected_chunks} audio chunks"
        )


async def parrot(
    transport,
    stt_factory,
    detector: MiniTurnDetector,
    journal: InMemoryRingBuffer,
    session_id: str,
) -> None:
    """On each VAD turn, stream audio into STT, wait for final, speak it."""
    stt = None
    collected_final = ""

    try:
        async for tag, chunk in detector.frames(transport.receive_audio()):
            if tag == "speech_started":
                if stt is None:
                    stt = stt_factory()
                    await stt.start_stream()
                    collected_final = ""
                    journal.append(
                        kind=JournalRecordKind.EVENT,
                        name="turn.started",
                        session_id=session_id,
                        data={"stage": "turn", "t_ms": time.monotonic() * 1000},
                    )

            elif tag == "frame" and stt is not None:
                await stt.send_audio(chunk)

            elif tag == "speech_ended" and stt is not None:
                # Drain the event queue until the sentinel from end_stream().
                # A VADStop before STT saw any speech is harmless — we just
                # close an empty stream and get no FINAL back.
                active_stt = stt
                try:
                    await active_stt.end_stream()
                    async for event in active_stt.events():
                        if event.type == STTEventType.FINAL:
                            collected_final = event.text
                finally:
                    stt = None
                    await close_if_supported(active_stt)

                journal.append(
                    kind=JournalRecordKind.EVENT,
                    name="turn.ended",
                    session_id=session_id,
                    data={
                        "stage": "turn",
                        "t_ms": time.monotonic() * 1000,
                        "text": collected_final,
                    },
                )

                if collected_final.strip():
                    print(f"  → parrot: {collected_final!r}")
                    accepted_chunks, rejected_chunks = await speak(transport, collected_final)
                    record_delivery(
                        journal,
                        session_id=session_id,
                        text=collected_final,
                        accepted_chunks=accepted_chunks,
                        rejected_chunks=rejected_chunks,
                    )
    finally:
        if stt is not None:
            try:
                await stt.end_stream()
            finally:
                await close_if_supported(stt)


async def main() -> None:
    parser = argparse.ArgumentParser()
    parser.add_argument(
        "--no-preroll",
        action="store_true",
        help="Disable pre-roll; omit cached frames received before VAD-on.",
    )
    args = parser.parse_args()

    oai_key = os.getenv("OPENAI_API_KEY")
    dg_key = os.getenv("DEEPGRAM_API_KEY")
    if not oai_key or not dg_key:
        raise SystemExit("Set OPENAI_API_KEY (TTS) and DEEPGRAM_API_KEY (STT).")

    preroll = 0 if args.no_preroll else PREROLL_FRAMES
    session_id = f"ch04-vad-{'nopreroll' if args.no_preroll else 'preroll'}-{int(time.time())}"
    print(f"Pre-roll: {preroll * 20} ms" if preroll else "Pre-roll: OFF")

    journal = InMemoryRingBuffer(capacity=10_000)
    transport = LocalTransport(LocalTransportConfig(audio_format=PCM16_MONO_24K))

    def stt_factory():
        return create_stt_provider(
            STTProviderConfig(
                provider="deepgram",
                api_key=dg_key,
                params={"sample_rate": 24000, "event_bus": EventBus()},
            )
        )

    async with AsyncExitStack() as resources:
        resources.push_async_callback(transport.disconnect)
        await transport.connect()

        vad = create_vad(VADConfig())
        resources.push_async_callback(close_if_supported, vad)
        detector = MiniTurnDetector(vad, preroll_frames=preroll)

        print("Speak. The bot parrots back after each VAD turn. Ctrl-C to stop.")
        try:
            await parrot(transport, stt_factory, detector, journal, session_id)
        except (KeyboardInterrupt, asyncio.CancelledError):
            pass

    RUNS_DIR.mkdir(exist_ok=True)
    bundle_path = RUNS_DIR / f"{session_id}.bundle"
    session_stub = types.SimpleNamespace(journal=journal)
    export_debug_bundle(session_stub, bundle_path, overwrite=True)
    print(f"\nWrote bundle → {bundle_path.relative_to(Path.cwd())}")


if __name__ == "__main__":
    try:
        asyncio.run(main())
    except KeyboardInterrupt:
        pass
