"""The same Pipecat voice agent with and without Bell, on the same recorded turn.

  pip install -r requirements.txt
  python compare.py                # 3 runs on recording.wav (16 kHz, mono, 16-bit)

Keys, from your environment: RABEL_API_KEY, DEEPGRAM_API_KEY, ANTHROPIC_API_KEY.

The agent: Deepgram Flux (speech-to-text; it decides where the turn ends) -> Claude Haiku 5.5 -> Deepgram Aura-2.
  without Bell: the reply's model call starts at Flux's end of turn.
  with Bell:    RabelBellProcessor right after Flux sends Flux's words to the Rabel API; the reply's model call starts on
                Bell's guess, and is kept only if Flux's final words match it.

Each run streams the recording in 80 ms chunks at real-time pace, then silence, as a caller's microphone would, and times
the reply's first audio from the speaker's last word. The speaker's last word is where Deepgram Nova-3 (prerecorded, word
timestamps) puts the end of the recording's last word.
"""
import argparse
import asyncio
import json
import os
import statistics
import sys
import time
import urllib.request
import wave

from loguru import logger
from pipecat.frames.frames import InputAudioRawFrame, TTSAudioRawFrame, TTSTextFrame
from pipecat.pipeline.pipeline import Pipeline
from pipecat.pipeline.worker import PipelineParams, PipelineWorker
from pipecat.processors.aggregators.llm_context import LLMContext
from pipecat.processors.aggregators.llm_response_universal import LLMContextAggregatorPair, LLMUserAggregatorParams
from pipecat.processors.frame_processor import FrameDirection, FrameProcessor
from pipecat.services.anthropic.llm import AnthropicLLMService
from pipecat.services.deepgram.flux.stt import DeepgramFluxSTTService
from pipecat.services.deepgram.tts import DeepgramTTSService
from pipecat.turns.user_stop import NormalizedMatch
from pipecat.turns.user_turn_strategies import EagerUserTurnStrategies
from pipecat.workers.runner import WorkerRunner

from rabel_bell_stt import RabelBellProcessor

API = "https://api.rabellabs.com"
RATE = 16000
CHUNK = 2560            # 80 ms of PCM16 mono at 16 kHz
CHUNK_S = 0.08
SYSTEM = "You are a helpful voice assistant. Answer in one short spoken sentence."


class FirstAudio(FrameProcessor):
    """Notes when the reply's first audio and its text pass, then lets every frame through."""

    def __init__(self):
        super().__init__()
        self.first_audio = None
        self.text = []
        self.done = asyncio.Event()

    async def process_frame(self, frame, direction: FrameDirection):
        await super().process_frame(frame, direction)
        if isinstance(frame, TTSAudioRawFrame) and self.first_audio is None:
            self.first_audio = asyncio.get_running_loop().time()
            self.done.set()
        if isinstance(frame, TTSTextFrame):
            self.text.append(frame.text)
        await self.push_frame(frame, direction)


def build_agent(with_bell: bool, api_url: str):
    stt = DeepgramFluxSTTService(api_key=os.environ["DEEPGRAM_API_KEY"])
    bell = RabelBellProcessor(api_key=os.environ["RABEL_API_KEY"], url=api_url) if with_bell else None
    llm = AnthropicLLMService(
        api_key=os.environ["ANTHROPIC_API_KEY"],
        settings=AnthropicLLMService.Settings(model="claude-haiku-5-5", max_tokens=80, effort="low",
                                              system_instruction=SYSTEM),
    )
    tts = DeepgramTTSService(api_key=os.environ["DEEPGRAM_API_KEY"],
                             settings=DeepgramTTSService.Settings(voice="aura-2-thalia-en"))
    context = LLMContext()
    # With Bell, the user aggregator answers Bell's guesses early (Pipecat's EagerUserTurnStrategies).
    user_params = LLMUserAggregatorParams(user_turn_strategies=EagerUserTurnStrategies()) if with_bell else None
    user_aggregator, assistant_aggregator = LLMContextAggregatorPair(context, user_params=user_params)
    tap = FirstAudio()
    pipeline = Pipeline([stt] + ([bell] if bell else []) + [user_aggregator, llm, tts, tap, assistant_aggregator])
    worker = PipelineWorker(pipeline, params=PipelineParams(audio_in_sample_rate=RATE, audio_out_sample_rate=RATE),
                            idle_timeout_secs=None)
    return stt, bell, tap, worker


