import ssl
import json
import asyncio
import websockets

from config.logger import setup_logging
from typing import Optional, Tuple, List
from core.providers.asr.base import ASRProviderBase
from core.providers.asr.utils import lang_tag_filter
from core.providers.asr.dto.dto import InterfaceType

TAG = __name__
logger = setup_logging()


class ASRProvider(ASRProviderBase):
    def __init__(self, config: dict, delete_audio_file: bool):
        """
        Initialize the ASRProvider with server configuration.
        :param config: Dictionary containing 'host', 'port', and 'is_ssl'.
        :param delete_audio_file: Boolean to indicate whether to delete audio files after processing.
        """
        super().__init__()
        self.interface_type = InterfaceType.NON_STREAM
        self.host = config.get("host", "localhost")
        self.port = config.get("port", 10095)
        self.api_key = config.get("api_key", "none")
        self.is_ssl = str(config.get("is_ssl", True)).lower() in (
            "true",
            "1",
            "yes",
        )
        self.output_dir = config.get("output_dir")
        self.delete_audio_file = delete_audio_file
        self.uri = (
            f"wss://{self.host}:{self.port}"
            if self.is_ssl
            else f"ws://{self.host}:{self.port}"
        )
        self.ssl_context = ssl.SSLContext() if self.is_ssl else None
        if self.ssl_context:
            self.ssl_context.check_hostname = False
            self.ssl_context.verify_mode = ssl.CERT_NONE

    async def _receive_responses(self, ws) -> None:
        """
        Asynchronous generator to receive messages from the WebSocket.
        Yields each message as it is received.
        """
        text = ""
        while True:
            try:
                response = await asyncio.wait_for(ws.recv(), timeout=5)
                response_data = json.loads(response)
                logger.bind(tag=TAG).debug(f"Received response: {response_data}")
                if response_data.get("is_final", True):
                    text += response_data.get("text", "")
                    break
                else:
                    text += response_data.get("text", "")
            except asyncio.TimeoutError:
                logger.bind(tag=TAG).error(
                    "Timeout while waiting for response from WebSocket."
                )
                break
            except websockets.exceptions.ConnectionClosed as e:
                logger.bind(tag=TAG).error(f"WebSocket connection closed: {e}")
                break
        return text

    async def _send_data(self, ws, pcm_data: bytes, session_id: str) -> tuple:
        """
        Internal method to handle WebSocket communication.
        Reuses the persistent WebSocket connection if available.
        :param pcm_data: PCM audio data to send.
        :param session_id: Unique session identifier.
        :return: Tuple containing recognized text and optional timestamp.
        """

        # Send initial configuration message
        config_message = json.dumps(
            {
                "mode": "offline",
                "chunk_size": [5, 10, 5],
                "chunk_interval": 10,
                "wav_name": session_id,
                "is_speaking": True,
                "itn": False,
            }
        )
        await ws.send(config_message)
        logger.bind(tag=TAG).debug(f"Sent configuration message: {config_message}")

        # Send PCM data
        await ws.send(pcm_data)
        logger.bind(tag=TAG).debug(f"Sent PCM data of length: {len(pcm_data)} bytes")

        # Indicate end of speech
        end_message = json.dumps({"is_speaking": False})
        await ws.send(end_message)
        logger.bind(tag=TAG).debug(f"Sent end message: {end_message}")

    async def speech_to_text(
        self, opus_data: List[bytes], session_id: str, artifacts=None
    ) -> Tuple[Optional[str], Optional[str]]:
        """
        Convert speech data to text using FunASR.
        :param opus_data: List of Opus-encoded audio data chunks.
        :param session_id: Unique session identifier.
        :return: Tuple containing recognized text and optional timestamp.
        """
        
        if artifacts is None:
            return "", None
        auth_header = {"Authorization": "Bearer; {}".format(self.api_key)}
        async with websockets.connect(
            self.uri,
            additional_headers=auth_header,
            subprotocols=["binary"],
            ping_interval=None,
            ssl=self.ssl_context,
        ) as ws:
            try:
                # Use asyncio to handle WebSocket communication
                send_task = asyncio.create_task(
                    self._send_data(ws, artifacts.pcm_bytes, session_id)
                )
                receive_task = asyncio.create_task(self._receive_responses(ws))

                # Gather tasks with error handling
                done, pending = await asyncio.wait(
                    [send_task, receive_task], return_when=asyncio.FIRST_EXCEPTION
                )

                # Cancel any pending tasks
                for task in pending:
                    task.cancel()

                # Check for exceptions in completed tasks
                for task in done:
                    if task.exception():
                        raise task.exception()

                # Get the result from the receive task
                result = receive_task.result()
                
                # match = re.match(r"<\|(.*?)\|><\|(.*?)\|><\|(.*?)\|>(.*)", result)
                # if match:
                #     result = match.group(4).strip()

                # Handle tags
                result = lang_tag_filter(result)
                return (
                    result,
                    artifacts.file_path,
                )  # Return the recognized text and timestamp (if any)

            except websockets.exceptions.ConnectionClosed as e:
                logger.bind(tag=TAG).error(f"WebSocket connection closed: {e}")
                return "", artifacts.file_path
            except Exception as e:
                logger.bind(tag=TAG).error(
                    f"Error during speech-to-text conversion: {e}", exc_info=True
                )
                return "", artifacts.file_path