import asyncio
import json
import logging
import time
from typing import Dict, List, Optional
from uuid import uuid4

from pydantic import Field

from inputs.base import Message, SensorConfig
from inputs.base.loop import FuserInput
from providers.asr_rtsp_provider import ASRRTSPProvider
from providers.io_provider import IOProvider
from providers.sleep_ticker_provider import SleepTickerProvider
from providers.teleops_conversation_provider import TeleopsConversationProvider
from zenoh_msgs import ASRText, open_zenoh_session, prepare_header


class RivaASRRTSPSensorConfig(SensorConfig):
    """
    Configuration for Riva ASR Sensor.

    Parameters
    ----------
    api_key : Optional[str]
        API Key.
    rtsp_url : str
        RTSP URL for the audio stream. Default is "rtsp://localhost:8554/audio".
    rate : int
        Sampling rate. Default is 16000.
    base_url : str
        Base URL for the ASR service. Default is "wss://api-asr.openmind.com".
    enable_tts_interrupt : bool
        Enable TTS interrupt when ASR detects speech during playback.
    """

    api_key: Optional[str] = Field(default=None, description="API Key")
    rtsp_url: str = Field(
        default="rtsp://localhost:8554/audio",
        description="RTSP URL for the audio stream",
    )
    rate: int = Field(default=16000, description="Sampling rate")
    base_url: str = Field(default="wss://api-asr.openmind.com", description="Base URL for the ASR service")
    enable_tts_interrupt: bool = Field(
        default=False,
        description="Enable TTS interrupt (does not mute mic during TTS playback)",
    )


