"""VDO.Ninja WebRTC client — receive video frames from a VDO.Ninja stream.

Connects to VDO.Ninja's signaling server (wss://wss.vdo.ninja:443),
performs the WebRTC handshake via aiortc, and exposes frames as PIL Images —
same interface as StreamClient so the compose loop can use either one.

    VDO.Ninja publisher (browser)
        │  WebRTC video track (VP8/H.264)
        ▼
    VDONinjaClient (this module)
        │  get_frame() → PIL.Image (non-blocking, thread-safe)
        ▼
    compose_canvas() crops regions from it

Supported URL formats:
    https://vdo.ninja/?view=eWWbtku                  → view stream "eWWbtku"
    https://vdo.ninja/?view=eWWbtku&codec=vp8        → with codec pref
    https://vdo.ninja/?view=eWWbtku&password=mypass   → with password
    vdo.ninja/?view=eWWbtku                           → without scheme

Dependencies (optional — only needed if using VDO.Ninja streams):
    pip install aiortc aiohttp cryptography

Protocol overview (VDO.Ninja signaling):
    1. Connect WebSocket to wss://wss.vdo.ninja:443
    2. Send {"join": hashedStreamID} to register as viewer
    3. Receive SDP offer from publisher, send SDP answer back
    4. Exchange ICE candidates
    5. Receive video track via WebRTC
    6. Decode frames to PIL Image
"""

import asyncio
import hashlib
import io
import json
import os
import re
import struct
import threading
import time
from typing import Optional, Tuple
from urllib.parse import parse_qs, urlparse

# ── Optional imports — fail gracefully ────────────────────────────────

_AIORTC_OK = False
_CRYPTO_OK = False

try:
    from aiortc import RTCPeerConnection, RTCSessionDescription, RTCIceCandidate, RTCConfiguration, RTCIceServer
    from aiortc.rtcrtpreceiver import RemoteStreamTrack
    import av
    _AIORTC_OK = True
except ImportError:
    pass

try:
    import aiohttp
except ImportError:
    aiohttp = None

try:
    from PIL import Image
except ImportError:
    Image = None

try:
    from cryptography.hazmat.primitives.ciphers import Cipher, algorithms, modes
    from cryptography.hazmat.primitives import padding as sym_padding
    from cryptography.hazmat.backends import default_backend
    _CRYPTO_OK = True
except ImportError:
    pass


# ── Constants ─────────────────────────────────────────────────────────

DEFAULT_WSS = "wss://wss.vdo.ninja:443"
DEFAULT_PASSWORD = "someEncryptionKey123"
DEFAULT_SALT = "vdo.ninja"

STUN_SERVERS = [
    "stun:stun.cloudflare.com:3478",
    "stun:stun.l.google.com:19302",
]


# ── URL parsing ───────────────────────────────────────────────────────

def parse_vdoninja_url(url: str) -> dict:
    """Extract view ID, password, codec, and server from a VDO.Ninja URL.

    Returns dict with keys: stream_id, password, codec, server, room
    """
    url = url.strip()
    if not url:
        return {}

    # Add scheme if missing
    if "://" not in url:
        url = "https://" + url

    parsed = urlparse(url)
    params = parse_qs(parsed.query)

    result = {}

    # Stream ID from ?view= parameter
    if "view" in params:
        result["stream_id"] = params["view"][0]
    elif "v" in params:
        result["stream_id"] = params["v"][0]

    # Password
    if "password" in params:
        pw = params["password"][0]
        if pw.lower() == "false":
            result["password"] = False
        else:
            result["password"] = pw
    elif "pw" in params:
        result["password"] = params["pw"][0]

    # Codec preference
    if "codec" in params:
        result["codec"] = params["codec"][0]

    # Custom WSS server
    if "wss" in params:
        result["server"] = params["wss"][0]

    # Room
    if "room" in params:
        result["room"] = params["room"][0]

    return result


def is_vdoninja_url(url: str) -> bool:
    """Check if a URL looks like a VDO.Ninja stream URL."""
    url = url.strip().lower()
    if "vdo.ninja" in url:
        return True
    # Also catch known alternative domains
    for domain in ("rtc.ninja", "backup.vdo.ninja", "isolated.vdo.ninja"):
        if domain in url:
            return True
    # Check for ?view= parameter with vdo.ninja-style format
    if "view=" in url:
        return True
    return False


