#
# Copyright (c) 2024–2025, Daily
#
# SPDX-License-Identifier: BSD 2-Clause License
#

"""Ringg Parrot speech-to-text service implementation.

This service connects to a Ringg Parrot STT WebSocket server for real-time
transcription. The server provides partial transcripts as audio streams in,
and final transcripts when explicitly requested via a stop message.
"""

import asyncio
import json
from typing import AsyncGenerator, Optional

from loguru import logger
from pydantic import BaseModel

from pipecat.frames.frames import (
    CancelFrame,
    EndFrame,
    ErrorFrame,
    Frame,
    InterimTranscriptionFrame,
    StartFrame,
    StartInterruptionFrame,
    TranscriptionFrame,
    UserStartedSpeakingFrame,
    UserStoppedSpeakingFrame,
)
from enum import Enum
from pipecat.processors.frame_processor import FrameDirection
from pipecat.services.stt_service import STTService
from pipecat.utils.time import time_now_iso8601
from pipecat.utils.tracing.service_decorators import traced_stt

try:
    import websockets
    from websockets.asyncio.client import connect as websocket_connect
    from websockets.protocol import State
except ModuleNotFoundError as e:
    logger.error(f"Exception: {e}")
    logger.error("In order to use Ringg Parrot STT, you need to `pip install websockets`.")
    raise Exception(f"Missing module: {e}")


STOP_MESSAGE = '{"type": "reset"}'

RESET_MESSAGE_V2 = '{"type": "reset"}'
END_MESSAGE_V2 = '{"type": "end"}'


class PredictionMethod(str, Enum):
    """Prediction method for Ringg Parrot STT service."""

    ON_FINAL = "on_final"


class RinggParrotInputParams(BaseModel):
    """Configuration parameters for Ringg Parrot STT service.

    Parameters:
        encoding: Audio encoding format (default: int16 for PCM signed 16-bit).
        language: Language for transcription (default: Hindi).
        punctuate: Whether to add punctuation to transcriptions.
        partial_chunk_ms: Interval in ms for partial transcription updates.
    """

    api_key: str = ""
    encoding: str = "int16"
    language: str = "hi"
    punctuate: bool = False
    partial_chunk_ms: int = 400
    prediction_method: PredictionMethod = PredictionMethod.ON_FINAL


