buzz/desktop/src-tauri/src/huddle/relay_api.rs

//! Relay HTTP helpers for huddle operations.
//!
//! Thin wrappers around the relay REST API for channel membership queries,
//! human participant counting, and the audio relay WebSocket connection.
//!
//! ```text
//! connect_audio_relay(channel_id)
//!   → WS /huddle/{id}/audio → challenge → NIP-42 auth → joined
//!   → send loop: pcm_rx → Opus encode → WS binary frame
//!   → recv loop: WS binary frame → Opus decode (per-peer) → rodio playback
//! ```

use futures_util::{SinkExt, StreamExt};
use reqwest::Method;
use serde::Deserialize;
use std::sync::{
    atomic::{AtomicBool, Ordering},
    Arc,
};
use tokio_tungstenite::{connect_async, tungstenite::Message as WsMsg};
use tokio_util::sync::CancellationToken;
use uuid::Uuid;

use crate::app_state::AppState;
use crate::relay::{api_path, build_authed_request, send_json_request};

/// Maximum number of agents that can be invited to a single huddle.
pub(crate) const MAX_HUDDLE_AGENTS: usize = 20;

/// Per-peer frame threshold: speech ≈ 25 frames/500ms, DTX noise ≈ 1.
pub(crate) const REMOTE_SPEECH_THRESHOLD: u16 = 5;

/// Validate that a string looks like a Nostr pubkey hex (64 hex chars).
pub(crate) fn validate_pubkey_hex(pubkey: &str) -> Result<(), String> {
    if pubkey.len() != 64 || !pubkey.chars().all(|c| c.is_ascii_hexdigit()) {
        let preview: String = pubkey.chars().take(16).collect();
        return Err(format!("invalid pubkey hex: {preview}"));
    }
    Ok(())
}

pub(crate) fn parse_channel_uuid(channel_id: &str) -> Result<Uuid, String> {
    Uuid::parse_str(channel_id).map_err(|_| format!("invalid channel UUID: {channel_id}"))
}

/// Handshake timeout — matches the server's AUTH_TIMEOUT (5 s).
const HANDSHAKE_TIMEOUT: std::time::Duration = std::time::Duration::from_secs(5);

