"""WebRTC signaling and connection management using aiortc."""

import asyncio
import json
import logging
import uuid

from aiortc import RTCPeerConnection, RTCSessionDescription
from aiortc.contrib.media import MediaRelay

from .tracks import AudioProcessorTrack

logger = logging.getLogger(__name__)

# Global state
peer_connections: dict[str, RTCPeerConnection] = {}
relay = MediaRelay()


async def create_offer_response(sdp: str, sdp_type: str, pipeline) -> dict:
    """Handle an SDP offer from the client and return an SDP answer.

    Args:
        sdp: The SDP offer string from the client.
        sdp_type: The SDP type (should be "offer").
        pipeline: The Pipeline instance for audio processing.

    Returns:
        Dict with session_id, sdp, and type for the answer.
    """
    session_id = str(uuid.uuid4())
    pc = RTCPeerConnection()
    peer_connections[session_id] = pc

    # Processor is created inside on_track (which fires synchronously during
    # setRemoteDescription) so that addTrack happens *before* createAnswer.
    # This is the canonical pattern from the official aiortc examples.
    processor_ref: list[AudioProcessorTrack] = []
    pending_channel: list = []

    def _attach_channel(processor: AudioProcessorTrack, channel) -> None:
        """Wire up the data channel to the processor and start streaming STT."""
        logger.info("Attaching data channel to processor (session=%s)", session_id)
        processor.data_channel = channel

        # Start streaming STT if available
        asyncio.ensure_future(processor.start_streaming_stt())

        @channel.on("message")
        def on_message(message):
            try:
                data = json.loads(message)
            except (json.JSONDecodeError, TypeError):
                return

            if data.get("type") == "playback_done":
                processor._barge_in.playback_done()
            elif data.get("type") in ("pause_audio", "resume_audio"):
                processor.handle_pause_resume(data["type"])
            elif data.get("type") == "ping":
                if channel.readyState == "open":
                    channel.send(json.dumps({"type": "pong"}))

    @pc.on("connectionstatechange")
    async def on_connectionstatechange():
        logger.info("Connection state: %s (session=%s)", pc.connectionState, session_id)
        if pc.connectionState in ("failed", "closed"):
            await cleanup_connection(session_id)

    @pc.on("track")
    def on_track(track):
        """Fires synchronously during setRemoteDescription — before createAnswer."""
        logger.info("Received track: %s (session=%s)", track.kind, session_id)
        if track.kind == "audio":
            processor = AudioProcessorTrack(pipeline, session_id=session_id)
            processor_ref.append(processor)

            # addTrack here — the transceiver already exists from the offer,
            # so aiortc upgrades it to sendrecv in the answer SDP.
            pc.addTrack(processor)

            # Wire inbound audio — starts _process_loop
            processor.set_input_track(relay.subscribe(track))

            # If data channel arrived first, attach it now
            if pending_channel:
                _attach_channel(processor, pending_channel[0])
                pending_channel.clear()

            @track.on("ended")
            async def on_ended():
                logger.info("Track ended (session=%s)", session_id)
                processor.stop()

    @pc.on("datachannel")
    def on_datachannel(channel):
        logger.info("Data channel received: %s (session=%s)", channel.label, session_id)
        if channel.label == "status":
            if processor_ref:
                _attach_channel(processor_ref[0], channel)
            else:
                logger.info("Data channel before track, buffering (session=%s)", session_id)
                pending_channel.append(channel)

    # setRemoteDescription fires on_track synchronously → processor created + addTrack called
    offer = RTCSessionDescription(sdp=sdp, type=sdp_type)
    await pc.setRemoteDescription(offer)

    # createAnswer now sees the sendrecv transceiver
    answer = await pc.createAnswer()
    await pc.setLocalDescription(answer)

    return {
        "session_id": session_id,
        "sdp": pc.localDescription.sdp,
        "type": pc.localDescription.type,
    }


async def cleanup_connection(session_id: str):
    """Close and remove a peer connection."""
    pc = peer_connections.pop(session_id, None)
    if pc:
        await pc.close()
        logger.info("Cleaned up connection: %s", session_id)


async def cleanup_all():
    """Close all peer connections. Called on server shutdown."""
    coros = [pc.close() for pc in peer_connections.values()]
    await asyncio.gather(*coros, return_exceptions=True)
    peer_connections.clear()
    logger.info("All peer connections closed")