class RinggParrotSTTService(STTService):
    """Speech-to-Text service using Ringg Parrot's WebSocket API.

    This service connects to a Ringg Parrot STT WebSocket server for real-time
    transcription. It follows a VAD-forced endpoint pattern where finalization
    is triggered when Pipecat's VAD detects the user stopped speaking.

    The server:
    - Accepts raw PCM audio frames (16-bit signed, 16kHz)
    - Sends partial transcripts automatically as audio accumulates
    - Sends final transcript when client sends a stop message

    Example usage::

        stt = RinggParrotSTTService(
            url="ws://13.234.40.75:7864/v1/audio/stream",
            language="hi",
        )

        pipeline = Pipeline([
            transport.input(),
            stt,
            context_aggregator.user(),
            llm,
            tts,
            transport.output(),
        ])
    """

    def __init__(
        self,
        *,
        url: str = "ws://13.234.40.75:8888/v1/audio/stream",
        sample_rate: Optional[int] = 16000,
        params: Optional[RinggParrotInputParams] = None,
        **kwargs,
    ):
        """Initialize the Ringg Parrot STT service.

        Args:
            url: WebSocket URL for the Ringg Parrot STT server.
            sample_rate: Audio sample rate in Hz (default: 16000).
            params: Additional configuration parameters.
            **kwargs: Additional arguments passed to the STTService.
        """
        super().__init__(sample_rate=sample_rate, **kwargs)
        params = params or RinggParrotInputParams()

        self._url = url
        self._params = params
        self._websocket = None
        self._receive_task = None
        self._client_id: Optional[str] = None

        self.set_model_name("ringg-parrot")

    async def start(self, frame: StartFrame):
        """Start the Ringg Parrot STT websocket connection.

        Connects to the WebSocket server, sends the start configuration,
        waits for ready confirmation, and starts the receive task.

        Args:
            frame: The start frame containing initialization parameters.
        """
        await super().start(frame)
        if self._websocket:
            return

        try:
            self._websocket = await websocket_connect(self._url)

            if not self._websocket:
                logger.error(f"Unable to connect to Ringg Parrot STT at {self._url}")
                await self._call_event_handler("on_connection_error", "Connection failed")
                return

            # Send the start configuration message
            start_config = {
                "type": "start",
                "sample_rate": self.sample_rate,
                "encoding": self._params.encoding,
                "language": self._params.language,
                "punctuate": self._params.punctuate,
                "partial_chunk_ms": self._params.partial_chunk_ms,
                "prediction_method": self._params.prediction_method,
                "api_key": self._params.api_key,
            }
            await self._websocket.send(json.dumps(start_config))

            # Wait for ready response
            ready_response = await asyncio.wait_for(self._websocket.recv(), timeout=10.0)
            ready_data = json.loads(ready_response)

            if ready_data.get("type") != "ready":
                error_msg = f"Unexpected response: {ready_data}"
                logger.error(f"Ringg Parrot STT: {error_msg}")
                await self._call_event_handler("on_connection_error", error_msg)
                await self._cleanup()
                return

            self._client_id = ready_data.get("client_id")
            logger.debug(f"Ringg Parrot STT connected , client_id : {self._client_id}")

            # Start the receive task
            if self._websocket and not self._receive_task:
                self._receive_task = self.create_task(self._receive_task_handler())

            await self._call_event_handler("on_connected")

        except asyncio.TimeoutError:
            logger.error("Ringg Parrot STT: Connection timeout waiting for ready response")
            await self._call_event_handler("on_connection_error", "Connection timeout")
            await self._cleanup()
        except Exception as e:
            logger.error(f"Ringg Parrot STT: Connection error: {e}")
            await self._call_event_handler("on_connection_error", str(e))
            await self._cleanup()

    async def _cleanup(self):
        """Clean up WebSocket connection and tasks."""
        if self._websocket:
            try:
                await self._websocket.close()
            except Exception:
                pass
            self._websocket = None

        if self._receive_task:
            # Task cannot cancel itself
            if self._receive_task != asyncio.current_task():
                await self._receive_task
            self._receive_task = None

        self._client_id = None
        logger.debug("Disconnected from Ringg Parrot STT")
        await self._call_event_handler("on_disconnected")

    async def stop(self, frame: EndFrame):
        """Stop the Ringg Parrot STT websocket connection.

        Called when the call/pipeline ends. Closes the connection without
        sending a stop message (final transcripts are triggered by
        UserStoppedSpeakingFrame during the call).

        Args:
            frame: The end frame.
        """
        await super().stop(frame)
        await self._cleanup()

    async def cancel(self, frame: CancelFrame):
        """Cancel the Ringg Parrot STT websocket connection immediately.

        Closes the connection without waiting for final transcript.

        Args:
            frame: The cancel frame.
        """
        await super().cancel(frame)
        await self._cleanup()

    async def run_stt(self, audio: bytes) -> AsyncGenerator[Frame, None]:
        """Send audio data to Ringg Parrot STT Service.

        Audio is sent as raw binary bytes to the WebSocket. The server
        will process and send back transcripts via the receive task.

        Args:
            audio: Raw audio bytes to transcribe (PCM 16-bit signed).

        Yields:
            Frame: None (transcription results come via WebSocket callbacks).
        """
        await self.start_processing_metrics()
        if self._websocket and self._websocket.state is State.OPEN:
            try:
                await self._websocket.send(audio)
            except websockets.exceptions.ConnectionClosed:
                logger.warning("Ringg Parrot STT: WebSocket closed while sending audio")
            except Exception as e:
                logger.error(f"Ringg Parrot STT: Error sending audio: {e}")
        await self.stop_processing_metrics()

        yield None

    @traced_stt
    async def _handle_transcription(
        self, transcript: str, is_final: bool, language: Optional[str] = None
    ):
        """Handle a transcription result with tracing.

        This method is decorated with @traced_stt to capture STT metrics
        including transcription latency and text content for observability.

        Args:
            transcript: The transcribed text.
            is_final: Whether this is a final (not interim) transcription.
            language: The detected/configured language.
        """
        # Tracing is handled by the @traced_stt decorator
        # This method exists to provide the hook for metrics collection
        pass

    async def _send_end_message(self):
        """Send end message to finalize current transcription."""
        if self._websocket and self._websocket.state is State.OPEN:
            try:
                await self._websocket.send(STOP_MESSAGE)
                logger.debug("Ringg Parrot STT: Sent end message to finalize transcription")
            except Exception as e:
                logger.error(f"Ringg Parrot STT: Error sending end message: {e}")

    async def process_frame(self, frame: Frame, direction: FrameDirection):
        """Process frames, handling interruption and finalization.

        Handles:
        - UserStoppedSpeakingFrame: Sends end message to get final transcript
        - StartInterruptionFrame: Sends end message to finalize any pending
          transcript before the interruption context switch

        Args:
            frame: The frame to process.
            direction: The direction of frame processing.
        """
        await super().process_frame(frame, direction)

        if isinstance(frame, UserStoppedSpeakingFrame):
            # User stopped speaking - finalize to get the transcript
            await self._send_end_message()

    async def _receive_task_handler(self):
        """Handle incoming messages from the WebSocket server.

        Processes transcript messages (partial and final), chunk acknowledgments,
        and error responses. Emits appropriate frames for each message type.
        """
        if not self._websocket:
            return

        try:
            async for message in self._websocket:
                try:
                    content = json.loads(message)
                except json.JSONDecodeError:
                    logger.warning(f"Ringg Parrot STT: Invalid JSON: {message}")
                    continue

                msg_type = content.get("type")

                if msg_type == "transcript":
                    transcription = content.get("transcription", "")
                    is_partial = content.get("is_partial", True)

                    if is_partial:
                        # Partial/interim transcription
                        if transcription.strip():
                            await self.push_frame(
                                InterimTranscriptionFrame(
                                    text=transcription,
                                    user_id=self._user_id,
                                    timestamp=time_now_iso8601(),
                                )
                            )
                    else:
                        # Final transcription
                        if transcription.strip():
                            await self.push_frame(
                                TranscriptionFrame(
                                    text=transcription,
                                    user_id=self._user_id,
                                    timestamp=time_now_iso8601(),
                                )
                            )
                            await self._handle_transcription(
                                transcription, is_final=True, language=self._params.language
                            )
                            await self.stop_processing_metrics()

                elif msg_type == "chunk":
                    # Acknowledgment of received audio chunk - no action needed
                    pass

                elif msg_type == "keepalive":
                    # Keepalive response - no action needed
                    pass

                elif msg_type == "error":
                    error_detail = content.get("detail", "Unknown error")
                    logger.error(f"Ringg Parrot STT error: {error_detail}")
                    await self.push_error(ErrorFrame(f"Ringg Parrot STT: {error_detail}"))

                elif msg_type == "ready":
                    # Already handled in start(), but log if received again
                    logger.debug("Ringg Parrot STT: Received additional ready message")

                else:
                    logger.warning(f"Ringg Parrot STT: Unknown message type: {msg_type}")

        except websockets.exceptions.ConnectionClosed:
            # Expected when closing the connection
            logger.debug("Ringg Parrot STT: WebSocket connection closed")
        except Exception as e:
            logger.error(f"Ringg Parrot STT receive error: {e}")
            await self.push_error(ErrorFrame(f"Ringg Parrot STT receive error: {e}"))