/// Connect to the relay's audio WebSocket and run the Opus encode/decode pipeline.
///
/// Returns `(cancel_token, pcm_sender)` — caller stores both in `HuddleState`.
/// Dropping the sender or calling `cancel.cancel()` shuts down the relay task.
pub(crate) async fn connect_audio_relay(
    channel_id: &str,
    parent_channel_id: Option<&str>,
    state: &AppState,
) -> Result<(CancellationToken, tokio::sync::mpsc::Sender<Vec<u8>>), String> {
    use nostr::JsonUtil;

    let relay_url = crate::relay::relay_ws_url_with_override(state);
    let ws_url = format!("{relay_url}/huddle/{channel_id}/audio");

    let keys = state.keys.lock().map_err(|e| e.to_string())?.clone();

    // TTS interrupt flags — recv task cancels TTS when remote humans speak.
    let (tts_cancel, tts_active) = {
        let hs = state.huddle()?;
        (Arc::clone(&hs.tts_cancel), Arc::clone(&hs.tts_active))
    };

    let app_handle = state.app_handle.lock().ok().and_then(|g| g.clone());

    let (ws_stream, _) = connect_async(&ws_url)
        .await
        .map_err(|e| format!("audio WS connect failed: {e}"))?;
    let (mut ws_tx, mut ws_rx) = ws_stream.split();

    let challenge = tokio::time::timeout(HANDSHAKE_TIMEOUT, async {
        loop {
            match ws_rx.next().await {
                Some(Ok(WsMsg::Text(text))) => {
                    let v: serde_json::Value = serde_json::from_str(&text)
                        .map_err(|e| format!("bad challenge JSON: {e}"))?;
                    if v["type"] == "challenge" {
                        break v["challenge"]
                            .as_str()
                            .ok_or_else(|| "missing challenge string".to_string())
                            .map(|s| s.to_string());
                    }
                }
                Some(Ok(WsMsg::Close(_))) | None => {
                    break Err("connection closed before challenge".into());
                }
                _ => continue,
            }
        }
    })
    .await
    .map_err(|_| "timeout waiting for challenge from relay".to_string())?
    .map_err(|e: String| e)?;

    let tags = vec![
        nostr::Tag::parse(["relay", &relay_url]).map_err(|e| format!("tag relay: {e}"))?,
        nostr::Tag::parse(["challenge", &challenge]).map_err(|e| format!("tag challenge: {e}"))?,
    ];
    let event = nostr::EventBuilder::new(nostr::Kind::Custom(22242), "")
        .tags(tags)
        .sign_with_keys(&keys)
        .map_err(|e| format!("sign: {e}"))?;

    let event_json: serde_json::Value = serde_json::from_str(&event.as_json())
        .map_err(|e| format!("failed to serialize auth event: {e}"))?;
    let auth_msg = serde_json::json!({
        "type": "auth",
        "event": event_json,
        "parent_channel_id": parent_channel_id,
    });
    ws_tx
        .send(WsMsg::Text(auth_msg.to_string().into()))
        .await
        .map_err(|e| format!("send auth: {e}"))?;

    let initial_peers: Vec<(u8, String)> = tokio::time::timeout(HANDSHAKE_TIMEOUT, async {
        loop {
            match ws_rx.next().await {
                Some(Ok(WsMsg::Text(text))) => {
                    let v: serde_json::Value = serde_json::from_str(&text).unwrap_or_default();
                    match v["type"].as_str() {
                        Some("joined") => {
                            let peers = v["peers"]
                                .as_array()
                                .map(|arr| {
                                    arr.iter()
                                        .filter_map(|p| {
                                            Some((
                                                p["peer_index"].as_u64()? as u8,
                                                p["pubkey"].as_str()?.to_string(),
                                            ))
                                        })
                                        .collect::<Vec<_>>()
                                })
                                .unwrap_or_default();
                            break Ok(peers);
                        }
                        Some("error") => {
                            break Err(format!("audio relay auth error: {}", v["message"]));
                        }
                        _ => continue,
                    }
                }
                Some(Ok(WsMsg::Close(_))) | None => {
                    break Err("connection closed before joined".into());
                }
                _ => continue,
            }
        }
    })
    .await
    .map_err(|_| "timeout waiting for joined from relay".to_string())?
    .map_err(|e: String| e)?;

    let cancel = CancellationToken::new();
    let cancel_clone = cancel.clone();
    let (pcm_tx, pcm_rx) = tokio::sync::mpsc::channel::<Vec<u8>>(50);
    let output_device_name = state
        .audio_output_device
        .lock()
        .unwrap_or_else(|e| e.into_inner())
        .clone();

    tokio::spawn(async move {
        if let Err(e) = audio_relay_pipeline(
            ws_tx,
            ws_rx,
            pcm_rx,
            cancel_clone.clone(),
            app_handle.clone(),
            initial_peers,
            tts_cancel,
            tts_active,
            output_device_name,
        )
        .await
        {
            eprintln!("sprout-desktop: audio relay pipeline exited: {e}");
        }

        // Only emit the disconnect event for UNEXPECTED exits.
        // Skip if already cancelled (teardown_huddle in progress).
        if !cancel_clone.is_cancelled() {
            cancel_clone.cancel();
            if let Some(ref app) = app_handle {
                use tauri::Emitter;
                let _ = app.emit("huddle-audio-disconnected", ());
            }
        }
    });

    Ok((cancel, pcm_tx))
}

/// Background Opus encode/decode pipeline spawned by `connect_audio_relay`.
type WsStream =
    tokio_tungstenite::WebSocketStream<tokio_tungstenite::MaybeTlsStream<tokio::net::TcpStream>>;

