inferencelayer 0.2.4

//! The WebRTC media leg for `moshi-serve`: accept a browser's SDP offer, answer with an Opus
//! audio track, and pump audio both ways through the engine. Inbound RTP → Opus decode → 24 kHz
//! → engine; engine 24 kHz → Opus encode → outbound RTP. The transcript rides the browser's
//! `oai-events` data channel (the OpenAI Realtime convention). Feature-gated (`webrtc-media`).
//!
//! Transport-only and engine-agnostic: the caller wires the two channels to its engine loop, so
//! this module never touches Mimi/Moshi. The peer connection is kept alive by moving its `Arc`
//! into the spawned pump tasks; it tears down when the connection closes.

use std::sync::{Arc, Mutex};

use anyhow::{Context, Result};
use webrtc::api::APIBuilder;
use webrtc::api::interceptor_registry::register_default_interceptors;
use webrtc::api::media_engine::{MIME_TYPE_OPUS, MediaEngine};
use webrtc::data_channel::RTCDataChannel;
use webrtc::ice_transport::ice_server::RTCIceServer;
use webrtc::interceptor::registry::Registry;
use webrtc::media::Sample;
use webrtc::peer_connection::configuration::RTCConfiguration;
use webrtc::peer_connection::sdp::session_description::RTCSessionDescription;
use webrtc::rtp_transceiver::rtp_codec::RTCRtpCodecCapability;
use webrtc::track::track_local::track_local_static_sample::TrackLocalStaticSample;

use crate::opus_rtc::OpusRtc;

/// Engine → peer: an audio frame (24 kHz mono) or a transcript fragment.
pub enum RtcOut {
    Audio24(Vec<f32>),
    Text(String),
}

fn ensure_crypto_provider() {
    static INSTALL: std::sync::Once = std::sync::Once::new();
    INSTALL.call_once(|| {
        let _ = rustls::crypto::ring::default_provider().install_default();
    });
}

