"""WebRTC PeerConnection management using aiortc.

Creates a PeerConnection on demand when an offer arrives via MQTT,
attaches a video track from scrcpy H.264 stream, and handles DataChannel
for touch/key control signals.
"""

from __future__ import annotations

import asyncio
import json
import logging
from typing import Any, Callable, Awaitable

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

from .data_channel import DataChannelHandler

PrepareTrackCallback = Callable[[], Awaitable[Any]]

logger = logging.getLogger(__name__)


class WebRTCPeer:
    """Manages a single WebRTC PeerConnection for device streaming."""

    def __init__(
        self,
        ice_servers: list[dict[str, Any]] | None = None,
        on_viewer_connected: Callable[[], Awaitable[None]] | None = None,
        on_viewer_disconnected: Callable[[], Awaitable[None]] | None = None,
        on_data_channel_message: Callable[[dict[str, Any]], Awaitable[None]] | None = None,
        on_prepare_track: PrepareTrackCallback | None = None,
    ) -> None:
        self._ice_servers = ice_servers or []
        self._on_viewer_connected = on_viewer_connected
        self._on_viewer_disconnected = on_viewer_disconnected
        self._on_data_channel_message = on_data_channel_message
        self._on_prepare_track = on_prepare_track

        self._pc: RTCPeerConnection | None = None
        self._data_channel_handler: DataChannelHandler | None = None
        self._relay = MediaRelay()
        self._video_track: Any = None

    @property
    def connected(self) -> bool:
        return self._pc is not None and self._pc.connectionState == "connected"

    def set_video_track(self, track: Any) -> None:
        """Set the video track (from scrcpy adapter) to be sent to the peer."""
        self._video_track = track

    async def handle_offer(self, offer_payload: dict[str, Any]) -> dict[str, Any]:
        """Process a WebRTC offer and return an answer payload.

        If on_prepare_track is set, it is called first so that the video track
        (e.g. from scrcpy) is ready *before* the answer SDP is generated —
        ensuring the answer includes the video media section.
        """
        sdp = offer_payload.get("sdp", "")
        if not sdp:
            raise ValueError("Offer missing SDP")

        if self._on_prepare_track and not self._video_track:
            track = await self._on_prepare_track()
            if track:
                self._video_track = track

        if self._pc:
            logger.info("Closing existing PeerConnection before handling new offer")
            await self.close()

        ice_config = RTCConfiguration(
            iceServers=[
                RTCIceServer(
                    urls=s.get("urls", []),
                    username=s.get("username"),
                    credential=s.get("credential"),
                )
                for s in self._ice_servers
            ]
        )

        self._pc = RTCPeerConnection(ice_config)
        self._setup_event_handlers()

        if self._video_track:
            self._pc.addTrack(self._relay.subscribe(self._video_track))

        offer = RTCSessionDescription(sdp=sdp, type="offer")
        await self._pc.setRemoteDescription(offer)

        answer = await self._pc.createAnswer()
        await self._pc.setLocalDescription(answer)

        return {
            "type": "answer",
            "sdp": self._pc.localDescription.sdp,
        }

    async def handle_ice_candidate(self, candidate_payload: dict[str, Any]) -> None:
        """Add a remote ICE candidate to the PeerConnection."""
        if not self._pc:
            logger.warning("Received ICE candidate but no PeerConnection exists")
            return

        raw = candidate_payload.get("candidate")
        if not raw:
            return

        if isinstance(raw, str):
            sdp_mid = candidate_payload.get("sdpMid", "0")
            sdp_mline_index = candidate_payload.get("sdpMLineIndex", 0)
            try:
                from aiortc.sdp import candidate_from_sdp
                candidate_obj = candidate_from_sdp(raw.removeprefix("candidate:"))
                candidate_obj.sdpMid = sdp_mid
                candidate_obj.sdpMLineIndex = sdp_mline_index
            except Exception:
                logger.warning("Failed to parse ICE candidate SDP line, skipping")
                return
        elif isinstance(raw, dict):
            try:
                from aiortc.sdp import candidate_from_sdp
                sdp_line = raw.get("candidate", "")
                candidate_obj = candidate_from_sdp(sdp_line.removeprefix("candidate:"))
                candidate_obj.sdpMid = raw.get("sdpMid", "0")
                candidate_obj.sdpMLineIndex = raw.get("sdpMLineIndex", 0)
            except Exception:
                logger.warning("Failed to parse ICE candidate dict, skipping")
                return
        else:
            return

        logger.debug("Adding remote ICE candidate: %s", candidate_obj)
        await self._pc.addIceCandidate(candidate_obj)

    def _setup_event_handlers(self) -> None:
        if not self._pc:
            return

        @self._pc.on("connectionstatechange")
        async def on_state() -> None:
            assert self._pc is not None
            state = self._pc.connectionState
            logger.info("WebRTC connection state: %s", state)
            if state == "connected":
                if self._on_viewer_connected:
                    await self._on_viewer_connected()
            elif state in ("failed", "closed", "disconnected"):
                if self._on_viewer_disconnected:
                    await self._on_viewer_disconnected()

        @self._pc.on("datachannel")
        def on_datachannel(channel: Any) -> None:
            logger.info("DataChannel opened: %s", channel.label)
            self._data_channel_handler = DataChannelHandler(
                channel, self._on_data_channel_message
            )

    async def close(self) -> None:
        """Close the PeerConnection and release resources."""
        if self._pc:
            await self._pc.close()
            self._pc = None
            self._data_channel_handler = None
            self._video_track = None

    async def send_data(self, data: dict[str, Any]) -> None:
        """Send data through the DataChannel to the viewer."""
        if self._data_channel_handler:
            await self._data_channel_handler.send(data)
