#!/usr/bin/env python3
"""
WebRTC media flow test for La Suite Meet + LiveKit.

Tests the full media path between two participants:
  1. Preliminary connectivity checks (TCP 7881, UDP 7882, UDP 443 TURN, WSS signaling)
  2. Both users authenticate via Keycloak OIDC
  3. User 1 creates a room via the Meet API
  4. User 2 joins the room
  5. Both connect to LiveKit via the SDK
  6. User 1 publishes a dummy audio track
  7. User 2 verifies audio frames arrive
  8. Clean up: disconnect both, delete the room

This verifies:
  - WSS signaling through Traefik (port 7880)
  - ICE negotiation (direct or TURN relay)
  - Media transport over host-exposed ports (TCP 7881 / UDP 7882)
  - TURN relay via UDP 443 for clients behind restrictive NATs

Network requirements:
  The test can succeed via direct ICE (ports 7881/7882) or via TURN relay
  (UDP 443). Users behind CGNAT/symmetric NAT will use TURN automatically.
  The preliminary checks report which paths are available.

Prerequisites:
  pip install livekit requests

Usage:
  python3 webrtc-media.py
"""

import asyncio
import json
import os
import socket
import ssl
import struct
import sys
import time
from pathlib import Path
from urllib.parse import urlparse

sys.path.insert(0, os.path.join(os.path.dirname(__file__), '..', '..', '..'))

import requests
from livekit import rtc
from utils.tests.helpers import resolve_domain

SAMPLE_RATE = 48000
NUM_CHANNELS = 1
SAMPLES_PER_CHANNEL = 480  # 10ms at 48kHz
TIMEOUT = 45  # seconds to wait for audio frames (allow time for ICE TCP fallback)


def load_credentials():
    """Load Keycloak test credentials from domain-specific TOML file."""
    import tomllib
    from utils.tests.helpers import resolve_domain_suffix
    suffix = resolve_domain_suffix()
    creds_file = Path(__file__).resolve().parent.parent / f"keycloak-test-credentials.{suffix}.toml"
    if not creds_file.exists():
        print(f"FAIL: Credentials file not found: {creds_file}")
        print("Run setup_keycloak_integration.py first.")
        sys.exit(1)

    with open(creds_file, 'rb') as f:
        return tomllib.load(f)


# ---------------------------------------------------------------------------
# Preliminary connectivity checks
# ---------------------------------------------------------------------------

def check_tcp(host, port, timeout=5):
    """Check if a TCP port is reachable."""
    s = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
    s.settimeout(timeout)
    try:
        s.connect((host, port))
        s.close()
        return True
    except (socket.timeout, ConnectionRefusedError, OSError):
        return False


def check_udp(host, port, timeout=5):
    """Check if a UDP port accepts packets (send-only, best effort)."""
    s = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
    s.settimeout(timeout)
    try:
        s.sendto(b"\x00", (host, port))
        s.close()
        return True
    except OSError:
        return False


def check_stun(host, port, timeout=5):
    """Send a STUN Binding Request and verify we get a valid response.

    This proves the TURN/STUN server is actually listening and responding,
    not just that a UDP send succeeded (which almost always does).
    """
    txn_id = os.urandom(12)
    # STUN Binding Request: type=0x0001, length=0, magic=0x2112A442
    msg = struct.pack("!HHI", 0x0001, 0, 0x2112A442) + txn_id

    s = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
    s.settimeout(timeout)
    try:
        s.sendto(msg, (host, port))
        data, _ = s.recvfrom(1024)
    except (socket.timeout, OSError):
        return False
    finally:
        s.close()

    if len(data) < 20:
        return False
    resp_type, _, magic = struct.unpack("!HHI", data[:8])
    resp_txn = data[8:20]
    # Valid STUN Binding Response: type=0x0101, correct magic + transaction ID
    return resp_type == 0x0101 and magic == 0x2112A442 and resp_txn == txn_id


def check_wss(host, port=443, timeout=5):
    """Check if a WSS (TLS) connection can be established."""
    s = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
    s.settimeout(timeout)
    ctx = ssl.create_default_context()
    ss = ctx.wrap_socket(s, server_hostname=host)
    try:
        ss.connect((host, port))
        ss.close()
        return True
    except (socket.timeout, ConnectionRefusedError, OSError, ssl.SSLError):
        return False