/// Accept an SDP `offer`, set up the Opus media path wired to `inbound_pcm24` (decoded mic
/// audio → engine) and `outbound` (engine audio + transcript → peer), and return the SDP answer.
/// The peer connection lives on inside the spawned pumps until the connection closes.
pub async fn accept_call(
    offer: &str,
    inbound_pcm24: std::sync::mpsc::Sender<Vec<f32>>,
    outbound: tokio::sync::mpsc::UnboundedReceiver<RtcOut>,
    on_closed: Box<dyn Fn() + Send + Sync>,
) -> Result<String> {
    ensure_crypto_provider();

    let mut media = MediaEngine::default();
    media.register_default_codecs().context("register codecs")?;
    let mut registry = Registry::new();
    registry = register_default_interceptors(registry, &mut media).context("interceptors")?;
    let api = APIBuilder::new()
        .with_media_engine(media)
        .with_interceptor_registry(registry)
        .build();

    // Google STUN in production; `OSFKB_WEBRTC_HERMETIC=1` uses host candidates only (loopback
    // tests — fast, no external network, deterministic).
    let ice_servers = if std::env::var("OSFKB_WEBRTC_HERMETIC").as_deref() == Ok("1") {
        vec![]
    } else {
        vec![RTCIceServer {
            urls: vec!["stun:stun.l.google.com:19302".to_owned()],
            ..Default::default()
        }]
    };
    let config = RTCConfiguration { ice_servers, ..Default::default() };
    let pc = Arc::new(api.new_peer_connection(config).await.context("peer connection")?);

    // Outbound Opus track (48 kHz stereo, browser-standard) — the engine's voice.
    let track = Arc::new(TrackLocalStaticSample::new(
        RTCRtpCodecCapability {
            mime_type: MIME_TYPE_OPUS.to_owned(),
            clock_rate: 48000,
            channels: 2,
            sdp_fmtp_line: "minptime=10;useinbandfec=1".to_owned(),
            ..Default::default()
        },
        "audio".to_owned(),
        "moshi-webrtc".to_owned(),
    ));
    pc.add_track(track.clone()).await.context("add track")?;

    // Inbound audio: decode RTP Opus → 24 kHz → engine. Spawned per remote track.
    let inbound = inbound_pcm24.clone();
    pc.on_track(Box::new(move |remote, _receiver, _transceiver| {
        let inbound = inbound.clone();
        Box::pin(async move {
            if remote.codec().capability.mime_type.to_lowercase() != MIME_TYPE_OPUS.to_lowercase() {
                eprintln!("webrtc: ignoring non-Opus inbound track {}", remote.codec().capability.mime_type);
                return;
            }
            tokio::spawn(async move {
                let dbg = std::env::var_os("MOSHI_RTC_DEBUG").is_some();
                // Listen-probe: MOSHI_RTC_DUMP=<path> appends the DECODED inbound audio as raw
                // s16le @ 24 kHz — literally what the server hears from the peer's mic.
                let mut dump = std::env::var_os("MOSHI_RTC_DUMP").and_then(|p| {
                    std::fs::OpenOptions::new().create(true).append(true).open(p).ok()
                });
                let mut codec = match OpusRtc::new() {
                    Ok(c) => c,
                    Err(e) => {
                        eprintln!("webrtc: inbound opus init: {e}");
                        return;
                    }
                };
                let mut n = 0u64;
                let t0 = std::time::Instant::now();
                while let Ok((pkt, _)) = remote.read_rtp().await {
                    if pkt.payload.is_empty() {
                        continue;
                    }
                    match codec.decode_to_24k(&pkt.payload) {
                        Ok(pcm) => {
                            n += 1;
                            if let Some(f) = dump.as_mut() {
                                use std::io::Write;
                                let b: Vec<u8> = pcm
                                    .iter()
                                    .flat_map(|s| {
                                        (((s.clamp(-1.0, 1.0)) * 32767.0) as i16).to_le_bytes()
                                    })
                                    .collect();
                                let _ = f.write_all(&b);
                            }
                            if dbg && n % 100 == 1 {
                                eprintln!(
                                    "webrtc dbg: inbound pkt #{n} ({} samples) at {:.1}s = {:.1} pkt/s",
                                    pcm.len(),
                                    t0.elapsed().as_secs_f64(),
                                    n as f64 / t0.elapsed().as_secs_f64().max(0.001)
                                );
                            }
                            if inbound.send(pcm).is_err() {
                                break; // engine gone
                            }
                        }
                        Err(e) => eprintln!("webrtc: opus decode: {e}"),
                    }
                }
                if dbg {
                    eprintln!("webrtc dbg: inbound loop ended after {n} pkts");
                }
            });
        })
    }));

    // The browser creates the `oai-events` data channel; capture it for transcript output.
    let events_dc: Arc<Mutex<Option<Arc<RTCDataChannel>>>> = Arc::new(Mutex::new(None));
    let dc_slot = events_dc.clone();
    pc.on_data_channel(Box::new(move |dc: Arc<RTCDataChannel>| {
        let dc_slot = dc_slot.clone();
        Box::pin(async move {
            if dc.label() == "oai-events" {
                *dc_slot.lock().unwrap() = Some(dc.clone());
                dc.on_open(Box::new(|| Box::pin(async {})));
            }
        })
    }));

    let on_closed = Arc::new(on_closed);
    pc.on_peer_connection_state_change(Box::new(move |state| {
        let on_closed = on_closed.clone();
        Box::pin(async move {
            eprintln!("webrtc: peer state {state}");
            use webrtc::peer_connection::peer_connection_state::RTCPeerConnectionState as S;
            if matches!(state, S::Disconnected | S::Failed | S::Closed) {
                on_closed();
            }
        })
    }));

    // Non-trickle handshake: apply the offer, answer, gather all ICE candidates, return.
    pc.set_remote_description(RTCSessionDescription::offer(offer.to_owned())?)
        .await
        .context("set remote (offer)")?;
    let answer = pc.create_answer(None).await.context("create answer")?;
    let mut gather = pc.gathering_complete_promise().await;
    pc.set_local_description(answer).await.context("set local (answer)")?;
    let _ = gather.recv().await;
    let sdp = pc
        .local_description()
        .await
        .context("no local description after gathering")?
        .sdp;

    // Outbound pump: engine audio → Opus → track; transcript → data channel. Owns `pc` (keeps
    // the connection alive) + `track`; ends when the engine closes `outbound`.
    tokio::spawn(async move {
        let _keep = pc; // hold the connection open for the life of this pump
        let mut dump_out = std::env::var_os("MOSHI_RTC_DUMP_OUT").and_then(|p| {
            std::fs::OpenOptions::new().create(true).append(true).open(p).ok()
        });
        let mut codec = match OpusRtc::new() {
            Ok(c) => c,
            Err(e) => {
                eprintln!("webrtc: outbound opus init: {e}");
                return;
            }
        };
        let dbg = std::env::var_os("MOSHI_RTC_DEBUG").is_some();
        let mut n_out = 0u64;
        let mut outbound = outbound;
        while let Some(msg) = outbound.recv().await {
            match msg {
                RtcOut::Audio24(pcm) => {
                    if let Some(f) = dump_out.as_mut() {
                        use std::io::Write;
                        let b: Vec<u8> = pcm
                            .iter()
                            .flat_map(|s| (((s.clamp(-1.0, 1.0)) * 32767.0) as i16).to_le_bytes())
                            .collect();
                        let _ = f.write_all(&b);
                    }
                    match codec.encode_from_24k(&pcm) {
                    Ok(packets) => {
                        n_out += 1;
                        if dbg && n_out % 25 == 1 {
                            eprintln!("webrtc dbg: outbound frame #{n_out} -> {} pkts", packets.len());
                        }
                        for p in packets {
                            let sample = Sample {
                                data: p.into(),
                                duration: std::time::Duration::from_millis(20),
                                ..Default::default()
                            };
                            if track.write_sample(&sample).await.is_err() {
                                return; // peer gone
                            }
                        }
                    }
                    Err(e) => eprintln!("webrtc: opus encode: {e}"),
                }
                }
                RtcOut::Text(t) => {
                    // bind on its own statement so the std MutexGuard drops before the await
                    let dc = events_dc.lock().unwrap().clone();
                    if let Some(dc) = dc {
                        // \x01-prefixed = the CALLER's own transcript (whisper sidecar), which
                        // rides the same channel under the GA input-transcription event name
                        let ev = match t.strip_prefix('\u{1}') {
                            Some(user) => serde_json::json!({
                                "type": "conversation.item.input_audio_transcription.completed",
                                "item_id": "item_user",
                                "transcript": user,
                            }),
                            None => serde_json::json!({
                                "type": "response.output_audio_transcript.delta",
                                "response_id": "resp_duplex",
                                "delta": t,
                            }),
                        }
                        .to_string();
                        let _ = dc.send_text(ev).await;
                    }
                }
            }
        }
    });

    Ok(sdp)
}