async fn audio_relay_pipeline(
    ws_tx: futures_util::stream::SplitSink<WsStream, WsMsg>,
    mut ws_rx: futures_util::stream::SplitStream<WsStream>,
    mut pcm_rx: tokio::sync::mpsc::Receiver<Vec<u8>>,
    cancel: CancellationToken,
    app_handle: Option<tauri::AppHandle>,
    initial_peers: Vec<(u8, String)>,
    tts_cancel: Arc<AtomicBool>,
    tts_active: Arc<AtomicBool>,
    output_device_name: Option<String>,
) -> Result<(), String> {
    let mut encoder = opus::Encoder::new(48000, opus::Channels::Mono, opus::Application::Voip)
        .map_err(|e| format!("opus encoder: {e}"))?;
    encoder
        .set_bitrate(opus::Bitrate::Bits(32000))
        .map_err(|e| format!("opus bitrate: {e}"))?;
    encoder
        .set_dtx(true)
        .map_err(|e| format!("opus dtx: {e}"))?;

    let sink_handle = super::audio_output::open_output_sink_by_name(output_device_name.as_deref())?;
    let player = rodio::Player::connect_new(&sink_handle.mixer());

    let decoders: std::collections::HashMap<u8, opus::Decoder> = std::collections::HashMap::new();
    const FRAME_SAMPLES: usize = 960; // 20 ms at 48 kHz
    let decode_buf = vec![0f32; FRAME_SAMPLES];

    use std::sync::Arc as StdArc;
    let ws_tx = StdArc::new(tokio::sync::Mutex::new(ws_tx));
    let ws_tx_send = StdArc::clone(&ws_tx);
    let cancel_send = cancel.clone();

    let send_task = tokio::spawn(async move {
        let mut encoder = encoder; // Move encoder into task.
        const FRAME_SAMPLES: usize = 960;
        let mut out_buf = vec![0u8; 4000];

        loop {
            let pcm_bytes = {
                use futures_util::future::Either;
                let cancelled = std::pin::pin!(cancel_send.cancelled());
                let recv = std::pin::pin!(pcm_rx.recv());
                match futures_util::future::select(cancelled, recv).await {
                    Either::Left(_) => break, // Cancelled.
                    Either::Right((Some(b), _)) => b,
                    Either::Right((None, _)) => break, // Sender dropped.
                }
            };

            if pcm_bytes.len() % 4 != 0 {
                continue; // Malformed batch.
            }
            let samples: Vec<f32> = pcm_bytes
                .chunks_exact(4)
                .map(|b| f32::from_le_bytes([b[0], b[1], b[2], b[3]]))
                .collect();

            let mut tx = ws_tx_send.lock().await;
            for chunk in samples.chunks(FRAME_SAMPLES) {
                let encode_result = if chunk.len() == FRAME_SAMPLES {
                    encoder.encode_float(chunk, &mut out_buf)
                } else {
                    let mut padded = chunk.to_vec();
                    padded.resize(FRAME_SAMPLES, 0.0);
                    encoder.encode_float(&padded, &mut out_buf)
                };
                let n = match encode_result {
                    Ok(n) => n,
                    Err(e) => {
                        eprintln!("sprout-desktop: opus encode error: {e}");
                        continue;
                    }
                };
                if n > 0 {
                    // Send raw Opus bytes — the relay prepends the peer_index.
                    if tx
                        .send(WsMsg::Binary(out_buf[..n].to_vec().into()))
                        .await
                        .is_err()
                    {
                        return; // WS closed.
                    }
                }
            }
        }
        let mut tx = ws_tx_send.lock().await;
        let _ = tx.send(WsMsg::Close(None)).await;
    });

    let recv_task = tokio::spawn(async move {
        let mut decoders = decoders;
        let mut decode_buf = decode_buf;
        let cancel_recv = cancel;
        let ws_tx_recv = ws_tx;

        // Active speaker tracking: peer_index → pubkey, emit every 500ms.
        let mut index_to_pubkey: std::collections::HashMap<u8, String> =
            initial_peers.into_iter().collect();
        let mut active_indices: std::collections::HashSet<u8> = std::collections::HashSet::new();
        // Per-peer frame counter with Instant-based window (starvation-proof).
        let mut frame_counts: std::collections::HashMap<u8, u16> = std::collections::HashMap::new();
        let mut last_frame_reset = tokio::time::Instant::now();
        const FRAME_WINDOW: std::time::Duration = std::time::Duration::from_millis(500);
        let mut tts_was_active = false;
        let mut speaker_tick = tokio::time::interval(std::time::Duration::from_millis(500));
        speaker_tick.set_missed_tick_behavior(tokio::time::MissedTickBehavior::Skip);

        loop {
            use futures_util::future::Either;
            let cancelled = std::pin::pin!(cancel_recv.cancelled());
            let next = std::pin::pin!(ws_rx.next());
            let tick = std::pin::pin!(speaker_tick.tick());

            // Three-way select: cancel, WS message, or speaker tick.
            let ws_or_tick = std::pin::pin!(futures_util::future::select(next, tick));
            match futures_util::future::select(cancelled, ws_or_tick).await {
                Either::Left(_) => break, // Cancelled.
                Either::Right((Either::Right((_, _)), _)) => {
                    // Speaker tick: emit active speakers and reset.
                    if let Some(ref app) = app_handle {
                        use tauri::Emitter;
                        let pubkeys: Vec<String> = active_indices
                            .iter()
                            .filter_map(|idx| index_to_pubkey.get(idx).cloned())
                            .collect();
                        let _ = app.emit("huddle-active-speakers", &pubkeys);
                    }
                    active_indices.clear();
                    continue;
                }
                Either::Right((Either::Left((Some(Ok(msg)), _)), _)) => {
                    // WS message — process below.
                    match msg {
                        WsMsg::Binary(data) => {
                            if data.len() < 2 {
                                continue;
                            }
                            let peer_idx = data[0];
                            let opus_bytes = &data[1..];
                            active_indices.insert(peer_idx);

                            // Track TTS transitions at frame level (not decode level).
                            let tts_now = tts_active.load(Ordering::Acquire);
                            if tts_now && !tts_was_active {
                                frame_counts.clear();
                                last_frame_reset = tokio::time::Instant::now();
                            }
                            tts_was_active = tts_now;

                            let decoder = match decoders.entry(peer_idx) {
                                std::collections::hash_map::Entry::Occupied(e) => e.into_mut(),
                                std::collections::hash_map::Entry::Vacant(e) => {
                                    match opus::Decoder::new(48000, opus::Channels::Mono) {
                                        Ok(d) => e.insert(d),
                                        Err(err) => {
                                            eprintln!("sprout-desktop: opus decoder init peer {peer_idx}: {err}");
                                            continue;
                                        }
                                    }
                                }
                            };
                            match decoder.decode_float(opus_bytes, &mut decode_buf, false) {
                                Ok(n) if n > 0 => {
                                    if tts_now {
                                        if last_frame_reset.elapsed() >= FRAME_WINDOW {
                                            frame_counts.clear();
                                            last_frame_reset = tokio::time::Instant::now();
                                        }
                                        let count = frame_counts.entry(peer_idx).or_insert(0);
                                        *count = count.saturating_add(1);
                                        if *count >= REMOTE_SPEECH_THRESHOLD {
                                            tts_cancel.store(true, Ordering::Release);
                                        }
                                    }
                                    use rodio::buffer::SamplesBuffer;
                                    use std::num::NonZero;
                                    let channels = NonZero::new(1u16).unwrap();
                                    let rate = NonZero::new(48000u32).unwrap();
                                    player.append(SamplesBuffer::new(
                                        channels,
                                        rate,
                                        decode_buf[..n].to_vec(),
                                    ));
                                }
                                Ok(_) => {} // DTX silence — not counted.
                                Err(e) => {
                                    eprintln!("sprout-desktop: opus decode peer {peer_idx}: {e}");
                                }
                            }
                        }
                        WsMsg::Text(text) => {
                            if let Ok(v) = serde_json::from_str::<serde_json::Value>(&text) {
                                match v["type"].as_str() {
                                    Some("joined") => {
                                        if let Some(peers) = v["peers"].as_array() {
                                            for p in peers {
                                                if let (Some(pk), Some(idx)) =
                                                    (p["pubkey"].as_str(), p["peer_index"].as_u64())
                                                {
                                                    let key = idx as u8;
                                                    if index_to_pubkey.get(&key).map(|s| s.as_str())
                                                        != Some(pk)
                                                    {
                                                        decoders.remove(&key);
                                                        frame_counts.remove(&key);
                                                        active_indices.remove(&key);
                                                    }
                                                    index_to_pubkey.insert(key, pk.to_string());
                                                }
                                            }
                                        }
                                    }
                                    Some("left") => {
                                        if let Some(idx) = v["peer_index"].as_u64() {
                                            let key = idx as u8;
                                            index_to_pubkey.remove(&key);
                                            frame_counts.remove(&key);
                                            decoders.remove(&key);
                                        }
                                    }
                                    _ => {}
                                }
                            }
                        }
                        WsMsg::Ping(data) => {
                            let mut tx = ws_tx_recv.lock().await;
                            let _ = tx.send(WsMsg::Pong(data)).await;
                        }
                        WsMsg::Close(_) => break,
                        _ => {}
                    }
                }
                // WS error or closed.
                Either::Right((Either::Left(_), _)) => break,
            }
        }
    });

    // Wait for either task to finish, then abort the survivor.
    use futures_util::future::Either;
    match futures_util::future::select(std::pin::pin!(send_task), std::pin::pin!(recv_task)).await {
        Either::Left((_, recv_handle)) => recv_handle.abort(),
        Either::Right((_, send_handle)) => send_handle.abort(),
    }

    Ok(())
}