async def one_turn(with_bell: bool, audio: bytes, last_word_s: float, api_url: str):
    """One conversation, one turn. Returns the run's numbers."""
    stt, bell, tap, worker = build_agent(with_bell, api_url)
    loop = asyncio.get_running_loop()
    end_of_turn = {}
    connected = asyncio.Event()

    @stt.event_handler("on_connected")
    async def _connected(_):
        connected.set()

    @stt.event_handler("on_end_of_turn")
    async def _eot(_, transcript):
        end_of_turn.setdefault("t", loop.time())
        end_of_turn.setdefault("text", transcript)

    runner = WorkerRunner(handle_sigint=False)
    await runner.add_workers(worker)
    run = asyncio.create_task(runner.run())
    t_wait = loop.time()
    await asyncio.wait_for(connected.wait(), 120)
    if bell:
        await asyncio.wait_for(bell.ready.wait(), 120)  # Bell can take up to a minute to wake on a first call
    wake_s = loop.time() - t_wait

    silence = bytes(CHUNK)
    t0 = loop.time()
    i = 0
    total = len(audio) // CHUNK + int(len(audio) % CHUNK > 0)
    while True:                                    # the recording, then silence, at real-time pace
        if i < total:
            chunk = audio[i * CHUNK:(i + 1) * CHUNK]
        elif tap.done.is_set() or loop.time() - t0 > total * CHUNK_S + 10:
            break
        else:
            chunk = silence
        await worker.queue_frame(InputAudioRawFrame(audio=chunk, sample_rate=RATE, num_channels=1))
        i += 1
        await asyncio.sleep(max(0.0, t0 + i * CHUNK_S - loop.time()))
    await asyncio.sleep(2.0)                       # let the reply's text finish
    await worker.cancel()
    try:
        await asyncio.wait_for(run, 20)
    except asyncio.TimeoutError:
        logger.warning("the pipeline did not finish within 20 s of its cancel; going on")

    last_word = t0 + last_word_s
    out = {"agent": "with Bell" if with_bell else "without Bell", "wake_s": round(wake_s, 2),
           "ms": round((tap.first_audio - last_word) * 1000) if tap.first_audio else None,
           "reply": "".join(tap.text).strip()}
    if end_of_turn:
        out.update(end_of_turn_text=end_of_turn["text"], end_of_turn_ms=round((end_of_turn["t"] - last_word) * 1000))
    if bell:
        guesses = [m for _, m in bell.events if m.get("type") == "guess"]
        bell_eot = next((m for _, m in bell.events if m.get("type") == "end_of_turn"), None)
        held = bell.held_guesses[0] if bell.held_guesses else None
        out.update(guesses=len(guesses), held_guess=held,
                   reply_from_guess=bool(held and end_of_turn and NormalizedMatch().matches(held, end_of_turn["text"])),
                   bell_guess_id=bell_eot.get("guess_id") if bell_eot else None,
                   bell_matched=bool(bell_eot and bell_eot.get("guess_id") in {g.get("id") for g in guesses}),
                   messages=[[round((t - last_word) * 1000), m.get("type"), m.get("text"), m.get("id", m.get("guess_id"))]
                             for t, m in bell.events])
    return out