class RivaASRRTSPInput(FuserInput[RivaASRRTSPSensorConfig, Optional[str]]):
    """
    Automatic Speech Recognition (ASR) input handler.

    This class manages the input stream from an ASR service, buffering messages
    and providing text conversion capabilities.
    """

    def __init__(self, config: RivaASRRTSPSensorConfig):
        """
        Initialize RivaASRInput instance.

        Parameters
        ----------
        config : RivaASRRTSPSensorConfig
            Configuration for the ASR input handler.
        """
        super().__init__(config)

        # Buffer for storing the final output
        self.messages: List[str] = []

        # Set IO Provider
        self.descriptor_for_LLM = "Voice"
        self.io_provider = IOProvider()

        # Message buffer for incoming ASR messages
        self.message_buffer: asyncio.Queue[str] = asyncio.Queue()

        # Initialize ASR provider
        api_key = self.config.api_key
        rtsp_url = self.config.rtsp_url
        rate = self.config.rate
        base_url = self.config.base_url

        enable_tts_interrupt = self.config.enable_tts_interrupt

        self.asr: ASRRTSPProvider = ASRRTSPProvider(
            rtsp_url=rtsp_url,
            rate=rate,
            ws_url=base_url,
            language_code="en-US",
            enable_tts_interrupt=enable_tts_interrupt,
        )
        self.asr.start()
        self.asr.register_message_callback(self._handle_asr_message)

        # Initialize sleep ticker provider
        self.global_sleep_ticker_provider = SleepTickerProvider()

        # Initialize conversation provider
        self.conversation_provider = TeleopsConversationProvider(api_key=api_key)

        # Initialize Zenoh session
        self.asr_topic = "om/asr/text"
        self.session = None
        self.asr_publisher = None

        try:
            self.session = open_zenoh_session()
            self.asr_publisher = self.session.declare_publisher(self.asr_topic)
            logging.info("Zenoh ASR publisher initialized on topic 'om/asr/text'")
        except Exception as e:
            logging.warning(f"Could not initialize Zenoh for ASR broadcast: {e}")
            self.session = None
            self.asr_publisher = None

        # Guard flag: when True, this instance ignores incoming ASR messages.
        self._stopped = False

    def _handle_asr_message(self, raw_message: str):
        """
        Process incoming ASR messages.

        Parameters
        ----------
        raw_message : str
            Raw message received from ASR service
        """
        if self._stopped:
            return

        try:
            json_message: Dict = json.loads(raw_message)
            if "asr_reply" in json_message:
                asr_reply = json_message["asr_reply"]
                if len(asr_reply.split()) > 2:
                    self.message_buffer.put_nowait(asr_reply)
                    logging.info("Detected ASR message: %s", asr_reply)
        except json.JSONDecodeError:
            pass

    async def _poll(self) -> Optional[str]:
        """
        Poll for new messages in the buffer.

        Returns
        -------
        Optional[str]
            Message from the buffer if available, None otherwise
        """
        try:
            message = self.message_buffer.get_nowait()
            return message
        except asyncio.QueueEmpty:
            await asyncio.sleep(0.01)
            return None

    async def _raw_to_text(self, raw_input: Optional[str]) -> Optional[Message]:
        """
        Convert raw input to text format.

        Parameters
        ----------
        raw_input : Optional[str]
            Raw input to be processed

        Returns
        -------
        Optional[str]
            Processed text message or None if input is None
        """
        if raw_input is None:
            return None

        return Message(timestamp=time.time(), message=raw_input)

    async def raw_to_text(self, raw_input: Optional[str]):
        """
        Convert raw input to processed text and manage buffer.

        Parameters
        ----------
        raw_input : Optional[str]
            Raw input to be processed
        """
        pending_message = await self._raw_to_text(raw_input)
        if pending_message is None:
            if len(self.messages) != 0:
                self.global_sleep_ticker_provider.skip_sleep = True

        if pending_message is not None:
            if len(self.messages) == 0:
                self.messages.append(pending_message.message)
            else:
                self.messages[-1] = f"{self.messages[-1]} {pending_message.message}"

    def formatted_latest_buffer(self) -> Optional[str]:
        """
        Format and clear the latest buffer contents.

        Returns
        -------
        Optional[str]
            Formatted string of buffer contents or None if buffer is empty
        """
        if len(self.messages) == 0:
            return None

        result = f"""
{self.descriptor_for_LLM}: "{self.messages[-1]}"
"""
        # Add to IO provider and conversation provider
        self.io_provider.add_input(self.descriptor_for_LLM, self.messages[-1], time.time())
        self.io_provider.add_mode_transition_input(self.messages[-1])
        self.conversation_provider.store_user_message(self.messages[-1])

        # Publish to Zenoh
        if self.asr_publisher:
            try:
                asr_msg = ASRText(
                    header=prepare_header(str(uuid4())),
                    text=self.messages[-1],
                )
                self.asr_publisher.put(asr_msg.serialize())
                logging.info(f"Published ASR to Zenoh: {self.messages[-1]}")
            except Exception as e:
                logging.warning(f"Failed to publish ASR to Zenoh: {e}")

        # Reset messages buffer
        self.messages = []
        return result

    def stop(self):
        """
        Stop the ASR input handler and clean up resources.
        """
        logging.info("Stopping RivaASRRTSPInput, disabling callback")

        self._stopped = True

        while not self.message_buffer.empty():
            try:
                self.message_buffer.get_nowait()
            except asyncio.QueueEmpty:
                break

        self.messages = []

        if self.asr:
            try:
                self.asr.unregister_message_callback(self._handle_asr_message)
                logging.info("Unregistered ASR callback")
            except Exception as e:
                logging.warning(f"Failed to unregister ASR callback: {e}")

        if self.asr_publisher:
            try:
                self.asr_publisher.undeclare()
                logging.info("Zenoh ASR publisher undeclared")
            except Exception as e:
                logging.warning(f"Failed to undeclare Zenoh publisher: {e}")
            self.asr_publisher = None

        if self.session:
            try:
                self.session.close()
                logging.info("Zenoh ASR session closed")
            except Exception as e:
                logging.warning(f"Failed to close Zenoh session: {e}")
            self.session = None