/// Fetch channel members with roles from the relay. Returns (pubkey, role) tuples.
pub(crate) async fn fetch_channel_members_with_roles(
    channel_id: &str,
    state: &AppState,
) -> Result<Vec<(String, Option<String>)>, String> {
    #[derive(Deserialize)]
    struct Member {
        pubkey: String,
        role: Option<String>,
    }
    #[derive(Deserialize)]
    struct MembersResponse {
        members: Vec<Member>,
    }

    let path = api_path(&["channels", channel_id, "members"]);
    let request = build_authed_request(&state.http_client, Method::GET, &path, state)?;
    let resp: MembersResponse = send_json_request(request).await.map_err(|e| {
        eprintln!("sprout-desktop: fetch channel members failed: {e}");
        e
    })?;

    Ok(resp
        .members
        .into_iter()
        .map(|m| (m.pubkey, m.role))
        .collect())
}

/// Fetch channel members, optionally filtered by role (e.g., "bot" for agents).
pub(crate) async fn fetch_channel_members(
    channel_id: &str,
    role_filter: Option<&str>,
    state: &AppState,
) -> Result<Vec<String>, String> {
    let all = fetch_channel_members_with_roles(channel_id, state).await?;
    Ok(all
        .into_iter()
        .filter(|(_, role)| role_filter.map_or(true, |r| role.as_deref() == Some(r)))
        .map(|(pubkey, _)| pubkey)
        .collect())
}

/// Count human (non-bot) members remaining in a channel.
pub(crate) async fn count_human_members(
    channel_id: &str,
    state: &AppState,
) -> Result<usize, String> {
    let all = fetch_channel_members_with_roles(channel_id, state).await?;
    Ok(all
        .iter()
        .filter(|(_, role)| role.as_deref() != Some("bot"))
        .count())
}