guardian/crates/guardian_core/src/tts.rs

use std::io::Cursor;

use bevy::log;
use dasp::interpolate::sinc::Sinc;
use dasp::ring_buffer;
use dasp::signal;
use dasp::Sample;
use dasp::Signal;

use hound::WavSpec;
use hound::WavWriter;
use tokio::io::AsyncReadExt;
use tokio::io::AsyncWriteExt;
use tokio::net::TcpStream;
use tokio::net::ToSocketAddrs;

use serde::Deserialize;
use serde::Serialize;

#[derive(Debug, Deserialize)]
pub struct Packet {
    #[serde(rename = "type")]
    pub r#type: String,
    pub version: String,
    pub data_length: Option<usize>,
    pub payload_length: Option<usize>,
}

#[derive(Serialize)]
struct RequestData {
    text: String,
}

#[derive(Serialize)]
struct Request {
    r#type: String,
    data: RequestData,
}

pub async fn synthesize<A: ToSocketAddrs>(
    addr: A,
    text: &str,
) -> Result<Vec<Vec<u8>>, Box<dyn std::error::Error + Send + Sync + 'static>> {
    let stream = TcpStream::connect(addr).await?;

    let (mut read, mut write) = tokio::io::split(stream);

    let request = Request {
        r#type: "synthesize".to_string(),
        data: RequestData {
            text: text.to_string(),
        },
    };

    let mut request = serde_json::to_string(&request)?;
    request.push('\n');

    write.write_all(request.as_bytes()).await?;

    let mut str = String::new();
    let mut pcm_buffer: Vec<u8> = Vec::new();
    loop {
        let Ok(c) = read.read_u8().await else {
            log::error!("Unable to read byte");
            break;
        };

        if c == b'\n' {
            let Ok(packet) = serde_json::from_str::<Packet>(&str) else {
                log::error!("Unable to read packet from: {}", str);
                break;
            };
            str = String::new();

            if let Some(length) = packet.data_length {
                let mut buf = (0..length).map(|_| 0u8).collect::<Vec<_>>();

                if let Err(e) = read.read_exact(&mut buf).await {
                    log::error!("Unable to read data: {:?}", e);
                    break;
                }
            };

            if let Some(length) = packet.payload_length {
                let mut buf = (0..length).map(|_| 0u8).collect::<Vec<_>>();

                if let Err(e) = read.read_exact(&mut buf).await {
                    log::error!("Unable to read data: {:?}", e);
                    break;
                }

                pcm_buffer.append(&mut buf);
            }

            match packet.r#type.as_str() {
                "audio-start" => {
                    pcm_buffer.clear();
                }
                "audio-stop" => {
                    log::info!("Audio received, start post-processing");
                    break;
                }
                _ => {}
            }
        }
        str.push(c.into());
    }

    let wav = to_wav(&pcm_buffer)?;

    let Ok(reader) = hound::WavReader::new(&*wav) else {
        log::error!("Error opening reader");
        return Err("Error opening reader".into());
    };

    // pcm_buffer is 22050Hz, need to convert this to 16000Hz
    let samples = reader
        .into_samples()
        .filter_map(Result::ok)
        .map(i16::to_sample::<f64>);

    let signal = signal::from_interleaved_samples_iter(samples);

    let ring_buffer = ring_buffer::Fixed::from([[0.0]; 100]);
    let sinc = Sinc::new(ring_buffer);

    let new_signal = signal.from_hz_to_hz(sinc, 22050.0, 16000.0);

    let spec = WavSpec {
        channels: 1,
        sample_rate: 16000,
        bits_per_sample: 16,
        sample_format: hound::SampleFormat::Int,
    };

    let new_wav: Vec<u8> = Vec::new();
    let mut new_wav_cursor = Cursor::new(new_wav);

    let mut writer = WavWriter::new(&mut new_wav_cursor, spec).unwrap();

    for frame in new_signal.until_exhausted() {
        writer.write_sample(frame[0].to_sample::<i16>()).unwrap();
    }

    writer.flush().unwrap();

    drop(writer);

    let wav = new_wav_cursor.into_inner();
    let wav = bytes::Bytes::copy_from_slice(&wav);

    Ok(wav_to_opus(wav).await?)
}

async fn wav_to_opus(wav: bytes::Bytes) -> Result<Vec<Vec<u8>>, audiopus::Error> {
    use audiopus::coder::Encoder;
    use audiopus::{Application, Channels, SampleRate};

    tokio::task::spawn_blocking(move || {
        let audio_stream = wav
            .chunks(2)
            .map(|bytes| i16::from_le_bytes(bytes.try_into().unwrap()))
            .collect::<Vec<_>>();

        const MONO_20MS: usize = 16000 /* 1 channel */ * 20 / 1000;
        let enc = Encoder::new(SampleRate::Hz16000, Channels::Mono, Application::Voip)?;
        let mut pos = 0;
        let mut output = [0; 256];
        let mut frames = Vec::new();

        while pos + MONO_20MS < audio_stream.len() {
            let len = enc.encode(&audio_stream[pos..(pos + MONO_20MS)], &mut output)?;
            frames.push(output[..len].to_vec());

            pos += MONO_20MS;
        }

        Ok::<_, audiopus::Error>(frames)
    })
    .await
    .unwrap()
}

fn to_wav(data: &[u8]) -> Result<Vec<u8>, Box<dyn std::error::Error + Send + Sync + 'static>> {
    let input_buffer: Vec<u8> = Vec::new();
    let mut input_buffer_cursor = Cursor::new(input_buffer);

    let input_spec = WavSpec {
        channels: 1,
        sample_rate: 22050,
        bits_per_sample: 16,
        sample_format: hound::SampleFormat::Int,
    };

    let mut input = hound::WavWriter::new(&mut input_buffer_cursor, input_spec)?;

    let audio_stream = data
        .chunks(2)
        .map(|bytes| i16::from_le_bytes(bytes.try_into().unwrap()))
        .collect::<Vec<_>>();

    for frame in audio_stream {
        input.write_sample(frame.to_sample::<i16>()).unwrap();
    }

    drop(input);

    Ok(input_buffer_cursor.into_inner())
}