path: root/src/tts.rs

use std::io::Cursor;
use std::path::PathBuf;
use std::sync::atomic::{AtomicUsize, Ordering};
use std::sync::{Arc, Mutex, OnceLock};
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: OnceLock<Arc<Mutex<Piper>>>,
    cache: DashMap<String, (Vec<u8>, Instant)>,
    cache_bytes: AtomicUsize,
    model_path: Option<PathBuf>,
}

impl TtsState {
    pub fn new(model_path: Option<PathBuf>) -> Self {
        Self {
            piper: OnceLock::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<Arc<Mutex<Piper>>> {
        if let Some(piper) = self.piper.get() {
            return Ok(piper.clone());
        }

        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 too.
        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");

        // set() may fail if another task initialized first — that's fine
        let _ = self.piper.set(Arc::new(Mutex::new(piper)));
        Ok(self
            .piper
            .get()
            .expect("piper was just initialized")
            .clone())
    }

    /// 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);
            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<u8>) {
        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<Vec<u8>> {
        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<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();
    }
}

// ── 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<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()
}

#[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());
    }
}