siprouter/rust/crates/proxy-engine/src/webrtc_engine.rs

//! WebRTC engine — manages browser PeerConnections.
//!
//! Audio bridging is now channel-based:
//! - Browser Opus audio → on_track → mixer inbound channel
//! - Mixer outbound channel → Opus RTP → TrackLocalStaticRTP → browser
//!
//! The mixer handles all transcoding. The WebRTC engine just shuttles raw Opus.

use crate::ipc::{emit_event, OutTx};
use crate::mixer::RtpPacket;
use codec_lib::PT_OPUS;
use std::collections::HashMap;
use std::sync::Arc;
use tokio::sync::{mpsc, Mutex};
use webrtc::api::media_engine::MediaEngine;
use webrtc::api::APIBuilder;
use webrtc::ice_transport::ice_candidate::RTCIceCandidateInit;
use webrtc::peer_connection::configuration::RTCConfiguration;
use webrtc::peer_connection::peer_connection_state::RTCPeerConnectionState;
use webrtc::peer_connection::sdp::session_description::RTCSessionDescription;
use webrtc::peer_connection::RTCPeerConnection;
use webrtc::rtp_transceiver::rtp_codec::RTCRtpCodecCapability;
use webrtc::track::track_local::track_local_static_rtp::TrackLocalStaticRTP;
use webrtc::track::track_local::{TrackLocal, TrackLocalWriter};

/// A managed WebRTC session.
struct WebRtcSession {
    pc: Arc<RTCPeerConnection>,
    local_track: Arc<TrackLocalStaticRTP>,
    call_id: Option<String>,
    /// Channel sender for forwarding browser Opus audio to the mixer.
    /// Set when the session is linked to a call via link_to_mixer().
    mixer_tx: Arc<Mutex<Option<mpsc::Sender<RtpPacket>>>>,
}

/// Manages all WebRTC sessions.
pub struct WebRtcEngine {
    sessions: HashMap<String, WebRtcSession>,
    out_tx: OutTx,
}

impl WebRtcEngine {
    pub fn new(out_tx: OutTx) -> Self {
        Self {
            sessions: HashMap::new(),
            out_tx,
        }
    }

    /// Handle a WebRTC offer from a browser — create PeerConnection, return SDP answer.
    pub async fn handle_offer(
        &mut self,
        session_id: &str,
        offer_sdp: &str,
    ) -> Result<String, String> {
        let mut media_engine = MediaEngine::default();
        media_engine
            .register_default_codecs()
            .map_err(|e| format!("register codecs: {e}"))?;

        let api = APIBuilder::new()
            .with_media_engine(media_engine)
            .build();

        let config = RTCConfiguration {
            ice_servers: vec![],
            ..Default::default()
        };

        let pc = api
            .new_peer_connection(config)
            .await
            .map_err(|e| format!("create peer connection: {e}"))?;
        let pc = Arc::new(pc);

        // Local audio track for sending audio to browser (Opus).
        let local_track = Arc::new(TrackLocalStaticRTP::new(
            RTCRtpCodecCapability {
                mime_type: "audio/opus".to_string(),
                clock_rate: 48000,
                channels: 1,
                ..Default::default()
            },
            "audio".to_string(),
            "siprouter".to_string(),
        ));

        let _sender = pc
            .add_track(local_track.clone() as Arc<dyn TrackLocal + Send + Sync>)
            .await
            .map_err(|e| format!("add track: {e}"))?;

        // Shared mixer channel sender (populated when linked to a call).
        let mixer_tx: Arc<Mutex<Option<mpsc::Sender<RtpPacket>>>> =
            Arc::new(Mutex::new(None));

        // ICE candidate handler.
        let out_tx_ice = self.out_tx.clone();
        let sid_ice = session_id.to_string();
        pc.on_ice_candidate(Box::new(move |candidate| {
            let out_tx = out_tx_ice.clone();
            let sid = sid_ice.clone();
            Box::pin(async move {
                if let Some(c) = candidate {
                    if let Ok(json) = c.to_json() {
                        emit_event(
                            &out_tx,
                            "webrtc_ice_candidate",
                            serde_json::json!({
                                "session_id": sid,
                                "candidate": json.candidate,
                                "sdp_mid": json.sdp_mid,
                                "sdp_mline_index": json.sdp_mline_index,
                            }),
                        );
                    }
                }
            })
        }));

        // Connection state handler.
        let out_tx_state = self.out_tx.clone();
        let sid_state = session_id.to_string();
        pc.on_peer_connection_state_change(Box::new(move |state| {
            let out_tx = out_tx_state.clone();
            let sid = sid_state.clone();
            Box::pin(async move {
                let state_str = match state {
                    RTCPeerConnectionState::Connected => "connected",
                    RTCPeerConnectionState::Disconnected => "disconnected",
                    RTCPeerConnectionState::Failed => "failed",
                    RTCPeerConnectionState::Closed => "closed",
                    RTCPeerConnectionState::New => "new",
                    RTCPeerConnectionState::Connecting => "connecting",
                    _ => "unknown",
                };
                emit_event(
                    &out_tx,
                    "webrtc_state",
                    serde_json::json!({ "session_id": sid, "state": state_str }),
                );
            })
        }));

        // Track handler — receives Opus audio from the browser.
        // Forwards raw Opus payload to the mixer channel (when linked).
        let out_tx_track = self.out_tx.clone();
        let sid_track = session_id.to_string();
        let mixer_tx_for_track = mixer_tx.clone();
        pc.on_track(Box::new(move |track, _receiver, _transceiver| {
            let out_tx = out_tx_track.clone();
            let sid = sid_track.clone();
            let mixer_tx = mixer_tx_for_track.clone();
            Box::pin(async move {
                let codec_info = track.codec();
                emit_event(
                    &out_tx,
                    "webrtc_track",
                    serde_json::json!({
                        "session_id": sid,
                        "kind": track.kind().to_string(),
                        "codec": codec_info.capability.mime_type,
                    }),
                );

                // Spawn browser→mixer forwarding task.
                tokio::spawn(browser_to_mixer_loop(track, mixer_tx, out_tx, sid));
            })
        }));

        // Set remote offer.
        let offer = RTCSessionDescription::offer(offer_sdp.to_string())
            .map_err(|e| format!("parse offer: {e}"))?;
        pc.set_remote_description(offer)
            .await
            .map_err(|e| format!("set remote description: {e}"))?;

        // Create answer.
        let answer = pc
            .create_answer(None)
            .await
            .map_err(|e| format!("create answer: {e}"))?;
        let answer_sdp = answer.sdp.clone();
        pc.set_local_description(answer)
            .await
            .map_err(|e| format!("set local description: {e}"))?;

        self.sessions.insert(
            session_id.to_string(),
            WebRtcSession {
                pc,
                local_track,
                call_id: None,
                mixer_tx,
            },
        );

        Ok(answer_sdp)
    }