def run_connectivity_checks(meet_domain, livekit_domain):
    """Run preliminary connectivity checks. Returns (all_critical_pass, turn_available)."""
    server_ip = socket.gethostbyname(meet_domain)
    livekit_ip = socket.gethostbyname(livekit_domain)

    print(f"  Server IP: {server_ip} ({meet_domain})")
    print(f"  LiveKit IP: {livekit_ip} ({livekit_domain})")

    all_pass = True
    turn_available = False

    # HTTPS / Traefik
    if check_tcp(server_ip, 443):
        print(f"  PASS: TCP {server_ip}:443 (HTTPS/Traefik)")
    else:
        print(f"  FAIL: TCP {server_ip}:443 (HTTPS/Traefik) unreachable")
        all_pass = False

    # WSS signaling
    if check_wss(livekit_domain):
        print(f"  PASS: WSS {livekit_domain}:443 (LiveKit signaling)")
    else:
        print(f"  FAIL: WSS {livekit_domain}:443 (LiveKit signaling) unreachable")
        all_pass = False

    # LiveKit TCP media port
    if check_tcp(server_ip, 7881):
        print(f"  PASS: TCP {server_ip}:7881 (LiveKit media/TCP)")
    else:
        print(f"  WARN: TCP {server_ip}:7881 (LiveKit media/TCP) unreachable")

    # LiveKit UDP media port
    if check_udp(server_ip, 7882):
        print(f"  PASS: UDP {server_ip}:7882 (LiveKit media/UDP) send OK")
    else:
        print(f"  WARN: UDP {server_ip}:7882 (LiveKit media/UDP) send failed")

    # TURN/UDP port — use STUN probe to verify the server actually responds
    if check_stun(server_ip, 443):
        print(f"  PASS: UDP {server_ip}:443 (TURN/UDP) STUN response received")
        turn_available = True
    else:
        print(f"  WARN: UDP {server_ip}:443 (TURN/UDP) no STUN response — TURN relay unavailable")

    if turn_available:
        print("  => TURN relay available — clients behind restrictive NATs can connect")
    else:
        print("  => TURN relay NOT available — only direct ICE connections will work")

    return all_pass, turn_available


# ---------------------------------------------------------------------------
# API helpers
# ---------------------------------------------------------------------------

def get_oidc_token(kc_url, realm, client_id, client_secret, username, password):
    """Get an access token from Keycloak via direct access grant."""
    resp = requests.post(
        f"{kc_url}/realms/{realm}/protocol/openid-connect/token",
        data={
            "grant_type": "password",
            "client_id": client_id,
            "client_secret": client_secret,
            "username": username,
            "password": password,
            "scope": "openid email",
        },
    )
    resp.raise_for_status()
    return resp.json()["access_token"]


def create_room(meet_url, token, room_name):
    """Create a public room. Returns (room_id, livekit_config)."""
    resp = requests.post(
        f"{meet_url}/api/v1.0/rooms/",
        headers={"Authorization": f"Bearer {token}"},
        json={"name": room_name, "access_level": "public"},
    )
    if resp.status_code != 201:
        print(f"  FAIL: Room creation returned HTTP {resp.status_code}: {resp.text}")
        sys.exit(1)
    data = resp.json()
    return data["id"], data["livekit"]


def join_room(meet_url, token, room_id):
    """Join an existing room. Returns livekit config."""
    resp = requests.get(
        f"{meet_url}/api/v1.0/rooms/{room_id}/",
        headers={"Authorization": f"Bearer {token}"},
    )
    resp.raise_for_status()
    return resp.json()["livekit"]


def delete_room(meet_url, token, room_id):
    """Delete a room."""
    requests.delete(
        f"{meet_url}/api/v1.0/rooms/{room_id}/",
        headers={"Authorization": f"Bearer {token}"},
    )


# ---------------------------------------------------------------------------
# WebRTC test
# ---------------------------------------------------------------------------

async def push_audio(source, stop_event):
    """Push silent PCM frames until stopped."""
    frame = rtc.AudioFrame.create(SAMPLE_RATE, NUM_CHANNELS, SAMPLES_PER_CHANNEL)
    while not stop_event.is_set():
        await source.capture_frame(frame)


