import asyncio
from typing import AsyncGenerator

from loguru import logger

from pipecat.frames.frames import (
    CancelFrame,
    EndFrame,
    ErrorFrame,
    Frame,
    InterimTranscriptionFrame,
    StartFrame,
    TranscriptionFrame,
)
from pipecat.services.ai_services import STTService
from pipecat.transcriptions.language import Language
from pipecat.utils.time import time_now_iso8601

try:
    import assemblyai as aai
    from assemblyai import AudioEncoding
except ModuleNotFoundError as e:
    logger.error(f"Exception: {e}")
    logger.error(
        "In order to use AssemblyAI, you need to `pip install pipecat-ai[assemblyai]`. Also, set `ASSEMBLYAI_API_KEY` environment variable."
    )
    raise Exception(f"Missing module: {e}")


class AssemblyAISTTService(STTService):
    def __init__(
        self,
        *,
        api_key: str,
        sample_rate: int = 16000,
        encoding: AudioEncoding = AudioEncoding("pcm_s16le"),
        language=Language.EN,  # Only English is supported for Realtime
        **kwargs,
    ):
        super().__init__(**kwargs)

        aai.settings.api_key = api_key
        self._transcriber: aai.RealtimeTranscriber | None = None
        # Store reference to the main event loop for use in callback functions
        self._loop = asyncio.get_event_loop()

        self._settings = {
            "sample_rate": sample_rate,
            "encoding": encoding,
            "language": language,
        }

    async def set_language(self, language: Language):
        logger.info(f"Switching STT language to: [{language}]")
        self._settings["language"] = language

    async def start(self, frame: StartFrame):
        await super().start(frame)
        await self._connect()

    async def stop(self, frame: EndFrame):
        await super().stop(frame)
        await self._disconnect()

    async def cancel(self, frame: CancelFrame):
        await super().cancel(frame)
        await self._disconnect()

    async def run_stt(self, audio: bytes) -> AsyncGenerator[Frame, None]:
        """
        Process an audio chunk for STT transcription.

        This method streams the audio data to AssemblyAI for real-time transcription.
        Transcription results are handled asynchronously via callback functions.

        :param audio: Audio data as bytes
        :yield: None (transcription frames are pushed via self.push_frame in callbacks)
        """
        if self._transcriber:
            await self.start_processing_metrics()
            self._transcriber.stream(audio)
            await self.stop_processing_metrics()
        yield None

    async def _connect(self):
        """
        Establish a connection to the AssemblyAI real-time transcription service.

        This method sets up the necessary callback functions and initializes the
        AssemblyAI transcriber.
        """

        def on_open(session_opened: aai.RealtimeSessionOpened):
            """Callback for when the connection to AssemblyAI is opened."""
            logger.info(f"{self}: Connected to AssemblyAI")

        def on_data(transcript: aai.RealtimeTranscript):
            """
            Callback for handling incoming transcription data.

            This function runs in a separate thread from the main asyncio event loop.
            It creates appropriate transcription frames and schedules them to be
            pushed to the next stage of the pipeline in the main event loop.
            """
            if not transcript.text:
                return

            timestamp = time_now_iso8601()

            if isinstance(transcript, aai.RealtimeFinalTranscript):
                frame = TranscriptionFrame(
                    transcript.text, "", timestamp, self._settings["language"]
                )
            else:
                frame = InterimTranscriptionFrame(
                    transcript.text, "", timestamp, self._settings["language"]
                )

            # Schedule the coroutine to run in the main event loop
            # This is necessary because this callback runs in a different thread
            asyncio.run_coroutine_threadsafe(self.push_frame(frame), self._loop)

        def on_error(error: aai.RealtimeError):
            """
            Callback for handling errors from AssemblyAI.

            Like on_data, this runs in a separate thread and schedules error
            handling in the main event loop.
            """
            logger.error(f"{self}: An error occurred: {error}")
            # Schedule the coroutine to run in the main event loop
            asyncio.run_coroutine_threadsafe(self.push_frame(ErrorFrame(str(error))), self._loop)

        def on_close():
            """Callback for when the connection to AssemblyAI is closed."""
            logger.info(f"{self}: Disconnected from AssemblyAI")

        self._transcriber = aai.RealtimeTranscriber(
            sample_rate=self._settings["sample_rate"],
            encoding=self._settings["encoding"],
            on_data=on_data,
            on_error=on_error,
            on_open=on_open,
            on_close=on_close,
        )
        self._transcriber.connect()

    async def _disconnect(self):
        """Disconnect from the AssemblyAI service and clean up resources."""
        if self._transcriber:
            self._transcriber.close()
            self._transcriber = None