def last_word_end(wav_bytes: bytes) -> tuple[float, str]:
    """Where the recording's last word ends, in seconds (Deepgram Nova-3, prerecorded, word timestamps)."""
    req = urllib.request.Request("https://api.deepgram.com/v1/listen?model=nova-3", data=wav_bytes,
                                 headers={"Authorization": "Token " + os.environ["DEEPGRAM_API_KEY"],
                                          "Content-Type": "audio/wav", "User-Agent": "rabel-api-example"})
    r = json.loads(urllib.request.urlopen(req, timeout=60).read())
    words = r["results"]["channels"][0]["alternatives"][0]["words"]
    return float(words[-1]["end"]), " ".join(w["word"] for w in words)


async def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--runs", type=int, default=3)
    ap.add_argument("--recording", default="recording.wav", help="the recorded turn: 16 kHz, mono, 16-bit WAV")
    ap.add_argument("--url", default=API)
    ap.add_argument("--json", help="write every run's numbers to this file")
    ap.add_argument("--verbose", action="store_true")
    a = ap.parse_args()
    logger.remove()
    logger.add(sys.stderr, level="DEBUG" if a.verbose else "WARNING")

    with open(a.recording, "rb") as f:
        wav_bytes = f.read()
    w = wave.open(a.recording)
    if (w.getframerate(), w.getnchannels(), w.getsampwidth()) != (RATE, 1, 2):
        raise SystemExit(f"{a.recording} must be 16 kHz, mono, 16-bit")
    audio = w.readframes(w.getnframes())
    last_s, words = last_word_end(wav_bytes)
    print(f"recorded turn: {len(audio) / 2 / RATE:.2f} s, last word ends at {last_s:.2f} s ({words!r})", flush=True)

    rows = []
    for r in range(1, a.runs + 1):
        for with_bell in (False, True):
            try:
                x = await asyncio.wait_for(one_turn(with_bell, audio, last_s, a.url), 240)
            except (asyncio.TimeoutError, TimeoutError):
                x = {"agent": "with Bell" if with_bell else "without Bell", "ms": None, "wake_s": None,
                     "reply_from_guess": False, "error": "no reply within 240 s"}
            x["run"] = r
            rows.append(x)
            note = f" (Bell woke in {x['wake_s']} s)" if with_bell and (x["wake_s"] or 0) > 2 else ""
            kept = "" if not with_bell else ("  reply made from Bell's guess" if x["reply_from_guess"]
                                             else "  no guess matched: reply made from the final words")
            print(f"run {r}  {x['agent']:<13} {x['ms']} ms from the last word to the reply's first audio{note}{kept}",
                  flush=True)
            if a.verbose:
                print("   ", json.dumps(x), flush=True)

    print("\nms from the speaker's last word to the reply's first audio")
    print(f"{'run':<5}{'without Bell':>14}{'with Bell':>12}{'saved':>8}")
    saved = []
    for r in range(1, a.runs + 1):
        wo = next(x["ms"] for x in rows if x["run"] == r and x["agent"] == "without Bell")
        wi = next(x["ms"] for x in rows if x["run"] == r and x["agent"] == "with Bell")
        s = wo - wi if wo is not None and wi is not None else None
        if s is not None:
            saved.append(s)
        print(f"{r:<5}{str(wo):>14}{str(wi):>12}{str(s):>8}")
    if saved:
        print(f"median saved a turn: {statistics.median(saved):.0f} ms over {len(saved)} runs")
    if a.json:
        with open(a.json, "w") as f:
            json.dump({"last_word_s": last_s, "words": words, "runs": rows,
                       "utc": time.strftime("%Y-%m-%dT%H:%M:%SZ", time.gmtime())}, f, indent=1)


if __name__ == "__main__":
    code = 0
    try:
        asyncio.run(main())
    except BaseException:  # noqa: BLE001 -- printed, then the exit code says it
        import traceback
        traceback.print_exc()
        code = 1
    sys.stdout.flush()
    sys.stderr.flush()
    os._exit(code)      # a pipeline left mid-close by a timeout never holds the command open