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

"""AWS Transcribe Speech-to-Text service implementation.

This module provides a WebSocket-based connection to AWS Transcribe for real-time
speech-to-text transcription with support for multiple languages and audio formats.
"""

import asyncio
import json
import os
import random
import string
from typing import AsyncGenerator, Optional

from loguru import logger

from pipecat.frames.frames import (
    CancelFrame,
    EndFrame,
    ErrorFrame,
    Frame,
    InterimTranscriptionFrame,
    StartFrame,
    TranscriptionFrame,
)
from pipecat.services.aws.utils import build_event_message, decode_event, get_presigned_url
from pipecat.services.stt_service import STTService
from pipecat.transcriptions.language import Language
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 AWS services, you need to `pip install pipecat-ai[aws]`.")
    raise Exception(f"Missing module: {e}")


class AWSTranscribeSTTService(STTService):
    """AWS Transcribe Speech-to-Text service using WebSocket streaming.

    Provides real-time speech transcription using AWS Transcribe's streaming API.
    Supports multiple languages, configurable sample rates, and both interim and
    final transcription results.
    """

    def __init__(
        self,
        *,
        api_key: Optional[str] = None,
        aws_access_key_id: Optional[str] = None,
        aws_session_token: Optional[str] = None,
        region: Optional[str] = "us-east-1",
        sample_rate: int = 16000,
        language: Language = Language.EN,
        **kwargs,
    ):
        """Initialize the AWS Transcribe STT service.

        Args:
            api_key: AWS secret access key. If None, uses AWS_SECRET_ACCESS_KEY environment variable.
            aws_access_key_id: AWS access key ID. If None, uses AWS_ACCESS_KEY_ID environment variable.
            aws_session_token: AWS session token for temporary credentials. If None, uses AWS_SESSION_TOKEN environment variable.
            region: AWS region for the service. Defaults to "us-east-1".
            sample_rate: Audio sample rate in Hz. Must be 8000 or 16000. Defaults to 16000.
            language: Language for transcription. Defaults to English.
            **kwargs: Additional arguments passed to parent STTService class.
        """
        super().__init__(**kwargs)

        self._settings = {
            "sample_rate": sample_rate,
            "language": language,
            "media_encoding": "linear16",  # AWS expects raw PCM
            "number_of_channels": 1,
            "show_speaker_label": False,
            "enable_channel_identification": False,
        }

        # Validate sample rate - AWS Transcribe only supports 8000 Hz or 16000 Hz
        if sample_rate not in [8000, 16000]:
            logger.warning(
                f"AWS Transcribe only supports 8000 Hz or 16000 Hz sample rates. Converting from {sample_rate} Hz to 16000 Hz."
            )
            self._settings["sample_rate"] = 16000

        self._credentials = {
            "aws_access_key_id": aws_access_key_id or os.getenv("AWS_ACCESS_KEY_ID"),
            "aws_secret_access_key": api_key or os.getenv("AWS_SECRET_ACCESS_KEY"),
            "aws_session_token": aws_session_token or os.getenv("AWS_SESSION_TOKEN"),
            "region": region or os.getenv("AWS_REGION", "us-east-1"),
        }

        self._ws_client = None
        self._connection_lock = asyncio.Lock()
        self._connecting = False
        self._receive_task = None

    def get_service_encoding(self, encoding: str) -> str:
        """Convert internal encoding format to AWS Transcribe format.

        Args:
            encoding: Internal encoding format string.

        Returns:
            AWS Transcribe compatible encoding format.
        """
        encoding_map = {
            "linear16": "pcm",  # AWS expects "pcm" for 16-bit linear PCM
        }
        return encoding_map.get(encoding, encoding)

    async def start(self, frame: StartFrame):
        """Initialize the connection when the service starts.

        Args:
            frame: Start frame signaling service initialization.

        Raises:
            RuntimeError: If WebSocket connection cannot be established after retries.
        """
        await super().start(frame)
        logger.info("Starting AWS Transcribe service...")
        retry_count = 0
        max_retries = 3

        while retry_count < max_retries:
            try:
                await self._connect()
                if self._ws_client and self._ws_client.state is State.OPEN:
                    logger.info("Successfully established WebSocket connection")
                    return
                logger.warning("WebSocket connection not established after connect")
            except Exception as e:
                logger.error(f"Failed to connect (attempt {retry_count + 1}/{max_retries}): {e}")
                retry_count += 1
                if retry_count < max_retries:
                    await asyncio.sleep(1)  # Wait before retrying

        raise RuntimeError("Failed to establish WebSocket connection after multiple attempts")

    async def stop(self, frame: EndFrame):
        """Stop the service and disconnect from AWS Transcribe.

        Args:
            frame: End frame signaling service shutdown.
        """
        await super().stop(frame)
        await self._disconnect()

    async def cancel(self, frame: CancelFrame):
        """Cancel the service and disconnect from AWS Transcribe.

        Args:
            frame: Cancel frame signaling service cancellation.
        """
        await super().cancel(frame)
        await self._disconnect()

    async def run_stt(self, audio: bytes) -> AsyncGenerator[Frame, None]:
        """Process audio data and send to AWS Transcribe.

        Args:
            audio: Raw audio bytes to transcribe.

        Yields:
            ErrorFrame: If processing fails or connection issues occur.
        """
        try:
            # Ensure WebSocket is connected
            if not self._ws_client or self._ws_client.state is State.CLOSED:
                logger.debug("WebSocket not connected, attempting to reconnect...")
                try:
                    await self._connect()
                except Exception as e:
                    logger.error(f"Failed to reconnect: {e}")
                    yield ErrorFrame("Failed to reconnect to AWS Transcribe", fatal=False)
                    return

            # Format the audio data according to AWS event stream format
            event_message = build_event_message(audio)

            # Send the formatted event message
            try:
                await self._ws_client.send(event_message)
                # Start metrics after first chunk sent
                await self.start_processing_metrics()
                await self.start_ttfb_metrics()
            except websockets.exceptions.ConnectionClosed as e:
                logger.warning(f"Connection closed while sending: {e}")
                await self._disconnect()
                # Don't yield error here - we'll retry on next frame
            except Exception as e:
                logger.error(f"Error sending audio: {e}")
                yield ErrorFrame(f"AWS Transcribe error: {str(e)}", fatal=False)
                await self._disconnect()

        except Exception as e:
            logger.error(f"Error in run_stt: {e}")
            yield ErrorFrame(f"AWS Transcribe error: {str(e)}", fatal=False)
            await self._disconnect()

    async def _connect(self):
        """Connect to AWS Transcribe with connection state management."""
        if self._ws_client and self._ws_client.state is State.OPEN and self._receive_task:
            logger.debug(f"{self} Already connected")
            return

        async with self._connection_lock:
            if self._connecting:
                logger.debug(f"{self} Connection already in progress")
                return

            try:
                self._connecting = True
                logger.debug(f"{self} Starting connection process...")

                if self._ws_client:
                    await self._disconnect()

                language_code = self.language_to_service_language(
                    Language(self._settings["language"])
                )
                if not language_code:
                    raise ValueError(f"Unsupported language: {self._settings['language']}")

                # Generate random websocket key
                websocket_key = "".join(
                    random.choices(
                        string.ascii_uppercase + string.ascii_lowercase + string.digits, k=20
                    )
                )

                # Add required headers
                additional_headers = {
                    "Origin": "https://localhost",
                    "Sec-WebSocket-Key": websocket_key,
                    "Sec-WebSocket-Version": "13",
                    "Connection": "keep-alive",
                }

                # Get presigned URL
                presigned_url = get_presigned_url(
                    region=self._credentials["region"],
                    credentials={
                        "access_key": self._credentials["aws_access_key_id"],
                        "secret_key": self._credentials["aws_secret_access_key"],
                        "session_token": self._credentials["aws_session_token"],
                    },
                    language_code=language_code,
                    media_encoding=self.get_service_encoding(
                        self._settings["media_encoding"]
                    ),  # Convert to AWS format
                    sample_rate=self._settings["sample_rate"],
                    number_of_channels=self._settings["number_of_channels"],
                    enable_partial_results_stabilization=True,
                    partial_results_stability="high",
                    show_speaker_label=self._settings["show_speaker_label"],
                    enable_channel_identification=self._settings["enable_channel_identification"],
                )

                logger.debug(f"{self} Connecting to WebSocket with URL: {presigned_url[:100]}...")

                # Connect with the required headers and settings
                self._ws_client = await websocket_connect(
                    presigned_url,
                    additional_headers=additional_headers,
                    subprotocols=["mqtt"],
                    ping_interval=None,
                    ping_timeout=None,
                    compression=None,
                )

                logger.debug(f"{self} WebSocket connected, starting receive task...")

                # Start receive task
                self._receive_task = self.create_task(self._receive_loop())

                logger.info(f"{self} Successfully connected to AWS Transcribe")

            except Exception as e:
                logger.error(f"{self} Failed to connect to AWS Transcribe: {e}")
                await self._disconnect()
                raise

            finally:
                self._connecting = False

    async def _disconnect(self):
        """Disconnect from AWS Transcribe."""
        if self._receive_task:
            await self.cancel_task(self._receive_task)
            self._receive_task = None

        try:
            if self._ws_client and self._ws_client.state is State.OPEN:
                # Send end-stream message
                end_stream = {"message-type": "event", "event": "end"}
                await self._ws_client.send(json.dumps(end_stream))
            await self._ws_client.close()
        except Exception as e:
            logger.warning(f"{self} Error closing WebSocket connection: {e}")
        finally:
            self._ws_client = None

    def language_to_service_language(self, language: Language) -> str | None:
        """Convert internal language enum to AWS Transcribe language code.

        Source:
        https://docs.aws.amazon.com/transcribe/latest/dg/supported-languages.html
        All language codes that support streaming are included.

        Args:
            language: Internal language enumeration value.

        Returns:
            AWS Transcribe compatible language code, or None if unsupported.
        """
        language_map = {
            # Afrikaans
            Language.AF: "af-ZA",
            Language.AF_ZA: "af-ZA",
            # Arabic
            Language.AR: "ar-SA",  # Default to Modern Standard Arabic
            Language.AR_AE: "ar-AE",  # Gulf Arabic
            Language.AR_SA: "ar-SA",  # Modern Standard Arabic
            # Basque
            Language.EU: "eu-ES",
            Language.EU_ES: "eu-ES",
            # Catalan
            Language.CA: "ca-ES",
            Language.CA_ES: "ca-ES",
            # Chinese
            Language.ZH: "zh-CN",  # Default to Simplified
            Language.ZH_CN: "zh-CN",  # Simplified
            Language.ZH_TW: "zh-TW",  # Traditional
            Language.ZH_HK: "zh-HK",  # Cantonese (also yue-HK)
            Language.YUE: "zh-HK",  # Cantonese fallback
            # Croatian
            Language.HR: "hr-HR",
            Language.HR_HR: "hr-HR",
            # Czech
            Language.CS: "cs-CZ",
            Language.CS_CZ: "cs-CZ",
            # Danish
            Language.DA: "da-DK",
            Language.DA_DK: "da-DK",
            # Dutch
            Language.NL: "nl-NL",
            Language.NL_NL: "nl-NL",
            # English
            Language.EN: "en-US",  # Default to US
            Language.EN_AU: "en-AU",  # Australian
            Language.EN_GB: "en-GB",  # British
            Language.EN_IN: "en-IN",  # Indian
            Language.EN_IE: "en-IE",  # Irish
            Language.EN_NZ: "en-NZ",  # New Zealand
            # Note: Scottish (en-AB) and Welsh (en-WL) don't have direct Language enum matches
            Language.EN_ZA: "en-ZA",  # South African
            Language.EN_US: "en-US",  # US
            # Persian/Farsi
            Language.FA: "fa-IR",
            Language.FA_IR: "fa-IR",
            # Finnish
            Language.FI: "fi-FI",
            Language.FI_FI: "fi-FI",
            # French
            Language.FR: "fr-FR",  # Default to France
            Language.FR_FR: "fr-FR",
            Language.FR_CA: "fr-CA",  # Canadian
            # Galician
            Language.GL: "gl-ES",
            Language.GL_ES: "gl-ES",
            # Georgian
            Language.KA: "ka-GE",
            Language.KA_GE: "ka-GE",
            # German
            Language.DE: "de-DE",  # Default to Germany
            Language.DE_DE: "de-DE",
            Language.DE_CH: "de-CH",  # Swiss
            # Greek
            Language.EL: "el-GR",
            Language.EL_GR: "el-GR",
            # Hebrew
            Language.HE: "he-IL",
            Language.HE_IL: "he-IL",
            # Hindi
            Language.HI: "hi-IN",
            Language.HI_IN: "hi-IN",
            # Indonesian
            Language.ID: "id-ID",
            Language.ID_ID: "id-ID",
            # Italian
            Language.IT: "it-IT",
            Language.IT_IT: "it-IT",
            # Japanese
            Language.JA: "ja-JP",
            Language.JA_JP: "ja-JP",
            # Korean
            Language.KO: "ko-KR",
            Language.KO_KR: "ko-KR",
            # Latvian
            Language.LV: "lv-LV",
            Language.LV_LV: "lv-LV",
            # Malay
            Language.MS: "ms-MY",
            Language.MS_MY: "ms-MY",
            # Norwegian
            Language.NB: "no-NO",  # Norwegian Bokmål
            Language.NB_NO: "no-NO",
            Language.NO: "no-NO",
            # Polish
            Language.PL: "pl-PL",
            Language.PL_PL: "pl-PL",
            # Portuguese
            Language.PT: "pt-PT",  # Default to Portugal
            Language.PT_PT: "pt-PT",
            Language.PT_BR: "pt-BR",  # Brazilian
            # Romanian
            Language.RO: "ro-RO",
            Language.RO_RO: "ro-RO",
            # Russian
            Language.RU: "ru-RU",
            Language.RU_RU: "ru-RU",
            # Serbian
            Language.SR: "sr-RS",
            Language.SR_RS: "sr-RS",
            # Slovak
            Language.SK: "sk-SK",
            Language.SK_SK: "sk-SK",
            # Somali
            Language.SO: "so-SO",
            Language.SO_SO: "so-SO",
            # Spanish
            Language.ES: "es-ES",  # Default to Spain
            Language.ES_ES: "es-ES",
            Language.ES_US: "es-US",  # US Spanish
            # Swedish
            Language.SV: "sv-SE",
            Language.SV_SE: "sv-SE",
            # Tagalog/Filipino
            Language.TL: "tl-PH",
            Language.FIL: "tl-PH",  # Filipino maps to Tagalog
            Language.FIL_PH: "tl-PH",
            # Thai
            Language.TH: "th-TH",
            Language.TH_TH: "th-TH",
            # Ukrainian
            Language.UK: "uk-UA",
            Language.UK_UA: "uk-UA",
            # Vietnamese
            Language.VI: "vi-VN",
            Language.VI_VN: "vi-VN",
            # Zulu
            Language.ZU: "zu-ZA",
            Language.ZU_ZA: "zu-ZA",
        }

        return language_map.get(language)

    @traced_stt
    async def _handle_transcription(
        self, transcript: str, is_final: bool, language: Optional[str] = None
    ):
        pass

    async def _receive_loop(self):
        """Background task to receive and process messages from AWS Transcribe."""
        while True:
            if not self._ws_client or self._ws_client.state is State.CLOSED:
                logger.warning(f"{self} WebSocket closed in receive loop")
                break

            try:
                response = await asyncio.wait_for(self._ws_client.recv(), timeout=1.0)

                headers, payload = decode_event(response)

                if headers.get(":message-type") == "event":
                    # Process transcription results
                    results = payload.get("Transcript", {}).get("Results", [])
                    if results:
                        result = results[0]
                        alternatives = result.get("Alternatives", [])
                        if alternatives:
                            transcript = alternatives[0].get("Transcript", "")
                            is_final = not result.get("IsPartial", True)

                            if transcript:
                                await self.stop_ttfb_metrics()
                                if is_final:
                                    await self.push_frame(
                                        TranscriptionFrame(
                                            transcript,
                                            self._user_id,
                                            time_now_iso8601(),
                                            self._settings["language"],
                                            result=result,
                                        )
                                    )
                                    await self._handle_transcription(
                                        transcript,
                                        is_final,
                                        self._settings["language"],
                                    )
                                    await self.stop_processing_metrics()
                                else:
                                    await self.push_frame(
                                        InterimTranscriptionFrame(
                                            transcript,
                                            self._user_id,
                                            time_now_iso8601(),
                                            self._settings["language"],
                                            result=result,
                                        )
                                    )
                elif headers.get(":message-type") == "exception":
                    error_msg = payload.get("Message", "Unknown error")
                    logger.error(f"{self} Exception from AWS: {error_msg}")
                    await self.push_frame(
                        ErrorFrame(f"AWS Transcribe error: {error_msg}", fatal=False)
                    )
                else:
                    logger.debug(f"{self} Other message type received: {headers}")
                    logger.debug(f"{self} Payload: {payload}")
            except asyncio.TimeoutError:
                self.reset_watchdog()
            except websockets.exceptions.ConnectionClosed as e:
                logger.error(
                    f"{self} WebSocket connection closed in receive loop with code {e.code}: {e.reason}"
                )
                break
            except Exception as e:
                logger.error(f"{self} Unexpected error in receive loop: {e}")
                break