# ── Encryption helpers (match VDO.Ninja SDK) ─────────────────────────

def _derive_key(password: str) -> bytes:
    """Derive a 256-bit AES key from password using SHA-256.

    VDO.Ninja uses the raw password string hashed with SHA-256 as the key.
    """
    return hashlib.sha256(password.encode("utf-8")).digest()


def _generate_hash(stream_id: str, salt: str, password: str) -> str:
    """Generate the 6-char hex suffix that VDO.Ninja appends to stream IDs.

    The hash is derived from stream_id + salt + password to create a
    unique identifier that viewers use to find publishers.
    """
    combined = stream_id + salt
    if password:
        combined += password
    h = hashlib.sha256(combined.encode("utf-8")).hexdigest()
    return h[:6]


def _encrypt_message(plaintext: str, password: str) -> Tuple[str, str]:
    """Encrypt a message using AES-CBC, matching VDO.Ninja's encryption.

    Returns (ciphertext_base64, iv_hex).
    """
    if not _CRYPTO_OK:
        return plaintext, ""

    key = _derive_key(password)
    iv = os.urandom(16)

    padder = sym_padding.PKCS7(128).padder()
    padded = padder.update(plaintext.encode("utf-8")) + padder.finalize()

    cipher = Cipher(algorithms.AES(key), modes.CBC(iv), backend=default_backend())
    encryptor = cipher.encryptor()
    ct = encryptor.update(padded) + encryptor.finalize()

    import base64
    ct_b64 = base64.b64encode(ct).decode("ascii")
    iv_hex = iv.hex()
    return ct_b64, iv_hex


def _decrypt_message(ciphertext_b64: str, iv_hex: str, password: str) -> str:
    """Decrypt a message encrypted by VDO.Ninja's AES-CBC encryption."""
    if not _CRYPTO_OK:
        return ciphertext_b64

    import base64
    key = _derive_key(password)
    iv = bytes.fromhex(iv_hex)
    ct = base64.b64decode(ciphertext_b64)

    cipher = Cipher(algorithms.AES(key), modes.CBC(iv), backend=default_backend())
    decryptor = cipher.decryptor()
    padded = decryptor.update(ct) + decryptor.finalize()

    unpadder = sym_padding.PKCS7(128).unpadder()
    plaintext = unpadder.update(padded) + unpadder.finalize()
    return plaintext.decode("utf-8")


# ── VDONinjaClient ───────────────────────────────────────────────────