    /// Link a WebRTC session to a call's mixer via channels.
    /// - `inbound_tx`: browser audio goes TO the mixer through this channel
    /// - `outbound_rx`: mixed audio comes FROM the mixer through this channel
    pub async fn link_to_mixer(
        &mut self,
        session_id: &str,
        call_id: &str,
        inbound_tx: mpsc::Sender<RtpPacket>,
        outbound_rx: mpsc::Receiver<Vec<u8>>,
    ) -> bool {
        let session = match self.sessions.get_mut(session_id) {
            Some(s) => s,
            None => return false,
        };

        session.call_id = Some(call_id.to_string());

        // Set the mixer sender so the on_track loop starts forwarding.
        {
            let mut tx = session.mixer_tx.lock().await;
            *tx = Some(inbound_tx);
        }

        // Spawn mixer→browser outbound task.
        let local_track = session.local_track.clone();
        tokio::spawn(mixer_to_browser_loop(outbound_rx, local_track));

        true
    }

    pub async fn add_ice_candidate(
        &self,
        session_id: &str,
        candidate: &str,
        sdp_mid: Option<&str>,
        sdp_mline_index: Option<u16>,
    ) -> Result<(), String> {
        let session = self
            .sessions
            .get(session_id)
            .ok_or_else(|| format!("session {session_id} not found"))?;

        let init = RTCIceCandidateInit {
            candidate: candidate.to_string(),
            sdp_mid: sdp_mid.map(|s| s.to_string()),
            sdp_mline_index,
            ..Default::default()
        };

        session
            .pc
            .add_ice_candidate(init)
            .await
            .map_err(|e| format!("add ICE: {e}"))
    }

    pub async fn close_session(&mut self, session_id: &str) -> Result<(), String> {
        if let Some(session) = self.sessions.remove(session_id) {
            session.pc.close().await.map_err(|e| format!("close: {e}"))?;
        }
        Ok(())
    }
}

/// Browser → Mixer audio forwarding loop.
/// Reads Opus RTP from the browser track, sends raw Opus payload to the mixer channel.
async fn browser_to_mixer_loop(
    track: Arc<webrtc::track::track_remote::TrackRemote>,
    mixer_tx: Arc<Mutex<Option<mpsc::Sender<RtpPacket>>>>,
    out_tx: OutTx,
    session_id: String,
) {
    let mut buf = vec![0u8; 1500];
    let mut count = 0u64;

    loop {
        match track.read(&mut buf).await {
            Ok((rtp_packet, _attributes)) => {
                count += 1;

                let payload = &rtp_packet.payload;
                if payload.is_empty() {
                    continue;
                }

                // Send raw Opus payload to mixer (if linked).
                let tx = mixer_tx.lock().await;
                if let Some(ref tx) = *tx {
                    let _ = tx
                        .send(RtpPacket {
                            payload: payload.to_vec(),
                            payload_type: PT_OPUS,
                            marker: false,
                            timestamp: 0,
                        })
                        .await;
                }
                drop(tx);

                if count == 1 || count == 50 || count % 500 == 0 {
                    emit_event(
                        &out_tx,
                        "webrtc_audio_rx",
                        serde_json::json!({
                            "session_id": session_id,
                            "direction": "browser_to_mixer",
                            "packet_count": count,
                        }),
                    );
                }
            }
            Err(_) => break, // Track ended.
        }
    }
}

/// Mixer → Browser audio forwarding loop.
/// Reads Opus-encoded RTP packets from the mixer and writes to the WebRTC track.
async fn mixer_to_browser_loop(
    mut outbound_rx: mpsc::Receiver<Vec<u8>>,
    local_track: Arc<TrackLocalStaticRTP>,
) {
    while let Some(rtp_data) = outbound_rx.recv().await {
        let _ = local_track.write(&rtp_data).await;
    }
}