#!/usr/bin/env python3
"""
TURN relay test for La Suite Meet + LiveKit.

Forces relay-only ICE transport (TRANSPORT_RELAY) to verify the TURN relay
data path works end-to-end. This simulates clients behind restrictive NATs
(CGNAT, symmetric NAT, corporate firewalls) that cannot establish direct
ICE connections and must use TURN relay.

The test:
  1. Verifies TURN server responds to STUN probe (UDP 443)
  2. Both users authenticate via Keycloak OIDC
  3. User 1 creates a room, User 2 joins
  4. Both connect to LiveKit with ice_transport_type=TRANSPORT_RELAY
  5. User 1 publishes a dummy audio track
  6. User 2 verifies audio frames arrive via relay
  7. Clean up

This test will FAIL if the TURN relay data path is broken (e.g., due to
docker-proxy source address mangling in Docker Swarm).

Prerequisites:
  pip install livekit requests

Usage:
  python3 webrtc-relay.py
"""

import asyncio
import os
import socket
import struct
import sys
import time
from pathlib import Path

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 = 30  # seconds — relay should connect quickly if it works at all


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)


def check_stun(host, port, timeout=5):
    """Send a STUN Binding Request and verify we get a valid response."""
    txn_id = os.urandom(12)
    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]
    return resp_type == 0x0101 and magic == 0x2112A442 and resp_txn == txn_id


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}"},
    )


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_relay_test(lk_url, token_a, token_b):
    """
    Connect two participants using TURN relay only.
    A publishes audio, B verifies reception.
    Returns (success, frames_received, publisher_identity, error_message).
    """
    ws_url = lk_url.replace("https://", "wss://").replace("http://", "ws://")

    # Force relay-only ICE — no direct host/srflx candidates
    relay_config = rtc.RtcConfiguration(
        ice_transport_type=rtc.IceTransportType.TRANSPORT_RELAY,
    )
    room_opts = rtc.RoomOptions(
        auto_subscribe=True,
        rtc_config=relay_config,
    )

    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)
        print("  Participant A connected (relay-only)")

        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)
        print("  Participant B connected (relay-only)")

        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 (relay path broken?)"
    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


def main():
    creds = load_credentials()
    meet_domain = resolve_domain("lasuite-meet")
    livekit_domain = f"livekit.{meet_domain}"
    meet_url = f"https://{meet_domain}"
    kc_url = f"https://{resolve_domain('keycloak')}"
    room_name = f"relay-test-{int(time.time())}"

    print("Testing TURN relay path: force relay-only ICE, publish audio, verify reception")
    print()

    # Check TURN is available
    print("Step 1: Verify TURN server responds ...")
    server_ip = socket.gethostbyname(livekit_domain)
    print(f"  LiveKit IP: {server_ip} ({livekit_domain})")

    if not check_stun(server_ip, 443):
        print(f"  FAIL: No STUN response from {server_ip}:443 — TURN server not reachable")
        print("  Cannot run relay test without TURN.")
        sys.exit(1)
    print(f"  PASS: TURN server responds on UDP {server_ip}:443")

    # 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"  LiveKit URL: {lk1['url']}")

    # Relay test
    print("Step 4: Testing TURN relay media flow (relay-only ICE) ...")

    success, frames, publisher, error = asyncio.run(
        run_relay_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: TURN relay media flow verified")
        print(f"  Audio frames received: {frames}")
        print(f"  Publisher identity: {publisher}")
        print("  Media was transported entirely via TURN relay (no direct ICE)")
    else:
        print(f"FAIL: TURN relay test failed — {error}")
        print()
        print("  The TURN server is reachable and relay candidates are generated,")
        print("  but media cannot flow through the relay path. This is typically")
        print("  caused by docker-proxy source address mangling in Docker Swarm:")
        print("  the TURN relay and SFU are in the same container but communicate")
        print("  via the external IP through docker-proxy, which changes the source")
        print("  address and breaks TURN permission checks.")
        sys.exit(1)


if __name__ == "__main__":
    main()