1 files changed, 264 insertions, 120 deletions

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()
+}