use std::io::Cursor; use std::path::PathBuf; use std::sync::atomic::{AtomicUsize, Ordering}; use std::sync::{Arc, Mutex}; use std::time::{Duration, Instant}; use dashmap::DashMap; use piper_rs::Piper; use tracing::{debug, info, warn}; use crate::error::{Error, Result}; // ── 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); /// Hard cap on cached audio so arbitrary `/tts` text cannot exhaust memory. const CACHE_MAX_ENTRIES: usize = 256; const CACHE_MAX_BYTES: usize = 64 * 1024 * 1024; /// Longest text accepted for synthesis. Keeps a single request (and its /// cached PCM) bounded. const MAX_TTS_CHARS: usize = 2_000; /// 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. /// /// 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: tokio::sync::OnceCell>>, cache: DashMap, Instant)>, cache_bytes: AtomicUsize, model_path: Option, } impl TtsState { pub fn new(model_path: Option) -> Self { Self { piper: tokio::sync::OnceCell::new(), cache: DashMap::new(), cache_bytes: AtomicUsize::new(0), model_path, } } /// Get or initialize the Piper model. Returns error if model path is /// unset or model fails to load. async fn get_piper(&self) -> Result>> { // `get_or_try_init` guarantees only one task loads the model even when // several `/tts` calls arrive at once. let piper = self .piper .get_or_try_init(|| async { 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" ); // Model loading is CPU/IO heavy; keep it off the async runtime. let piper = tokio::task::spawn_blocking(move || { Piper::new(&onnx_path, &config_path) .map_err(|e| Error::Custom(format!("Failed to load piper model: {e}"))) }) .await??; info!("piper model loaded successfully"); Ok::<_, Error>(Arc::new(Mutex::new(piper))) }) .await?; Ok(piper.clone()) } /// Look up cached audio for a text string. Returns `None` if uncached /// or expired. fn cache_lookup(&self, text: &str) -> Option> { let entry = self.cache.get(text)?; if entry.1.elapsed() > CACHE_TTL { drop(entry); if let Some((old, _)) = self.cache.remove(text) { self.cache_bytes.fetch_sub(old.len(), Ordering::Relaxed); } None } else { Some(entry.0.clone()) } } /// Store audio bytes in the cache, evicting expired/old entries first. fn cache_store(&self, text: String, audio: Vec) { let audio_len = audio.len(); if audio_len > CACHE_MAX_BYTES { return; } self.cache_evict(); while self.cache.len() >= CACHE_MAX_ENTRIES { if !self.cache_evict_oldest() { break; } } while self.cache_bytes.load(Ordering::Relaxed) + audio_len > CACHE_MAX_BYTES { if !self.cache_evict_oldest() { break; } } if let Some((old, _)) = self.cache.insert(text, (audio, Instant::now())) { self.cache_bytes.fetch_sub(old.len(), Ordering::Relaxed); } self.cache_bytes.fetch_add(audio_len, Ordering::Relaxed); } /// Evict all expired cache entries. fn cache_evict(&self) { let bytes = &self.cache_bytes; self.cache.retain(|_, (audio, inserted_at)| { let keep = inserted_at.elapsed() <= CACHE_TTL; if !keep { bytes.fetch_sub(audio.len(), Ordering::Relaxed); } keep }); } /// Evict the least-recently-inserted entry. Returns false when empty. fn cache_evict_oldest(&self) -> bool { let oldest = self .cache .iter() .min_by_key(|entry| entry.value().1) .map(|entry| entry.key().clone()); let Some(key) = oldest else { return false; }; if let Some((old, _)) = self.cache.remove(&key) { self.cache_bytes.fetch_sub(old.len(), Ordering::Relaxed); } true } /// 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> { if text.chars().count() > MAX_TTS_CHARS { return Err(Error::Custom(format!( "TTS text is too long (max {MAX_TTS_CHARS} characters)" ))); } // Cache hit if let Some(cached) = self.cache_lookup(text) { debug!(text_len = text.len(), "tts cache hit"); return Ok(cached); } // Cache miss — synthesize on a blocking thread. Piper inference is // synchronous and can take >100 ms; it must not stall a Tokio worker. let piper = self.get_piper().await?; let text_owned = text.to_string(); let started = Instant::now(); let (samples, sample_rate) = tokio::task::spawn_blocking(move || { let mut piper = piper .lock() .map_err(|_| Error::Custom("piper mutex poisoned".into()))?; piper .create(&text_owned, false, None, None, None, None) .map_err(|e| Error::Custom(format!("piper synthesis failed: {e}"))) }) .await??; let elapsed_ms = started.elapsed().as_millis() as u64; debug!( text_len = text.len(), samples = samples.len(), sample_rate, elapsed_ms, "piper synthesized" ); // 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) } /// 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 = 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(); } } // ── 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}") } // ── 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 { 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 { 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) -> 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() } #[cfg(test)] mod tests { use super::*; #[test] fn cache_is_bounded_by_entries_and_bytes() { let state = TtsState::new(None); for i in 0..(CACHE_MAX_ENTRIES + 32) { state.cache_store(format!("phrase-{i}"), vec![0u8; 1024]); } assert!(state.cache.len() <= CACHE_MAX_ENTRIES); assert!(state.cache_bytes.load(Ordering::Relaxed) <= CACHE_MAX_BYTES); } #[tokio::test] async fn overlong_text_is_rejected() { let state = TtsState::new(None); let text = "x".repeat(MAX_TTS_CHARS + 1); assert!(state.synthesize(&text).await.is_err()); } }