kothok-edge-tts 0.2.2

//! `EdgeTts` — the concrete Edge Read-Aloud TTS engine.
//!
//! Implements [`Engine`] by opening a fresh WebSocket per synthesis request,
//! sending the `speech.config` and `ssml` messages, then collecting streamed
//! [`TtsEvent`]s until `turn.end`.

use crate::auth::{current_date_string, random_hex};
use crate::connection::Ws;
use crate::error::TtsError;
use crate::event::TtsEvent;
use crate::protocol;
use crate::ssml;
use crate::tts::Engine;
use std::future::Future;
use std::time::Duration;
use tokio_tungstenite::tungstenite::Message;

/// MP3 output format requested from the server.
const OUTPUT_FORMAT: &str = "audio-24khz-48kbitrate-mono-mp3";

/// If no message arrives within this window, the receive loop aborts with a
/// timeout error.  Prevents a missing `turn.end` from hanging the caller.
const RECEIVE_IDLE_TIMEOUT: Duration = Duration::from_secs(30);

/// X-RequestId byte length (random hex → 32 hex chars).
const REQ_ID_BYTES: usize = 16;

/// Path header value marking the end of a synthesis turn.
const PATH_TURN_END: &str = "turn.end";

/// Path header value marking a metadata JSON payload.
const PATH_AUDIO_METADATA: &str = "audio.metadata";

/// A stateless Edge-TTS engine.
///
/// Construct `EdgeTts` (a unit struct) and call [`Engine::synthesize`].
/// Each call opens its own WebSocket; there is no connection reuse.
///
/// ```no_run
/// # async fn run() -> Result<(), Box<dyn std::error::Error>> {
/// use kothok_edge_tts::Engine;
/// kothok_edge_tts::init_tls();
/// let events = kothok_edge_tts::EdgeTts
///     .synthesize("Hello world.", "en-US-EmmaMultilingualNeural", "+0%", "en-US")
///     .await?;
/// # Ok(()) }
/// ```
pub struct EdgeTts;

impl Engine for EdgeTts {
    fn synthesize(
        &self,
        text: &str,
        voice: &str,
        rate: &str,
        lang: &str,
    ) -> impl Future<Output = Result<Vec<TtsEvent>, TtsError>> + Send {
        let text = text.to_string();
        let voice = voice.to_string();
        let rate = rate.to_string();
        let lang = lang.to_string();
        async move { Self::run_synthesis(&text, &voice, &rate, &lang).await }
    }
}

impl EdgeTts {
    /// Full synthesis pipeline: connect → configure → request → receive.
    async fn run_synthesis(
        text: &str,
        voice: &str,
        rate: &str,
        lang: &str,
    ) -> Result<Vec<TtsEvent>, TtsError> {
        let mut ws = Ws::connect().await?;
        Self::send_config(&mut ws).await?;
        Self::send_ssml(&mut ws, text, voice, rate, lang).await?;
        Self::receive_events(&mut ws).await
    }

    /// Send the `speech.config` message (output format + metadata flags).
    async fn send_config(ws: &mut Ws) -> Result<(), TtsError> {
        let timestamp = current_date_string();
        let config = format!(
            "X-Timestamp:{timestamp}\r\n\
             Content-Type:application/json; charset=utf-8\r\n\
             Path:speech.config\r\n\r\n\
             {{\"context\":{{\"synthesis\":{{\"audio\":{{\
             \"metadataoptions\":{{\
             \"sentenceBoundaryEnabled\":\"false\",\
             \"wordBoundaryEnabled\":\"true\"}},\
             \"outputFormat\":\"{OUTPUT_FORMAT}\"}}}}}}}}\r\n"
        );
        ws.send_text(config).await
    }

    /// Send the SSML request for the utterance.
    async fn send_ssml(
        ws: &mut Ws,
        text: &str,
        voice: &str,
        rate: &str,
        lang: &str,
    ) -> Result<(), TtsError> {
        let request_id = random_hex(REQ_ID_BYTES);
        let timestamp = current_date_string();
        let ssml = ssml::build_ssml(text, voice, rate, lang);

        let message = format!(
            "X-RequestId:{request_id}\r\n\
             Content-Type:application/ssml+xml\r\n\
             X-Timestamp:{timestamp}Z\r\n\
             Path:ssml\r\n\r\n{ssml}"
        );
        ws.send_text(message).await
    }

    /// Receive messages until `turn.end` or stream close, collecting events.
    /// Returns [`TtsError::NoAudio`] if the turn ends without any audio frames.
    async fn receive_events(ws: &mut Ws) -> Result<Vec<TtsEvent>, TtsError> {
        let mut events = Vec::new();
        let mut got_audio = false;

        while let Some(msg) = ws.recv_timeout(RECEIVE_IDLE_TIMEOUT).await? {
            match msg {
                Message::Binary(bin) => {
                    log::debug!("edge-tts BIN len={}", bin.len());
                    if let Some(audio) = protocol::parse_binary_audio(&bin) {
                        got_audio = true;
                        events.push(TtsEvent::Audio(audio));
                    }
                }
                Message::Text(text) => {
                    if Self::handle_text_message(&text, &mut events) {
                        break;
                    }
                }
                Message::Close(_) => break,
                Message::Ping(data) => ws.send_pong(data).await,
                ref other => log::debug!("edge-tts OTHER {other:?}"),
            }
        }

        if !got_audio {
            return Err(TtsError::NoAudio);
        }
        Ok(events)
    }

    /// Process one text message.  Returns `true` if the caller should stop
    /// receiving (i.e. `turn.end` was received).
    fn handle_text_message(text: &str, events: &mut Vec<TtsEvent>) -> bool {
        let (headers, body) = protocol::split_msg(text);
        let path = headers.get("Path").map(String::as_str).unwrap_or("");

        log::debug!(
            "edge-tts TXT path={path} body[..160]={}",
            body.chars().take(160).collect::<String>()
        );

        match path {
            PATH_TURN_END => {
                events.push(TtsEvent::TurnEnd);
                true
            }
            PATH_AUDIO_METADATA => {
                events.extend(protocol::parse_word_boundaries(body));
                false
            }
            _ => false,
        }
    }
}