"""Bell beside your agent's own speech-to-text, as a Pipecat processor.

Put it in the pipeline right after your speech-to-text service. It sends that service's words to the Rabel API
(`WS /v1/conversation`) and passes Bell's guesses to Pipecat:

  InterimTranscriptionFrame (the turn's words so far)  -> {"type": "heard", "text": ...}
  TranscriptionFrame (the turn's final words)          -> {"type": "end_of_turn", "text": ...}
  Bell's guess                                         -> EagerTranscriptionFrame: Pipecat starts the reply's model call on it

Every frame passes through unchanged; your speech-to-text still decides where the turn ends.

The guesses go through Pipecat's own early-reply path (EagerUserTurnStrategies, Pipecat 1.15), the one Deepgram Flux's
EagerEndOfTurn uses, so the user aggregator needs `user_turn_strategies=EagerUserTurnStrategies()`:

- the reply made from a guess is held by the LLM service until the turn ends;
- it is kept only when the turn's final words match the guess (NormalizedMatch: case and punctuation ignored);
- otherwise it is dropped and the reply is made again from the final words;
- Pipecat holds one reply made from a guess at a time: the newest guess withdraws the reply made from the older one,
  except when the older guess equals the words heard so far; and when the words heard reach a guess Bell made
  earlier in the turn, the reply is started on that guess;
- a held reply that asks for a tool call is cancelled by Pipecat, so a guess never triggers an action.
"""
# Where it goes, as Pipecat's Smart Turn page puts its model beside the speech-to-text (the turn strategy in the user
# aggregator's params): https://docs.pipecat.ai/api-reference/server/utilities/turn-detection/smart-turn-overview
import asyncio
import json
import os

import websockets
from loguru import logger
from pipecat.frames.frames import (
    CancelFrame,
    EagerTranscriptionFrame,
    EndFrame,
    Frame,
    InterimTranscriptionFrame,
    StartFrame,
    TranscriptionFrame,
    UserTurnResponseCancelFrame,
)
from pipecat.processors.frame_processor import FrameDirection, FrameProcessor
from pipecat.turns.user_stop import NormalizedMatch
from pipecat.utils.time import time_now_iso8601

RABEL_API_URL = "https://api.rabellabs.com"


class RabelBellProcessor(FrameProcessor):
    """Sends your speech-to-text's words to the Rabel API and starts the reply on Bell's guesses.

    Args:
        api_key: Your Rabel API key (defaults to the RABEL_API_KEY environment variable).
        url: The API's base address.
    """

    def __init__(self, *, api_key: str | None = None, url: str = RABEL_API_URL, **kwargs):
        super().__init__(**kwargs)
        self._api_key = api_key or os.environ.get("RABEL_API_KEY", "")
        self._url = url.rstrip("/").replace("https://", "wss://").replace("http://", "ws://") + "/v1/conversation"
        self._ws = None
        self._receive_task = None
        self._user_id = ""
        self._in_turn = False
        self._guesses, self._speculated, self._heard = [], None, ""
        self.ready = asyncio.Event()   # set when Bell is awake and the API takes words
        self.events = []               # (monotonic time, message) for every message, sent ("sent", ...) or received
        self.held_guesses = []         # per turn: the guess whose reply was held when the turn ended (or None)

    async def process_frame(self, frame: Frame, direction: FrameDirection):
        await super().process_frame(frame, direction)
        if isinstance(frame, StartFrame):
            await self.push_frame(frame, direction)
            self._ws = await websockets.connect(self._url, max_size=None)
            await self._ws.send(json.dumps({"key": self._api_key}))
            self._receive_task = self.create_task(self._receive())
            return
        if isinstance(frame, (EndFrame, CancelFrame)):
            await self._close()
        elif isinstance(frame, InterimTranscriptionFrame) and frame.text:
            await self._on_heard(frame)
        elif isinstance(frame, TranscriptionFrame) and frame.text:
            await self._on_final(frame)
        await self.push_frame(frame, direction)

    async def _send(self, kind: str, text: str):
        if self._ws is not None and self.ready.is_set():
            self.events.append((asyncio.get_running_loop().time(), {"type": "sent " + kind, "text": text}))
            await self._ws.send(json.dumps({"type": kind, "text": text}))

    async def _on_heard(self, frame: InterimTranscriptionFrame):
        if not self._in_turn:
            self._in_turn = True
            self._guesses, self._speculated, self._heard = [], None, ""
        self._user_id = frame.user_id
        self._heard = frame.text
        await self._send("heard", frame.text)
        # The words heard have reached a guess Bell made earlier in the turn: answer that guess.
        earlier = next((g for g in self._guesses if self._same(g, self._heard)), None)
        if earlier and not self._same(earlier, self._speculated):
            await self._speculate(earlier)

    async def _on_final(self, frame: TranscriptionFrame):
        self._in_turn = False
        self.held_guesses.append(self._speculated)   # the guess whose reply was held at the turn's end
        self._speculated = None                       # the turn's end settles it, through Pipecat's match
        await self._send("end_of_turn", frame.text)

    @staticmethod
    def _same(a: str | None, b: str | None) -> bool:
        return bool(a) and bool(b) and NormalizedMatch().matches(a, b)

    async def _speculate(self, text: str):
        """Withdraw the reply made from the previous guess, if any, and start one on this guess."""
        if self._speculated is not None:
            await self.push_frame(UserTurnResponseCancelFrame())
        self._speculated = text
        await self.push_frame(EagerTranscriptionFrame(text, self._user_id, time_now_iso8601()))

    async def _receive(self):
        loop = asyncio.get_running_loop()
        async for raw in self._ws:
            m = json.loads(raw)
            self.events.append((loop.time(), m))
            kind = m.get("type")
            if kind == "ready":
                self.ready.set()
            elif kind == "guess" and self._in_turn:
                self._guesses.append(m["text"])
                # A guess equal to the words heard so far stays; otherwise the newest guess replaces the older one.
                if not (self._speculated and self._same(self._speculated, self._heard)):
                    await self._speculate(m["text"])
            elif kind == "error":
                logger.error(f"{self}: {m.get('error')}: {m.get('message')}")

    async def _close(self):
        if self._receive_task:
            await self.cancel_task(self._receive_task)
            self._receive_task = None
        if self._ws:
            try:
                await self._ws.close()
            except Exception:  # noqa: BLE001 -- already gone
                pass
            self._ws = None