path: root/src/tts.rs

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<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,
        }
    }

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

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

    /// 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<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();

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