class RinggParrotV2InputParams(BaseModel):
    """Configuration parameters for Ringg Parrot STT V2 service.

    Parameters:
        api_key: API key for authentication.
        encoding: Audio encoding format (default: int16 for PCM signed 16-bit).
        language: Language for transcription (default: Hindi).
        mode: Processing mode (default: on_final).
        drop_last_chunk: Whether to drop the last audio chunk (default: False).
    """

    api_key: str = ""
    encoding: str = "int16"
    language: str = "hi"
    mode: str = "on_final"
    drop_last_chunk: bool = False


class RinggParrotSTTV2Service(STTService):
    """Speech-to-Text service using Ringg Parrot V2 WebSocket API.

    This service connects to a Ringg Parrot STT V2 WebSocket server for real-time
    transcription. It operates in on_final mode and forwards VAD events to the
    server. Finalization is triggered when Pipecat's VAD detects the user stopped
    speaking, at which point a vad_event and end message are sent.

    Protocol differences from V1:
    - Uses ``mode`` instead of ``prediction_method``
    - Uses ``is_final`` (bool) instead of ``is_partial``
    - Sends VAD events (user_start_speaking / user_stop_speaking) to the server
    - Supports ``drop_last_chunk`` parameter
    - Server responds with ``ack`` (not ``chunk``) and ``pong`` (not ``keepalive``)
    """

    def __init__(
        self,
        *,
        url: str = "ws://43.204.215.249:7860/v1/audio/stream",
        sample_rate: Optional[int] = 16000,
        params: Optional[RinggParrotV2InputParams] = None,
        **kwargs,
    ):
        super().__init__(sample_rate=sample_rate, **kwargs)
        params = params or RinggParrotV2InputParams()

        self._url = url
        self._params = params
        self._websocket = None
        self._receive_task = None
        self._client_id: Optional[str] = None

        self.set_model_name("ringg-parrot-v2")

    async def start(self, frame: StartFrame):
        """Start the Ringg Parrot STT V2 websocket connection."""
        await super().start(frame)
        if self._websocket:
            return

        try:
            self._websocket = await websocket_connect(self._url)

            if not self._websocket:
                logger.error(f"Unable to connect to Ringg Parrot STT V2 at {self._url}")
                await self._call_event_handler("on_connection_error", "Connection failed")
                return

            start_config = {
                "type": "start",
                "sample_rate": self.sample_rate,
                "encoding": self._params.encoding,
                "language": self._params.language,
                "api_key": self._params.api_key,
                "mode": self._params.mode,
                "drop_last_chunk": self._params.drop_last_chunk,
            }
            await self._websocket.send(json.dumps(start_config))

            ready_response = await asyncio.wait_for(self._websocket.recv(), timeout=10.0)
            ready_data = json.loads(ready_response)

            if ready_data.get("type") != "ready":
                error_msg = f"Unexpected response: {ready_data}"
                logger.error(f"Ringg Parrot STT V2: {error_msg}")
                await self._call_event_handler("on_connection_error", error_msg)
                await self._cleanup()
                return

            self._client_id = ready_data.get("client_id")

            if self._websocket and not self._receive_task:
                self._receive_task = self.create_task(self._receive_task_handler())

            await self._call_event_handler("on_connected")

        except asyncio.TimeoutError:
            logger.error("Ringg Parrot STT V2: Connection timeout waiting for ready response")
            await self._call_event_handler("on_connection_error", "Connection timeout")
            await self._cleanup()
        except Exception as e:
            logger.error(f"Ringg Parrot STT V2: Connection error: {e}")
            await self._call_event_handler("on_connection_error", str(e))
            await self._cleanup()

    async def _cleanup(self):
        """Clean up WebSocket connection and tasks."""
        if self._websocket:
            try:
                await self._send_end_message()
                await self._websocket.close()
            except Exception:
                pass
            self._websocket = None

        if self._receive_task:
            if self._receive_task != asyncio.current_task():
                await self._receive_task
            self._receive_task = None

        self._client_id = None
        logger.debug("Disconnected from Ringg Parrot STT V2")
        await self._call_event_handler("on_disconnected")

    async def stop(self, frame: EndFrame):
        """Stop the Ringg Parrot STT V2 websocket connection."""
        await super().stop(frame)
        await self._cleanup()

    async def cancel(self, frame: CancelFrame):
        """Cancel the Ringg Parrot STT V2 websocket connection immediately."""
        await super().cancel(frame)
        await self._cleanup()

    async def run_stt(self, audio: bytes) -> AsyncGenerator[Frame, None]:
        """Send audio data to Ringg Parrot STT V2 Service."""
        await self.start_processing_metrics()
        if self._websocket and self._websocket.state is State.OPEN:
            try:
                await self._websocket.send(audio)
            except websockets.exceptions.ConnectionClosed:
                logger.warning("Ringg Parrot STT V2: WebSocket closed while sending audio")
            except Exception as e:
                logger.error(f"Ringg Parrot STT V2: Error sending audio: {e}")
        await self.stop_processing_metrics()

        yield None

    async def _send_vad_event(self, state: str):
        """Send a VAD event to the server."""
        if self._websocket and self._websocket.state is State.OPEN:
            try:
                msg = json.dumps({"type": "vad_event", "state": state})
                await self._websocket.send(msg)
            except Exception as e:
                logger.error(f"Ringg Parrot STT V2: Error sending vad_event: {e}")

    async def _send_reset_message(self):
        """Send reset message to finalize current transcription."""
        if self._websocket and self._websocket.state is State.OPEN:
            try:
                await self._websocket.send(RESET_MESSAGE_V2)
            except Exception as e:
                logger.error(f"Ringg Parrot STT V2: Error sending reset message: {e}")

    async def _send_end_message(self):
        """Send end message to finalize current transcription."""
        if self._websocket and self._websocket.state is State.OPEN:
            try:
                await self._websocket.send(END_MESSAGE_V2)
            except Exception as e:
                logger.error(f"Ringg Parrot STT V2: Error sending end message: {e}")

    @traced_stt
    async def _handle_transcription(
        self, transcript: str, is_final: bool, language: Optional[str] = None
    ):
        """Handle a transcription result with tracing."""
        pass

    async def process_frame(self, frame: Frame, direction: FrameDirection):
        """Process frames, handling VAD events and finalization.

        Handles:
        - UserStartedSpeakingFrame: Sends vad_event(user_start_speaking)
        - UserStoppedSpeakingFrame: Sends vad_event(user_stop_speaking) + end message
        """
        await super().process_frame(frame, direction)

        if isinstance(frame, UserStartedSpeakingFrame):
            await self._send_vad_event("user_start_speaking")
        elif isinstance(frame, UserStoppedSpeakingFrame):
            await self._send_reset_message()

    async def _receive_task_handler(self):
        """Handle incoming messages from the V2 WebSocket server."""
        if not self._websocket:
            return

        try:
            async for message in self._websocket:
                try:
                    content = json.loads(message)
                except json.JSONDecodeError:
                    logger.warning(f"Ringg Parrot STT V2: Invalid JSON: {message}")
                    continue

                msg_type = content.get("type")

                if msg_type == "transcript":
                    transcription = content.get("transcription", "")
                    logger.debug(f"Ringg Parrot STT V2: Received transcript: {transcription}")
                    is_final = content.get("is_final", False)

                    if is_final:
                        if transcription.strip():
                            await self.push_frame(
                                TranscriptionFrame(
                                    text=transcription,
                                    user_id=self._user_id,
                                    timestamp=time_now_iso8601(),
                                )
                            )
                            await self._handle_transcription(
                                transcription, is_final=True, language=self._params.language
                            )
                            await self.stop_processing_metrics()
                    else:
                        if transcription.strip():
                            await self.push_frame(
                                InterimTranscriptionFrame(
                                    text=transcription,
                                    user_id=self._user_id,
                                    timestamp=time_now_iso8601(),
                                )
                            )

                elif msg_type == "ack":
                    pass

                elif msg_type == "pong":
                    pass

                elif msg_type == "error":
                    error_detail = content.get("detail", "Unknown error")
                    error_code = content.get("code", "UNKNOWN")
                    logger.error(f"Ringg Parrot STT V2 error [{error_code}]: {error_detail}")
                    await self.push_error(ErrorFrame(f"Ringg Parrot STT V2: {error_detail}"))

                elif msg_type == "ready":
                    logger.debug(f"Ringg Parrot STT V2: Received additional ready message {content}")

                else:
                    logger.warning(f"Ringg Parrot STT V2: Unknown message type: {msg_type}")

        except websockets.exceptions.ConnectionClosed:
            logger.debug("Ringg Parrot STT V2: WebSocket connection closed")
        except Exception as e:
            logger.error(f"Ringg Parrot STT V2 receive error: {e}")
            await self.push_error(ErrorFrame(f"Ringg Parrot STT V2 receive error: {e}"))