class VDONinjaClient:
    """Connect to a VDO.Ninja stream and provide frames as PIL Images.

    Same public interface as StreamClient:
        start() / stop() / get_frame() / connected / frame_size

    Usage:
        client = VDONinjaClient("https://vdo.ninja/?view=eWWbtku")
        client.start()
        ...
        frame = client.get_frame()  # PIL Image or None
        ...
        client.stop()
    """

    def __init__(self, url: str, reconnect_delay: float = 3.0):
        if not _AIORTC_OK:
            raise ImportError(
                "aiortc is required for VDO.Ninja stream input. "
                "Install it with: pip install aiortc"
            )
        if aiohttp is None:
            raise ImportError(
                "aiohttp is required for VDO.Ninja stream input. "
                "Install it with: pip install aiohttp"
            )

        self._raw_url = url
        self._reconnect_delay = reconnect_delay

        # Parse URL parameters
        params = parse_vdoninja_url(url)
        self._stream_id = params.get("stream_id", "")
        if not self._stream_id:
            raise ValueError(
                f"Could not find stream ID (view=...) in URL: {url}"
            )

        # Password handling: False = no encryption, None/empty = default
        pw = params.get("password", None)
        if pw is False:
            self._password = None  # No encryption
        elif pw:
            self._password = pw
        else:
            self._password = DEFAULT_PASSWORD

        self._salt = DEFAULT_SALT
        self._codec = params.get("codec", "")
        self._wss_url = params.get("server", DEFAULT_WSS)
        self._room = params.get("room", "")

        # Compute the hashed stream ID used for signaling
        if self._password:
            self._hashcode = _generate_hash(self._stream_id, self._salt, self._password)
            self._hashed_stream_id = self._stream_id + self._hashcode
        else:
            self._hashcode = ""
            self._hashed_stream_id = self._stream_id

        # Frame state (thread-safe)
        self._frame: Optional["Image.Image"] = None
        self._frame_lock = threading.Lock()
        self._frame_size: Optional[tuple] = None
        self._frame_count: int = 0
        self._last_frame_time: float = 0.0

        # Connection state
        self._running = False
        self._connected = False
        self._thread: Optional[threading.Thread] = None
        self._loop: Optional[asyncio.AbstractEventLoop] = None
        self._error: str = ""
        self._pc: Optional["RTCPeerConnection"] = None

    # ── Public API (matches StreamClient) ─────────────────────────────

    def start(self):
        """Start the background WebRTC receiver thread."""
        if self._running:
            return
        self._running = True
        self._thread = threading.Thread(
            target=self._run_loop, daemon=True,
            name=f"VDONinjaClient-{self._stream_id}"
        )
        self._thread.start()

    def stop(self):
        """Stop the receiver and close the WebRTC connection."""
        self._running = False
        if self._loop and not self._loop.is_closed():
            self._loop.call_soon_threadsafe(self._loop.stop)
        if self._thread:
            self._thread.join(timeout=5.0)
            self._thread = None
        self._connected = False

    def get_frame(self) -> Optional["Image.Image"]:
        """Return the latest frame as a PIL Image, or None if no frame yet."""
        with self._frame_lock:
            return self._frame

    @property
    def connected(self) -> bool:
        return self._connected

    @property
    def is_running(self) -> bool:
        return self._running

    @property
    def url(self) -> str:
        return self._raw_url

    @property
    def raw_url(self) -> str:
        return self._raw_url

    @property
    def frame_size(self) -> Optional[tuple]:
        """(width, height) of the last received frame, or None."""
        return self._frame_size

    @property
    def frame_count(self) -> int:
        return self._frame_count

    @property
    def error(self) -> str:
        return self._error

    @property
    def stream_id(self) -> str:
        return self._stream_id

    # ── Background thread ─────────────────────────────────────────────

    def _run_loop(self):
        """Run the asyncio event loop in the background thread."""
        self._loop = asyncio.new_event_loop()
        asyncio.set_event_loop(self._loop)
        try:
            self._loop.run_until_complete(self._connect_loop())
        except Exception as e:
            self._error = str(e)
            print(f"[VDONinja] Fatal error: {e}")
        finally:
            try:
                self._loop.run_until_complete(self._cleanup())
            except Exception:
                pass
            try:
                self._loop.close()
            except Exception:
                pass
            self._loop = None
            self._connected = False

    async def _cleanup(self):
        """Close the peer connection if open."""
        if self._pc:
            try:
                await self._pc.close()
            except Exception:
                pass
            self._pc = None

    async def _connect_loop(self):
        """Connect, receive frames, reconnect on failure."""
        while self._running:
            try:
                await self._connect_and_receive()
            except asyncio.CancelledError:
                break
            except Exception as e:
                self._error = str(e)
                self._connected = False
                await self._cleanup()
                if self._running:
                    print(f"[VDONinja] Connection lost: {e}")
                    print(f"[VDONinja] Reconnecting in {self._reconnect_delay}s...")
                    await asyncio.sleep(self._reconnect_delay)

    async def _connect_and_receive(self):
        """Single connection attempt: signaling + WebRTC."""
        self._error = ""

        # Create WebRTC peer connection with STUN servers
        ice_servers = [RTCIceServer(urls=s) for s in STUN_SERVERS]
        config = RTCConfiguration(iceServers=ice_servers)
        self._pc = RTCPeerConnection(configuration=config)

        # Track the frame receiver task so we can wait on it
        frame_receiver_task = None

        @self._pc.on("track")
        def on_track(track):
            nonlocal frame_receiver_task
            if track.kind == "video":
                print(f"[VDONinja] Video track received from {self._stream_id}")
                self._connected = True
                frame_receiver_task = asyncio.ensure_future(
                    self._receive_frames(track)
                )

        @self._pc.on("connectionstatechange")
        async def on_connectionstatechange():
            state = self._pc.connectionState
            print(f"[VDONinja] Connection state: {state}")
            if state in ("failed", "closed"):
                self._connected = False

        # Connect to signaling server and perform handshake
        timeout = aiohttp.ClientTimeout(total=None, connect=10)
        async with aiohttp.ClientSession(timeout=timeout) as session:
            print(f"[VDONinja] Connecting to signaling server {self._wss_url}...")
            async with session.ws_connect(
                self._wss_url,
                heartbeat=30.0,
            ) as ws:
                print(f"[VDONinja] Connected to signaling server")

                # Send "play" request — this tells the server we want to VIEW this stream.
                # The server will then tell the publisher to send us an SDP offer.
                # (NOT "seed" — that's what a PUBLISHER sends to announce their stream)
                play_msg = {"request": "play", "streamID": self._hashed_stream_id}
                await ws.send_json(play_msg)
                print(f"[VDONinja] Sent play request for stream: {self._stream_id} (hashed: {self._hashed_stream_id})")

                # Process signaling messages
                async for msg in ws:
                    if not self._running:
                        break

                    if msg.type == aiohttp.WSMsgType.TEXT:
                        try:
                            data = json.loads(msg.data)
                        except json.JSONDecodeError:
                            continue

                        await self._handle_signaling(ws, data)

                    elif msg.type in (
                        aiohttp.WSMsgType.ERROR,
                        aiohttp.WSMsgType.CLOSE,
                        aiohttp.WSMsgType.CLOSING,
                        aiohttp.WSMsgType.CLOSED,
                    ):
                        break

                    # If we have a frame receiver, check if it's done
                    # (connection closed by remote)
                    if frame_receiver_task and frame_receiver_task.done():
                        break

        self._connected = False

    async def _handle_signaling(self, ws, data: dict):
        """Handle a signaling message from the VDO.Ninja server.

        VDO.Ninja signaling protocol (viewer side):
        1. We send: {"request": "play", "streamID": "hashedID"}
        2. Server tells publisher: {"request": "offerSDP", "UUID": ourUUID}
        3. Publisher sends us SDP offer (direct-routed message with "description" field)
        4. We send SDP answer back (direct-routed to publisher's UUID)
        5. ICE candidates are exchanged both ways
        """
        request = data.get("request", "")

        # Log all incoming signaling messages for debugging
        if request:
            print(f"[VDONinja] Signaling: request={request}")
        elif "description" in data:
            print(f"[VDONinja] Signaling: SDP message from UUID={data.get('UUID', '?')}")
        elif "candidate" in data:
            print(f"[VDONinja] Signaling: ICE candidate from UUID={data.get('UUID', '?')}")
        else:
            print(f"[VDONinja] Signaling: unknown message keys={list(data.keys())}")

        # Server alerts (e.g. "Stream ID not found")
        if request == "alert":
            msg = data.get("message", "")
            print(f"[VDONinja] Server alert: {msg}")
            self._error = f"Server: {msg}"
            return

        # Server error
        if request == "error":
            msg = data.get("message", "")
            code = data.get("code", "")
            print(f"[VDONinja] Server error: {msg} (code={code})")
            self._error = f"Server error: {msg}"
            return

        # SDP offer from the publisher (direct-routed message)
        if "description" in data or "sdp" in data:
            await self._handle_offer(ws, data)
            return

        # ICE candidate from publisher
        if "candidate" in data:
            await self._handle_ice_candidate(data)
            return

    async def _handle_offer(self, ws, data: dict):
        """Handle an SDP offer from the publisher."""
        desc_data = data.get("description") or data.get("sdp")
        if not desc_data:
            return

        # Decrypt if encrypted
        if isinstance(desc_data, str) and self._password:
            vector = data.get("vector", "")
            if vector:
                try:
                    desc_data = _decrypt_message(
                        desc_data, vector, self._password + self._salt
                    )
                    desc_data = json.loads(desc_data)
                except Exception as e:
                    print(f"[VDONinja] Failed to decrypt offer: {e}")
                    return
        elif isinstance(desc_data, str):
            try:
                desc_data = json.loads(desc_data)
            except json.JSONDecodeError:
                pass

        if isinstance(desc_data, dict):
            sdp_type = desc_data.get("type", "offer")
            sdp = desc_data.get("sdp", "")
        else:
            return

        if sdp_type != "offer":
            return

        print(f"[VDONinja] Received SDP offer from publisher")

        # Set remote description (the offer)
        offer = RTCSessionDescription(sdp=sdp, type="offer")
        await self._pc.setRemoteDescription(offer)

        # Create and set local description (our answer)
        answer = await self._pc.createAnswer()
        await self._pc.setLocalDescription(answer)

        # Send answer back via signaling
        answer_data = {
            "type": self._pc.localDescription.type,
            "sdp": self._pc.localDescription.sdp,
        }

        msg = {}
        if "UUID" in data:
            msg["UUID"] = data["UUID"]
        if "session" in data:
            msg["session"] = data["session"]

        answer_json = json.dumps(answer_data)

        if self._password:
            encrypted, vector = _encrypt_message(
                answer_json, self._password + self._salt
            )
            msg["description"] = encrypted
            msg["vector"] = vector
        else:
            msg["description"] = answer_data

        msg["streamID"] = self._hashed_stream_id

        await ws.send_json(msg)
        print(f"[VDONinja] Sent SDP answer")

    async def _handle_ice_candidate(self, data: dict):
        """Handle an ICE candidate from the publisher."""
        candidate_data = data.get("candidate")
        if not candidate_data:
            return

        # Decrypt if encrypted
        if isinstance(candidate_data, str) and self._password:
            vector = data.get("vector", "")
            if vector:
                try:
                    candidate_data = _decrypt_message(
                        candidate_data, vector, self._password + self._salt
                    )
                    candidate_data = json.loads(candidate_data)
                except Exception as e:
                    print(f"[VDONinja] Failed to decrypt ICE candidate: {e}")
                    return
            else:
                try:
                    candidate_data = json.loads(candidate_data)
                except (json.JSONDecodeError, TypeError):
                    pass

        if isinstance(candidate_data, dict):
            cand_str = candidate_data.get("candidate", "")
            sdp_mid = candidate_data.get("sdpMid", "")
            sdp_mline_index = candidate_data.get("sdpMLineIndex", 0)

            if cand_str:
                try:
                    candidate = RTCIceCandidate(
                        sdpMid=sdp_mid,
                        sdpMLineIndex=sdp_mline_index,
                        candidate=cand_str,
                    )
                    await self._pc.addIceCandidate(candidate)
                except Exception as e:
                    # Some candidates may be invalid, that's ok
                    pass

    async def _receive_frames(self, track):
        """Receive video frames from the WebRTC track and store as PIL Images."""
        print(f"[VDONinja] Starting frame receiver for {self._stream_id}")
        try:
            while self._running:
                try:
                    frame = await asyncio.wait_for(track.recv(), timeout=10.0)
                except asyncio.TimeoutError:
                    if not self._running:
                        break
                    # No frame for 10s — connection might be stale
                    if self._pc and self._pc.connectionState in ("failed", "closed"):
                        break
                    continue
                except Exception as e:
                    # Track ended or error
                    print(f"[VDONinja] Track recv error: {e}")
                    break

                # Convert av.VideoFrame → PIL Image
                try:
                    img = frame.to_image()  # Returns PIL Image
                    with self._frame_lock:
                        self._frame = img
                        self._frame_size = img.size
                    self._frame_count += 1
                    self._last_frame_time = time.monotonic()
                except Exception as e:
                    # Skip corrupted frames
                    pass

        except asyncio.CancelledError:
            pass
        finally:
            self._connected = False
            print(f"[VDONinja] Frame receiver stopped for {self._stream_id}")

    # ── Representation ────────────────────────────────────────────────

    def __repr__(self):
        state = "connected" if self._connected else "disconnected"
        size = f"{self._frame_size[0]}×{self._frame_size[1]}" if self._frame_size else "?"
        return f"VDONinjaClient({self._stream_id!r}, {state}, {size}, {self._frame_count} frames)"
