use std::io::Cursor; use std::path::PathBuf; use std::sync::OnceLock; use std::time::{Duration, Instant}; use dashmap::DashMap; use piper_rs::Piper; use tokio::sync::Mutex; 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); /// 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: OnceLock>, cache: DashMap, Instant)>, model_path: Option, } impl TtsState { pub fn new(model_path: Option) -> Self { Self { piper: OnceLock::new(), cache: DashMap::new(), 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<&Mutex> { 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()) } /// 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); self.cache.remove(text); None } else { Some(entry.0.clone()) } } /// Store audio bytes in the cache. fn cache_store(&self, text: String, audio: Vec) { self.cache.insert(text, (audio, Instant::now())); } /// Evict all expired cache entries. fn cache_evict(&self) { self.cache .retain(|_, (_, inserted_at)| inserted_at.elapsed() <= CACHE_TTL); } /// 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> { // 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(); 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 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() }