"""VoicegroundBridgeTransport - Drop-in transport replacement for simulation.

This module provides a bridge transport that connects two Pipecat pipelines,
enabling closed-loop simulation. It extends Pipecat's base transport classes
to ensure full compatibility with VAD, turn analysis, and interruption handling.
"""

import asyncio
from typing import Optional

from loguru import logger
from pipecat.frames.frames import (
    CancelFrame,
    CancelTaskFrame,
    EndFrame,
    InputAudioRawFrame,
    OutputAudioRawFrame,
    StartFrame,
)
from pipecat.processors.frame_processor import FrameDirection, FrameProcessor
from pipecat.transports.base_input import BaseInputTransport
from pipecat.transports.base_output import BaseOutputTransport
from pipecat.transports.base_transport import BaseTransport, TransportParams
from pydantic import ConfigDict


class VoicegroundBridgeTransportParams(TransportParams):
    """Configuration parameters for bridge transport.

    Extends TransportParams with bridge-specific settings.
    """

    model_config = ConfigDict(arbitrary_types_allowed=True)


class VoicegroundBridgeInputTransport(BaseInputTransport):
    """Input transport that receives audio from a bridge queue.

    This extends BaseInputTransport to leverage Pipecat's built-in VAD
    processing, turn analysis, and interruption handling. Audio frames
    received from the paired output transport are pushed to the standard
    audio processing pipeline.
    """

    _params: VoicegroundBridgeTransportParams

    # Timeouts for graceful shutdown stages
    DISCONNECT_TIMEOUT_SECONDS = 5.0  # Wait for bot to handle on_client_disconnected
    CANCEL_TIMEOUT_SECONDS = 5.0  # Wait for CancelTaskFrame to take effect

    # Audio settings for silence generation
    SAMPLE_RATE = 16000
    NUM_CHANNELS = 1
    SILENCE_DURATION_MS = 20  # Duration of silence frames in milliseconds

    def __init__(
        self,
        transport: "VoicegroundBridgeTransport",
        params: VoicegroundBridgeTransportParams,
        audio_queue: "asyncio.Queue[InputAudioRawFrame | None]",
        **kwargs,
    ):
        """Initialize the bridge input transport.

        Args:
            transport: The parent bridge transport.
            params: Transport configuration parameters.
            audio_queue: Queue to receive audio frames from.
        """
        super().__init__(params, **kwargs)
        self._transport = transport
        self._bridge_audio_queue = audio_queue
        self._audio_read_task: asyncio.Task | None = None
        self._disconnected = False
        self._force_stop_task: asyncio.Task | None = None

    async def start(self, frame: StartFrame):
        """Start the input transport and begin reading from the bridge queue."""
        await super().start(frame)
        await self._transport._call_event_handler("on_client_connected")

        if not self._audio_read_task:
            self._audio_read_task = self.create_task(self._read_audio_from_bridge())

        await self.set_transport_ready(frame)

    async def stop(self, frame: EndFrame):
        """Stop the input transport."""
        self._disconnected = True
        await self._cleanup_tasks()
        await super().stop(frame)

    async def cancel(self, frame: CancelFrame):
        """Cancel the input transport."""
        self._disconnected = True
        await self._cleanup_tasks()
        await super().cancel(frame)

    async def _cleanup_tasks(self):
        if self._audio_read_task:
            await self.cancel_task(self._audio_read_task)
            self._audio_read_task = None
        if self._force_stop_task:
            self._force_stop_task.cancel()
            try:
                await self._force_stop_task
            except asyncio.CancelledError:
                pass
            self._force_stop_task = None

    def _create_silence_frame(self) -> InputAudioRawFrame:
        """Create a silence audio frame to keep VAD happy during idle periods."""
        num_samples = int(self.SAMPLE_RATE * self.SILENCE_DURATION_MS / 1000)
        silence = bytes(num_samples * self.NUM_CHANNELS * 2)  # 16-bit audio = 2 bytes per sample
        return InputAudioRawFrame(
            audio=silence,
            sample_rate=self.SAMPLE_RATE,
            num_channels=self.NUM_CHANNELS,
        )

    async def _read_audio_from_bridge(self):
        """Read audio from bridge queue and push to VAD processing."""
        while True:
            try:
                try:
                    audio_frame = await asyncio.wait_for(
                        self._bridge_audio_queue.get(),
                        timeout=0.02,  # 20ms for smooth audio
                    )
                except asyncio.TimeoutError:
                    # Send silence to keep VAD happy during idle periods
                    await self.push_audio_frame(self._create_silence_frame())
                    continue

                if audio_frame is None:
                    logger.debug(f"{self} received end of stream signal")
                    await self._transport._call_event_handler("on_client_disconnected")
                    self._force_stop_task = asyncio.create_task(self._force_stop_after_timeout())
                    break

                await self.push_audio_frame(audio_frame)

            except asyncio.CancelledError:
                break
            except Exception as e:
                logger.error(f"{self} error reading from bridge: {e}")

    async def _force_stop_after_timeout(self):
        """Force stop the pipeline if bot doesn't respond to disconnect.

        Uses a 3-stage graceful shutdown:
        1. Wait for bot to handle on_client_disconnected
        2. Send CancelTaskFrame upstream
        3. Push CancelFrame to force termination
        """
        try:
            await asyncio.sleep(self.DISCONNECT_TIMEOUT_SECONDS)
            if self._disconnected:
                return

            logger.warning(f"{self} bot didn't stop, sending CancelTaskFrame")
            await self.push_frame(CancelTaskFrame(), direction=FrameDirection.UPSTREAM)

            await asyncio.sleep(self.CANCEL_TIMEOUT_SECONDS)
            if self._disconnected:
                return

            logger.warning(f"{self} forcing immediate termination")
            await self.push_frame(CancelFrame(), direction=FrameDirection.DOWNSTREAM)

        except asyncio.CancelledError:
            pass