async def run_webrtc_test(lk_url, token_a, token_b):
    """
    Connect two participants to LiveKit. A publishes audio, B verifies reception.
    Returns (success, frames_received, publisher_identity, error_message).
    """
    ws_url = lk_url.replace("https://", "wss://").replace("http://", "ws://")

    room_opts = rtc.RoomOptions(auto_subscribe=True)

    room_a = rtc.Room()
    room_b = rtc.Room()
    stop_audio = asyncio.Event()
    track_received = asyncio.Event()
    result = {"frames": 0, "publisher": ""}

    @room_b.on("track_subscribed")
    def on_track_subscribed(track, publication, participant):
        result["publisher"] = participant.identity

        async def read_frames():
            stream = rtc.AudioStream(track)
            async for _ in stream:
                result["frames"] += 1
                if result["frames"] >= 5:
                    break
            await stream.aclose()
            track_received.set()

        asyncio.ensure_future(read_frames())

    try:
        # Participant A connects and publishes
        await room_a.connect(ws_url, token_a, options=room_opts)

        source = rtc.AudioSource(SAMPLE_RATE, NUM_CHANNELS)
        audio_track = rtc.LocalAudioTrack.create_audio_track("test-audio", source)
        options = rtc.TrackPublishOptions()
        options.source = rtc.TrackSource.SOURCE_MICROPHONE
        await room_a.local_participant.publish_track(audio_track, options)
        audio_task = asyncio.ensure_future(push_audio(source, stop_audio))

        await asyncio.sleep(1)

        # Participant B connects and waits for track
        await room_b.connect(ws_url, token_b, options=room_opts)

        await asyncio.wait_for(track_received.wait(), timeout=TIMEOUT)
        return True, result["frames"], result["publisher"], ""

    except asyncio.TimeoutError:
        return False, 0, "", "timed out waiting for audio frames"
    except rtc.ConnectError as e:
        return False, 0, "", f"LiveKit connection failed: {e}"
    except Exception as e:
        return False, 0, "", f"unexpected error: {e}"
    finally:
        stop_audio.set()
        try:
            await room_b.disconnect()
        except Exception:
            pass
        try:
            await room_a.disconnect()
        except Exception:
            pass


# ---------------------------------------------------------------------------
# Main
# ---------------------------------------------------------------------------

def main():
    creds = load_credentials()
    meet_domain = resolve_domain("lasuite-meet")
    # LiveKit signaling runs on its own subdomain (LIVEKIT_DOMAIN in the recipe
    # .env), which is `livekit-meet.<suffix>` — not a `livekit.` prefix on the
    # meet domain.
    from utils.tests.helpers import resolve_domain_suffix
    livekit_domain = f"livekit-meet.{resolve_domain_suffix()}"
    meet_url = f"https://{meet_domain}"
    kc_url = f"https://{resolve_domain('keycloak')}"
    room_name = f"webrtc-test-{int(time.time())}"

    print("Testing WebRTC media flow: publish audio, verify reception")
    print()

    # Preliminary checks
    print("Step 1: Connectivity checks ...")
    critical_pass, turn_available = run_connectivity_checks(meet_domain, livekit_domain)
    if not critical_pass:
        print()
        print("FAIL: Critical connectivity checks failed — cannot proceed with WebRTC test")
        sys.exit(1)

    # Authenticate
    print("Step 2: Authenticating both users ...")
    token1 = get_oidc_token(
        kc_url, creds["kc_realm"], creds["kc_client_id"], creds["kc_client_secret"],
        creds["kc_test_user"], creds["kc_test_pass"],
    )
    token2 = get_oidc_token(
        kc_url, creds["kc_realm"], creds["kc_client_id"], creds["kc_client_secret"],
        creds["kc_test_user2"], creds["kc_test_pass2"],
    )
    print("  PASS: Both users authenticated")

    # Create room
    print(f"Step 3: Creating room '{room_name}' ...")
    room_id, lk1 = create_room(meet_url, token1, room_name)
    lk2 = join_room(meet_url, token2, room_id)
    print(f"  PASS: Room created (id={room_id})")
    print(f"  PASS: LiveKit tokens obtained for both users")
    print(f"  LiveKit URL: {lk1['url']}")

    # WebRTC media test
    print("Step 4: Testing WebRTC media flow ...")
    print("  Connecting participant A, publishing audio ...")
    print("  Connecting participant B, waiting for audio frames ...")

    success, frames, publisher, error = asyncio.run(
        run_webrtc_test(lk1["url"], lk1["token"], lk2["token"])
    )

    # Clean up
    print("Step 5: Cleaning up ...")
    delete_room(meet_url, token1, room_id)
    print("  Room deleted")

    # Result
    print()
    if success:
        print("PASS: WebRTC media flow verified")
        print(f"  Audio frames received: {frames}")
        print(f"  Publisher identity: {publisher}")
        if turn_available:
            print("  TURN was available — media may have used relay or direct ICE")
            print("  (check LiveKit logs for connectionType to confirm relay vs direct)")
    else:
        print(f"FAIL: WebRTC media flow test failed — {error}")
        if "timed out" in error or "connection" in error.lower():
            print()
            if not turn_available:
                print("  Hint: Neither direct ICE (7881/7882) nor TURN relay (UDP 443)")
                print("  are reachable. If the server has TURN enabled, verify UDP 443")
                print("  is published and not blocked by firewall.")
            else:
                print("  Hint: TURN relay port (UDP 443) is reachable but media still")
                print("  failed. Check LiveKit logs to verify TURN server started.")
                print("  Run: ssh <server> docker service logs <stack>_livekit | grep -i turn")
            print()
            print("  Try running from a machine with less restrictive network access,")
            print("  or verify TURN is enabled in the LiveKit config.")
        sys.exit(1)


if __name__ == "__main__":
    main()