"""Chapter 5 — The blocking agent.

Same pipeline as chapter 4, but instead of parroting the transcript
back, we send it to an LLM and wait for the complete response before
handing it to TTS. The bot falls silent for 2-4 seconds per turn.
That silence is the whole point of this chapter.

Dependencies:
    uv sync --extra quickstart --extra deepgram --group dev
    export OPENAI_API_KEY=...
    export DEEPGRAM_API_KEY=...
    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 asyncio
import collections
import os
import time
import types
from contextlib import AsyncExitStack
from pathlib import Path

from openai import AsyncOpenAI

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
MODEL = "gpt-5.6-luna"
RUNS_DIR = Path(__file__).parent / "runs"
SESSION_ID = f"ch05-blocking-{int(time.time())}"


class MiniTurnDetector:
    """Same as chapter 4."""

    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
                    yield "speech_started", None
                    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)


class FirstAudioProbe:
    """Forward audio while recording when the first chunk is accepted."""

    def __init__(self, transport) -> None:
        self._transport = transport
        self.first_audio_at: float | None = None

    async def send_audio(self, chunk: AudioChunk) -> bool:
        accepted = await self._transport.send_audio(chunk)
        normalized = accepted is None or bool(accepted)
        if normalized and self.first_audio_at is None:
            self.first_audio_at = time.monotonic()
        return normalized


def span(journal: InMemoryRingBuffer, name: str, t0: float, **extra) -> None:
    """Record a closed span with start→end wall time in ms."""
    elapsed_ms = (time.monotonic() - t0) * 1000
    journal.append(
        kind=JournalRecordKind.EVENT,
        name=name,
        session_id=SESSION_ID,
        data={"stage": name.split(".")[1], "elapsed_ms": elapsed_ms, **extra},
    )


async def blocking_agent(client: AsyncOpenAI, user_text: str) -> str:
    """One LLM call. Wait for the full response. Return the string."""
    resp = await client.chat.completions.create(
        model=MODEL,
        reasoning_effort="none",
        messages=[
            {"role": "system", "content": "You are a helpful voice assistant. Keep it brief."},
            {"role": "user", "content": user_text},
        ],
    )
    return resp.choices[0].message.content or ""


async def run_turn(transport, stt, client, journal) -> None:
    """Finalize the current STT stream, run the LLM, speak the reply.

    The STT stream has been receiving chunks from the parent caller's
    VAD loop already — we just close it here and drain the FINAL.
    """
    final_text = ""
    stt_final_t = None
    async for event in stt.events():
        if event.type == STTEventType.FINAL:
            final_text = event.text
            stt_final_t = time.monotonic()

    if not final_text.strip() or stt_final_t is None:
        return

    print(f"  user: {final_text!r}")

    # Sub-gap 1: STT final → we start the LLM call. Just our own
    # dispatch overhead; should be under a millisecond.
    agent_dispatch = time.monotonic()
    span(
        journal,
        "stage.stt_to_agent",
        stt_final_t,
        at_ms=(agent_dispatch - stt_final_t) * 1000,
    )

    # Sub-gap 2: the LLM call itself. The biggest sub-gap — usually
    # 1-3 seconds of silence on a small model, more on a large one.
    agent_start = time.monotonic()
    reply = await blocking_agent(client, final_text)
    agent_end = time.monotonic()
    span(
        journal,
        "stage.agent.execute",
        agent_start,
        prompt=final_text,
        reply=reply,
    )

    # Sub-gap 3: agent response → the first TTS audio chunk the transport
    # accepts. ``speak`` returns after every produced chunk has been offered,
    # so a forwarding probe captures the earlier first-accepted milestone.
    tts_start = time.monotonic()
    print(f"  bot:  {reply!r}")
    audio_probe = FirstAudioProbe(transport)
    accepted_chunks, rejected_chunks = await speak(audio_probe, reply)
    tts_end = time.monotonic()
    first_audio_t = audio_probe.first_audio_at
    tts_first_audio_ms = None if first_audio_t is None else (first_audio_t - tts_start) * 1000
    tts_enqueue_ms = (tts_end - tts_start) * 1000
    span(
        journal,
        "stage.tts.execute",
        tts_start,
        text=reply,
        first_audio_ms=tts_first_audio_ms,
        enqueue_ms=tts_enqueue_ms,
        accepted_chunks=accepted_chunks,
        rejected_chunks=rejected_chunks,
    )

    total_gap = None if first_audio_t is None else (first_audio_t - stt_final_t) * 1000
    journal.append(
        kind=JournalRecordKind.EVENT,
        name="turn.gap",
        session_id=SESSION_ID,
        data={
            "stage": "turn",
            "total_gap_ms": total_gap,
            "stt_to_agent_ms": (agent_dispatch - stt_final_t) * 1000,
            "agent_ms": (agent_end - agent_start) * 1000,
            "tts_ms": tts_first_audio_ms,
            "tts_enqueue_ms": tts_enqueue_ms,
            "tts_accepted_chunks": accepted_chunks,
            "tts_rejected_chunks": rejected_chunks,
            "text": reply,
        },
    )
    if total_gap is None:
        if accepted_chunks:
            print("  (turn gap unavailable — accepted TTS audio had no timestamp)")
        elif rejected_chunks:
            print(
                f"  (turn gap unavailable — transport rejected all {rejected_chunks} TTS chunks)"
            )
        else:
            print("  (turn gap unavailable — TTS produced no audio)")
    else:
        print(f"  (turn gap: {total_gap:.0f} ms — STT final → first audio accepted)")


async def collect_turns(transport, detector, stt_factory, client, journal) -> None:
    """Stream turns and close every per-turn STT, including on cancellation."""
    stt = None
    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()
            elif tag == "frame" and stt is not None:
                await stt.send_audio(chunk)
            elif tag == "speech_ended" and stt is not None:
                active_stt = stt
                try:
                    await active_stt.end_stream()
                    await run_turn(transport, active_stt, client, journal)
                finally:
                    stt = None
                    await close_if_supported(active_stt)
    finally:
        if stt is not None:
            try:
                await stt.end_stream()
            finally:
                await close_if_supported(stt)


async def main() -> None:
    if not (os.getenv("OPENAI_API_KEY") and os.getenv("DEEPGRAM_API_KEY")):
        raise SystemExit("Set OPENAI_API_KEY and DEEPGRAM_API_KEY.")

    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=os.environ["DEEPGRAM_API_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)

        client = AsyncOpenAI()
        resources.push_async_callback(close_if_supported, client)

        print("Talk. Each turn will feel slow. That is the lesson.\n")
        try:
            await collect_turns(transport, detector, stt_factory, client, journal)
        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