class VoicegroundBridgeOutputTransport(BaseOutputTransport):
    """Output transport that sends audio to a bridge queue."""

    _params: VoicegroundBridgeTransportParams

    def __init__(
        self,
        params: VoicegroundBridgeTransportParams,
        audio_queue: "asyncio.Queue[InputAudioRawFrame | None]",
        **kwargs,
    ):
        super().__init__(params, **kwargs)
        self._bridge_audio_queue = audio_queue

    async def start(self, frame: StartFrame):
        """Start the output transport."""
        await super().start(frame)
        await self.set_transport_ready(frame)

    async def stop(self, frame: EndFrame):
        """Stop the output transport and signal end of stream."""
        await self._bridge_audio_queue.put(None)
        await super().stop(frame)

    async def cancel(self, frame: CancelFrame):
        """Cancel the output transport."""
        await self._bridge_audio_queue.put(None)
        await super().cancel(frame)

    async def write_audio_frame(self, frame: OutputAudioRawFrame) -> bool:
        """Write an audio frame to the bridge queue."""
        input_frame = InputAudioRawFrame(
            audio=frame.audio,
            sample_rate=frame.sample_rate,
            num_channels=frame.num_channels,
        )
        await self._bridge_audio_queue.put(input_frame)
        return True


class VoicegroundBridgeTransport(BaseTransport):
    """Transport that connects two Pipecat pipelines for simulation.

    Drop-in replacement for LocalAudioTransport, WebsocketTransport, or SmallWebRTCTransport
    for audio-only bots. Supports on_client_connected/on_client_disconnected events.

    Note: This transport does not support video frames or app messages.
    """

    def __init__(
        self,
        params: VoicegroundBridgeTransportParams | None = None,
        *,
        paired_audio_queue_in: Optional["asyncio.Queue[InputAudioRawFrame | None]"] = None,
        paired_audio_queue_out: Optional["asyncio.Queue[InputAudioRawFrame | None]"] = None,
    ):
        super().__init__()

        self._params = params or VoicegroundBridgeTransportParams(
            audio_in_enabled=True,
            audio_out_enabled=True,
        )

        if paired_audio_queue_in is None and paired_audio_queue_out is None:
            self._audio_queue_in: asyncio.Queue[InputAudioRawFrame | None] = asyncio.Queue()
            self._audio_queue_out: asyncio.Queue[InputAudioRawFrame | None] = asyncio.Queue()
            self._is_primary = True
        else:
            assert paired_audio_queue_in is not None
            assert paired_audio_queue_out is not None
            self._audio_queue_in = paired_audio_queue_in
            self._audio_queue_out = paired_audio_queue_out
            self._is_primary = False

        self._input: VoicegroundBridgeInputTransport | None = None
        self._output: VoicegroundBridgeOutputTransport | None = None

        self._register_event_handler("on_client_connected")
        self._register_event_handler("on_client_disconnected")

    def input(self) -> FrameProcessor:
        """Get the input transport."""
        if not self._input:
            self._input = VoicegroundBridgeInputTransport(
                self,
                self._params,
                audio_queue=self._audio_queue_in,
                name="VoicegroundBridgeInput",
            )
        return self._input

    def output(self) -> FrameProcessor:
        """Get the output transport."""
        if not self._output:
            self._output = VoicegroundBridgeOutputTransport(
                self._params,
                audio_queue=self._audio_queue_out,
                name="VoicegroundBridgeOutput",
            )
        return self._output

    def create_paired_transport(
        self,
        params: VoicegroundBridgeTransportParams | None = None,
    ) -> "VoicegroundBridgeTransport":
        """Create a paired transport with swapped input/output queues."""
        if not self._is_primary:
            raise RuntimeError("Cannot create paired transport from a paired transport")

        return VoicegroundBridgeTransport(
            params=params or self._params,
            paired_audio_queue_in=self._audio_queue_out,
            paired_audio_queue_out=self._audio_queue_in,
        )

    async def close(self) -> None:
        """Clean up resources."""
        if self._is_primary:
            await self._audio_queue_in.put(None)
            await self._audio_queue_out.put(None)
