feat: piper now not a subprocess, cache system for TTS

6 files changed, 295 insertions, 154 deletions

diff --git a/src/commands/ai.rs b/src/commands/ai.rs
index 61fb471..5dd98cb 100644
--- a/src/commands/ai.rs
+++ b/src/commands/ai.rs
@@ -30,15 +30,14 @@ pub async fn senator(
             })
             .ok_or_else(lang::err_not_in_voice)?;
 
-        let (audio_data, _) =
-            crate::tts::synthesize(&ctx.data().config, &ctx.data().http_client, &response).await?;
+        let audio_data = ctx.data().tts.synthesize(&response).await?;
 
         let call = ctx
             .data()
             .voice_manager
             .connect(guild_id, channel_id)
             .await?;
-        let input = crate::voice::decode_audio(&audio_data)?;
+        let input = crate::tts::pcm_bytes_to_input(audio_data);
         let mut call_lock = call.lock().await;
         if call_lock.current_channel().is_some() {
             call_lock.play_input(input);
diff --git a/src/commands/tts.rs b/src/commands/tts.rs
index 8d20179..f45d2e8 100644
--- a/src/commands/tts.rs
+++ b/src/commands/tts.rs
@@ -15,7 +15,7 @@ pub enum OutputMode {
     Vc,
 }
 
-/// Generate TTS (Piper locally, Google as fallback) and output to file or voice channel
+/// Generate TTS (local Piper) and output to file or voice channel
 #[poise::command(slash_command)]
 pub async fn tts(
     ctx: Context<'_>,
@@ -25,21 +25,21 @@ pub async fn tts(
     ctx.defer().await?;
 
     let output = output.unwrap_or(OutputMode::Vc);
-    let (audio_data, ext) =
-        crate::tts::synthesize(&ctx.data().config, &ctx.data().http_client, &text).await?;
+    let audio_data = ctx.data().tts.synthesize(&text).await?;
 
     match output {
         OutputMode::File => {
+            let wav = crate::tts::encode_wav(&audio_data, 48000, 2);
+
             tokio::fs::create_dir_all(crate::config::TEMP_PATH)
                 .await
                 .ok();
             let path = format!(
-                "{}/tts_{}.{}",
+                "{}/tts_{}.wav",
                 crate::config::TEMP_PATH,
                 uuid::Uuid::new_v4(),
-                ext
             );
-            tokio::fs::write(&path, &audio_data).await?;
+            tokio::fs::write(&path, &wav).await?;
             if let Err(e) = ctx
                 .send(
                     poise::CreateReply::default()
@@ -69,7 +69,7 @@ pub async fn tts(
                 .voice_manager
                 .connect(guild_id, channel_id)
                 .await?;
-            let input = crate::voice::decode_audio(&audio_data)?;
+            let input = crate::tts::pcm_bytes_to_input(audio_data);
             let mut call_lock = call.lock().await;
             if call_lock.current_channel().is_some() {
                 call_lock.play_input(input);
diff --git a/src/events/mod.rs b/src/events/mod.rs
index e4f1b22..87333c8 100644
--- a/src/events/mod.rs
+++ b/src/events/mod.rs
@@ -28,15 +28,14 @@ fn format_duration(seconds: f64) -> String {
 }
 
 async fn play_tts_in_channel(
-    config: crate::config::Config,
-    http_client: reqwest::Client,
+    tts: std::sync::Arc<crate::tts::TtsState>,
     voice_manager: crate::voice::VoiceManager,
     guild_id: serenity::GuildId,
     channel_id: serenity::ChannelId,
     text: String,
 ) {
-    let audio_data = match crate::tts::synthesize(&config, &http_client, &text).await {
-        Ok((bytes, _)) => bytes,
+    let audio_data = match tts.synthesize(&text).await {
+        Ok(bytes) => bytes,
         Err(e) => {
             warn!("Failed to synthesize TTS: {e}");
             return;
@@ -51,13 +50,7 @@ async fn play_tts_in_channel(
         }
     };
 
-    let input = match crate::voice::decode_audio(&audio_data) {
-        Ok(i) => i,
-        Err(e) => {
-            warn!("Failed to decode TTS audio: {e}");
-            return;
-        }
-    };
+    let input = crate::tts::pcm_bytes_to_input(audio_data);
 
     let mut call_lock = call.lock().await;
     if call_lock.current_channel().is_some() {
@@ -97,6 +90,16 @@ pub async fn handle_voice_state_update(
         .map(|m| m.display_name().to_string())
         .unwrap_or_default();
 
+    // Pre-generate all 6 event phrases for this user (background, non-blocking).
+    // First event per user pays ~120ms; subsequent events are instant from cache.
+    {
+        let tts = data.tts.clone();
+        let name = display_name.clone();
+        tokio::spawn(async move {
+            tts.pre_generate(&name).await;
+        });
+    }
+
     let now = current_ts as f64;
 
     let user_id_str = user_id.to_string();
@@ -293,8 +296,7 @@ pub async fn handle_voice_state_update(
             && let Some(cid) = tts_channel.or(new_channel)
         {
             tokio::spawn(play_tts_in_channel(
-                data.config.clone(),
-                data.http_client.clone(),
+                data.tts.clone(),
                 data.voice_manager.clone(),
                 gid,
                 cid,
@@ -314,8 +316,7 @@ pub async fn handle_voice_state_update(
             && let Some(cid) = tts_channel.or(new_channel)
         {
             tokio::spawn(play_tts_in_channel(
-                data.config.clone(),
-                data.http_client.clone(),
+                data.tts.clone(),
                 data.voice_manager.clone(),
                 gid,
                 cid,
@@ -330,8 +331,7 @@ pub async fn handle_voice_state_update(
         && let Some(gid) = guild_id
     {
         tokio::spawn(play_tts_in_channel(
-            data.config.clone(),
-            data.http_client.clone(),
+            data.tts.clone(),
             data.voice_manager.clone(),
             gid,
             cid,
diff --git a/src/lang.rs b/src/lang.rs
index de0c7d7..53888eb 100644
--- a/src/lang.rs
+++ b/src/lang.rs
@@ -114,13 +114,6 @@ pub fn fmt_voice_join_failed(e: &dyn std::fmt::Display) -> String {
 pub const L_TTS_SAVED: &str = "TTS zapisany";
 pub const L_TTS_PLAYING: &str = "Grom TTS na kanale głosowym";
 
-pub fn fmt_tts_piper_launch(e: &str) -> String {
-    format!("Piper sie nie uruchomił w czas: `{e}`")
-}
-pub fn fmt_tts_piper_failed(e: &str) -> String {
-    format!("Piper TTS sie nie powiodło: `{e}`")
-}
-
 // ── Trigger command ───────────────────────────────────────────────────
 
 pub const L_TRIGGER_UPDATED: &str = "Trigger zaktualizowany";
diff --git a/src/main.rs b/src/main.rs
index 7f57421..5ab33ab 100644
--- a/src/main.rs
+++ b/src/main.rs
@@ -50,6 +50,8 @@ pub struct Data {
     pub ac_search_cache: AcSearchCache,
     /// Active music track per guild (only one music source at a time).
     pub music_tracks: Arc<tokio::sync::Mutex<HashMap<u64, songbird::tracks::TrackHandle>>>,
+    /// TTS engine: piper model + audio cache.
+    pub tts: Arc<tts::TtsState>,
 }
 
 impl Data {
@@ -187,7 +189,10 @@ async fn main() -> Result<(), Box<dyn std::error::Error>> {
 
                 let music_tracks: Arc<tokio::sync::Mutex<HashMap<u64, songbird::tracks::TrackHandle>>> = Arc::new(tokio::sync::Mutex::new(HashMap::new()));
 
+                let tts = Arc::new(tts::TtsState::new(config.piper_model_path()));
+
                 let data = Data {
+                    tts,
                     config,
                     db_pool,
                     voice_manager: vm,
diff --git a/src/tts.rs b/src/tts.rs
index ded76ad..3ae5140 100644
--- a/src/tts.rs
+++ b/src/tts.rs
@@ -1,147 +1,291 @@
+use std::io::Cursor;
 use std::path::PathBuf;
-use std::process::Stdio;
-use std::time::Instant;
+use std::sync::OnceLock;
+use std::time::{Duration, Instant};
 
-use tokio::io::AsyncWriteExt;
+use dashmap::DashMap;
+use piper_rs::Piper;
+use tokio::sync::Mutex;
 use tracing::{debug, info, warn};
 
-use crate::config;
-use crate::config::TEMP_PATH;
 use crate::error::{Error, Result};
 
-/// Synthesize speech for `text`.
-///
-/// Primary engine is local Piper TTS (via the `piper` CLI). If Piper is not
-/// installed, fails, or `TTS_MODEL_PATH` is unset, we fall back to Google
-/// Translate TTS over HTTP.
+// ── Constants ──────────────────────────────────────────────────────────
+
+/// Target sample rate for songbird playback.
+const TARGET_SAMPLE_RATE: u32 = 48_000;
+
+/// Cache entries older than this are evicted on access.
+const CACHE_TTL: Duration = Duration::from_secs(3600);
+
+/// All voice-state TTS event types.
+const EVENT_TYPES: &[&str] = &[
+    "joined",
+    "left",
+    "deafened",
+    "undeafened",
+    "muted",
+    "unmuted",
+];
+
+// ── TtsState ───────────────────────────────────────────────────────────
+
+/// Shared TTS state: piper model instance + audio cache.
 ///
-/// Returns the encoded audio bytes plus the file extension to use when saving
-/// (`"wav"` for Piper, `"mp3"` for Google).
-pub async fn synthesize(
-    config: &config::Config,
-    http_client: &reqwest::Client,
-    text: &str,
-) -> Result<(Vec<u8>, &'static str)> {
-    match config.piper_model_path() {
-        Some(model_path) => {
-            let started = Instant::now();
-            match piper_synthesize(&model_path, text).await {
-                Ok(bytes) => {
-                    info!(
-                        model = %model_path.display(),
-                        text_len = text.len(),
-                        bytes = bytes.len(),
-                        elapsed_ms = started.elapsed().as_millis() as u64,
-                        "Piper TTS synthesized"
-                    );
-                    return Ok((bytes, "wav"));
-                }
-                Err(e) => {
-                    warn!(
-                        model = %model_path.display(),
-                        text_len = text.len(),
-                        elapsed_ms = started.elapsed().as_millis() as u64,
-                        "Piper TTS failed, falling back to Google TTS: {e}"
-                    );
-                }
-            }
+/// Held in poise `Data`. Piper is loaded lazily on first synthesis call.
+/// Cache maps full text string → resampled 48 kHz stereo PCM bytes.
+pub struct TtsState {
+    piper: OnceLock<Mutex<Piper>>,
+    cache: DashMap<String, (Vec<u8>, Instant)>,
+    model_path: Option<PathBuf>,
+}
+
+impl TtsState {
+    pub fn new(model_path: Option<PathBuf>) -> Self {
+        Self {
+            piper: OnceLock::new(),
+            cache: DashMap::new(),
+            model_path,
         }
-        None => {
-            debug!("TTS_MODEL_PATH is not set; using Google TTS fallback");
+    }
+
+    /// Get or initialize the Piper model. Returns error if model path is
+    /// unset or model fails to load.
+    async fn get_piper(&self) -> Result<&Mutex<Piper>> {
+        if let Some(piper) = self.piper.get() {
+            return Ok(piper);
         }
+
+        let path = self
+            .model_path
+            .as_ref()
+            .ok_or_else(|| Error::Custom("TTS_MODEL_PATH is not set".into()))?;
+
+        // Resolve .onnx.json config path from model path
+        let onnx_path = path.clone();
+        let config_path = {
+            let s = path.to_string_lossy();
+            if s.ends_with(".onnx") {
+                PathBuf::from(format!("{}.json", s))
+            } else {
+                path.clone()
+            }
+        };
+
+        debug!(
+            onnx = %onnx_path.display(),
+            config = %config_path.display(),
+            "loading piper model"
+        );
+        let piper = Piper::new(&onnx_path, &config_path).map_err(|e| {
+            Error::Custom(format!("Failed to load piper model: {e}"))
+        })?;
+        info!("piper model loaded successfully");
+
+        // set() may fail if another task initialized first — that's fine
+        let _ = self.piper.set(Mutex::new(piper));
+        Ok(self.piper.get().unwrap())
     }
 
-    let started = Instant::now();
-    let bytes = google_synthesize(http_client, text).await?;
-    info!(
-        text_len = text.len(),
-        bytes = bytes.len(),
-        elapsed_ms = started.elapsed().as_millis() as u64,
-        "Google TTS synthesized"
-    );
-    Ok((bytes, "mp3"))
-}
+    /// Look up cached audio for a text string. Returns `None` if uncached
+    /// or expired.
+    fn cache_lookup(&self, text: &str) -> Option<Vec<u8>> {
+        let entry = self.cache.get(text)?;
+        if entry.1.elapsed() > CACHE_TTL {
+            drop(entry);
+            self.cache.remove(text);
+            None
+        } else {
+            Some(entry.0.clone())
+        }
+    }
+
+    /// Store audio bytes in the cache.
+    fn cache_store(&self, text: String, audio: Vec<u8>) {
+        self.cache.insert(text, (audio, Instant::now()));
+    }
 
-/// Run `piper` as a subprocess, feeding `text` on stdin and collecting the
-/// generated WAV from `--output_file`.
-async fn piper_synthesize(model_path: &PathBuf, text: &str) -> Result<Vec<u8>> {
-    tokio::fs::create_dir_all(TEMP_PATH).await.ok();
-
-    let out_path = format!("{}/tts_piper_{}.wav", TEMP_PATH, uuid::Uuid::new_v4());
-    let model = model_path.to_string_lossy().to_string();
-
-    debug!(model = %model, "spawning piper subprocess");
-    let started = Instant::now();
-
-    let mut child = tokio::process::Command::new("piper")
-        .arg("--model")
-        .arg(&model)
-        .arg("--output_file")
-        .arg(&out_path)
-        .stdin(Stdio::piped())
-        .stdout(Stdio::piped())
-        .stderr(Stdio::piped())
-        .spawn()
-        .map_err(|e| Error::Custom(crate::lang::fmt_tts_piper_launch(&e.to_string())))?;
-
-    {
-        let mut stdin = child
-            .stdin
-            .take()
-            .ok_or_else(|| Error::Custom("Piper stdin unavailable".into()))?;
-        stdin.write_all(text.as_bytes()).await?;
-        stdin.write_all(b"\n").await?;
+    /// Evict all expired cache entries.
+    fn cache_evict(&self) {
+        self.cache
+            .retain(|_, (_, inserted_at)| inserted_at.elapsed() <= CACHE_TTL);
     }
 
-    let output = child.wait_with_output().await?;
-    let elapsed_ms = started.elapsed().as_millis() as u64;
+    /// Synthesize text → resampled 48 kHz stereo PCM bytes.
+    ///
+    /// Checks cache first. On miss, runs piper inference, resamples, caches,
+    /// and returns.
+    pub async fn synthesize(&self, text: &str) -> Result<Vec<u8>> {
+        // Cache hit
+        if let Some(cached) = self.cache_lookup(text) {
+            debug!(text_len = text.len(), "tts cache hit");
+            return Ok(cached);
+        }
+
+        // Cache miss — synthesize
+        let piper = self.get_piper().await?;
+        let started = Instant::now();
 
-    if !output.status.success() {
-        let stderr = String::from_utf8_lossy(&output.stderr);
-        warn!(
-            model = %model,
+        let (samples, sample_rate) = {
+            let mut lock = piper.lock().await;
+            lock.create(text, false, None, None, None, None)
+                .map_err(|e| Error::Custom(format!("piper synthesis failed: {e}")))?
+        };
+
+        let elapsed_ms = started.elapsed().as_millis() as u64;
+        debug!(
+            text_len = text.len(),
+            samples = samples.len(),
+            sample_rate,
             elapsed_ms,
-            "piper exited with error: {}",
-            stderr.trim_end()
+            "piper synthesized"
         );
-        let _ = tokio::fs::remove_file(&out_path).await;
-        return Err(Error::Custom(crate::lang::fmt_tts_piper_failed(
-            stderr.trim_end(),
-        )));
+
+        // Resample mono → stereo at target sample rate
+        let pcm_bytes = resample_to_stereo_bytes(&samples, sample_rate, TARGET_SAMPLE_RATE);
+
+        self.cache_store(text.to_string(), pcm_bytes.clone());
+        Ok(pcm_bytes)
     }
 
-    let bytes = tokio::fs::read(&out_path).await?;
-    debug!(
-        model = %model,
-        bytes = bytes.len(),
-        elapsed_ms,
-        "piper succeeded"
-    );
-    let _ = tokio::fs::remove_file(&out_path).await;
-    Ok(bytes)
+    /// Pre-generate all 6 event phrases for a display name.
+    ///
+    /// Called on first voice event for a user. Synthesizes all phrases and
+    /// caches them so subsequent events are instant.
+    pub async fn pre_generate(&self, display_name: &str) {
+        let phrases: Vec<String> = EVENT_TYPES
+            .iter()
+            .map(|evt| make_event_phrase(display_name, evt))
+            .collect();
+
+        for phrase in &phrases {
+            if self.cache_lookup(phrase).is_some() {
+                continue;
+            }
+            if let Err(e) = self.synthesize(phrase).await {
+                warn!(
+                    display_name,
+                    phrase, "pre-generation failed: {e}"
+                );
+            }
+        }
+
+        // Evict expired entries while we're at it
+        self.cache_evict();
+    }
 }
 
-/// Fetch Google Translate TTS (Polish) over HTTP.
-async fn google_synthesize(client: &reqwest::Client, text: &str) -> Result<Vec<u8>> {
-    let tts_url = format!(
-        "https://translate.google.com/translate_tts?ie=UTF-8&q={}&tl=pl&client=tw-ob",
-        urlencoding(text)
-    );
-    debug!(url_len = tts_url.len(), "fetching Google TTS");
-    let bytes = client.get(&tts_url).send().await?.bytes().await?;
-    Ok(bytes.to_vec())
+// ── Phrase construction ────────────────────────────────────────────────
+
+/// Build the TTS phrase for a voice event + display name.
+pub fn make_event_phrase(display_name: &str, event: &str) -> String {
+    let suffix = match event {
+        "joined" => "sie zespawnił",
+        "left" => "jest chujem",
+        "deafened" => "jest bity przez ojca",
+        "undeafened" => "uciekł od ojca",
+        "muted" => "idzie robić jajecznice",
+        "unmuted" => "zrobił jajecznice",
+        _ => return format!("{display_name} {event}"),
+    };
+    format!("{display_name} {suffix}")
 }
 
-pub fn urlencoding(s: &str) -> String {
-    let mut out = String::with_capacity(s.len());
-    for &byte in s.as_bytes() {
-        match byte {
-            b'A'..=b'Z' | b'a'..=b'z' | b'0'..=b'9' | b'_' | b'-' | b'.' | b'~' => {
-                out.push(byte as char)
-            }
-            b' ' => out.push('+'),
-            _ => out.push_str(&format!("%{:02X}", byte)),
+// ── Resampling ─────────────────────────────────────────────────────────
+
+/// Resample mono f32 PCM from `src_rate` to `dst_rate`, duplicate to stereo,
+/// and return as little-endian bytes.
+fn resample_to_stereo_bytes(samples: &[f32], src_rate: u32, dst_rate: u32) -> Vec<u8> {
+    if src_rate == dst_rate {
+        // Just duplicate mono → stereo, no resampling needed
+        let mut out = Vec::with_capacity(samples.len() * 2 * 4);
+        for &s in samples {
+            out.extend_from_slice(&s.to_le_bytes());
+            out.extend_from_slice(&s.to_le_bytes());
         }
+        return out;
     }
+
+    let ratio = src_rate as f64 / dst_rate as f64;
+    let out_len = (samples.len() as f64 / ratio).ceil() as usize;
+    let mut out = Vec::with_capacity(out_len * 2 * 4);
+
+    for i in 0..out_len {
+        let src_idx = i as f64 * ratio;
+        let idx = src_idx as usize;
+        let frac = src_idx - idx as f64;
+
+        let sample = if idx + 1 < samples.len() {
+            // Linear interpolation
+            samples[idx] + (samples[idx + 1] - samples[idx]) * frac as f32
+        } else if idx < samples.len() {
+            samples[idx]
+        } else {
+            0.0
+        };
+
+        // Mono → stereo duplication
+        out.extend_from_slice(&sample.to_le_bytes());
+        out.extend_from_slice(&sample.to_le_bytes());
+    }
+
     out
 }
+
+// ── WAV encoding ───────────────────────────────────────────────────────
+
+/// Encode raw f32 stereo PCM bytes into a WAV file (16-bit PCM).
+///
+/// `pcm_f32_bytes` is the raw little-endian f32 interleaved stereo data
+/// (e.g. from `resample_to_stereo_bytes`).
+pub fn encode_wav(pcm_f32_bytes: &[u8], sample_rate: u32, channels: u16) -> Vec<u8> {
+    let bits_per_sample: u16 = 16;
+    let byte_rate = sample_rate * channels as u32 * bits_per_sample as u32 / 8;
+    let block_align = channels * bits_per_sample / 8;
+
+    // Convert f32 samples to i16
+    let sample_count = pcm_f32_bytes.len() / 4; // 4 bytes per f32
+    let mut pcm_i16 = Vec::with_capacity(sample_count * 2);
+    for chunk in pcm_f32_bytes.chunks_exact(4) {
+        let f = f32::from_le_bytes([chunk[0], chunk[1], chunk[2], chunk[3]]);
+        let i = (f * 32767.0).clamp(-32768.0, 32767.0) as i16;
+        pcm_i16.extend_from_slice(&i.to_le_bytes());
+    }
+
+    let data_size = pcm_i16.len() as u32;
+    let file_size = 36 + data_size;
+
+    let mut wav = Vec::with_capacity(44 + pcm_i16.len());
+
+    // RIFF header
+    wav.extend_from_slice(b"RIFF");
+    wav.extend_from_slice(&file_size.to_le_bytes());
+    wav.extend_from_slice(b"WAVE");
+
+    // fmt chunk
+    wav.extend_from_slice(b"fmt ");
+    wav.extend_from_slice(&16u32.to_le_bytes()); // chunk size
+    wav.extend_from_slice(&1u16.to_le_bytes()); // PCM format
+    wav.extend_from_slice(&channels.to_le_bytes());
+    wav.extend_from_slice(&sample_rate.to_le_bytes());
+    wav.extend_from_slice(&byte_rate.to_le_bytes());
+    wav.extend_from_slice(&block_align.to_le_bytes());
+    wav.extend_from_slice(&bits_per_sample.to_le_bytes());
+
+    // data chunk
+    wav.extend_from_slice(b"data");
+    wav.extend_from_slice(&data_size.to_le_bytes());
+    wav.extend_from_slice(&pcm_i16);
+
+    wav
+}
+
+// ── Songbird Input helper ──────────────────────────────────────────────
+
+/// Create a songbird `Input` from resampled f32 stereo PCM bytes.
+pub fn pcm_bytes_to_input(pcm_f32_bytes: Vec<u8>) -> songbird::input::Input {
+    use songbird::input::RawAdapter;
+    let cursor = Cursor::new(pcm_f32_bytes);
+    // 2 channels (stereo), 48 kHz — matches TARGET_SAMPLE_RATE
+    RawAdapter::new(cursor, TARGET_SAMPLE_RATE, 2).into()
+}