path: root/tools/stt-client/src/main.rs

//! STT WebSocket client for testing the Makima STT streaming endpoint.
//!
//! This tool reads an audio file and streams it to the server via WebSocket,
//! printing transcription results as they arrive. Large files are decoded
//! and streamed asynchronously without loading the entire file into memory.

use std::path::PathBuf;
use std::sync::atomic::{AtomicBool, Ordering};
use std::sync::Arc;

use anyhow::Result;
use clap::Parser;
use futures::{SinkExt, StreamExt};
use serde::{Deserialize, Serialize};
use tokio::sync::mpsc;
use tokio_tungstenite::{connect_async, tungstenite::Message};
use url::Url;

#[derive(Parser)]
#[command(name = "stt-client")]
#[command(about = "WebSocket client for testing the Makima STT streaming endpoint")]
struct Args {
    /// Audio file to stream (supports MP3, WAV, FLAC, OGG, AAC)
    #[arg(short, long)]
    file: PathBuf,

    /// Server WebSocket URL
    #[arg(short, long, default_value = "ws://localhost:8080/api/v1/listen")]
    url: String,

    /// Chunk size in milliseconds for streaming
    #[arg(short, long, default_value = "100")]
    chunk_ms: u32,

    /// Simulate real-time streaming (add delays between chunks)
    #[arg(long, default_value = "true")]
    realtime: bool,

    /// Show progress during streaming (may interleave with transcripts)
    #[arg(long, default_value = "false")]
    show_progress: bool,
}

#[derive(Serialize)]
#[serde(tag = "type", rename_all = "camelCase")]
enum ClientMessage {
    Start(StartMessage),
    Stop { reason: Option<String> },
}

#[derive(Serialize)]
#[serde(rename_all = "camelCase")]
struct StartMessage {
    sample_rate: u32,
    channels: u16,
    encoding: String,
}

#[derive(Deserialize, Debug)]
#[serde(tag = "type", rename_all = "camelCase")]
enum ServerMessage {
    Ready { session_id: String },
    Transcript(TranscriptMessage),
    Error { code: String, message: String },
    Stopped { reason: String },
}

#[derive(Deserialize, Debug)]
#[serde(rename_all = "camelCase")]
struct TranscriptMessage {
    speaker: String,
    start: f32,
    end: f32,
    text: String,
    is_final: bool,
}

/// Audio format information extracted from file header.
struct AudioFormat {
    sample_rate: u32,
    channels: u16,
}

/// A chunk of decoded audio samples.
struct AudioChunk {
    samples: Vec<f32>,
}

#[tokio::main]
async fn main() -> Result<()> {
    let args = Args::parse();

    // Probe audio file to get format info without decoding
    eprintln!("[INFO] Probing audio file: {:?}", args.file);
    let format = probe_audio_format(&args.file)?;
    eprintln!(
        "[INFO] Audio format: {}Hz, {} channel(s)",
        format.sample_rate, format.channels
    );

    // Connect to WebSocket
    let url = Url::parse(&args.url)?;
    eprintln!("[INFO] Connecting to {}...", url);
    let (ws_stream, _) = connect_async(url.as_str()).await?;
    let (mut write, mut read) = ws_stream.split();

    // Send start message
    let start_msg = ClientMessage::Start(StartMessage {
        sample_rate: format.sample_rate,
        channels: format.channels,
        encoding: "pcm32f".to_string(),
    });
    write
        .send(Message::Text(serde_json::to_string(&start_msg)?))
        .await?;
    eprintln!("[INFO] Sent start message");

    // Flag to signal when session is stopped
    let session_stopped = Arc::new(AtomicBool::new(false));
    let session_stopped_clone = session_stopped.clone();

    // Spawn task to receive and print messages
    let receiver = tokio::spawn(async move {
        while let Some(msg) = read.next().await {
            match msg {
                Ok(Message::Text(text)) => match serde_json::from_str::<ServerMessage>(&text) {
                    Ok(ServerMessage::Ready { session_id }) => {
                        eprintln!("[INFO] Session ready: {}", session_id);
                    }
                    Ok(ServerMessage::Transcript(t)) => {
                        let final_marker = if t.is_final { " [FINAL]" } else { "" };
                        println!(
                            "[{:.2}s - {:.2}s] {}: {}{}",
                            t.start, t.end, t.speaker, t.text, final_marker
                        );
                    }
                    Ok(ServerMessage::Error { code, message }) => {
                        eprintln!("[ERROR] {}: {}", code, message);
                    }
                    Ok(ServerMessage::Stopped { reason }) => {
                        eprintln!("[INFO] Session stopped: {}", reason);
                        session_stopped_clone.store(true, Ordering::SeqCst);
                        break;
                    }
                    Err(e) => {
                        eprintln!("[ERROR] Failed to parse message: {}", e);
                        eprintln!("[DEBUG] Raw message: {}", text);
                    }
                },
                Ok(Message::Close(_)) => {
                    eprintln!("[INFO] Connection closed by server");
                    break;
                }
                Err(e) => {
                    eprintln!("[ERROR] WebSocket error: {}", e);
                    break;
                }
                _ => {}
            }
        }
    });

    // Calculate chunk size in samples
    let chunk_samples =
        (format.sample_rate * args.chunk_ms / 1000) as usize * format.channels as usize;

    // Create channel for streaming decoded audio chunks
    let (audio_tx, mut audio_rx) = mpsc::channel::<AudioChunk>(32);

    // Spawn blocking task to decode audio file and send chunks
    let file_path = args.file.clone();
    let decoder_handle = tokio::task::spawn_blocking(move || {
        decode_audio_streaming(&file_path, chunk_samples, audio_tx)
    });

    eprintln!(
        "[INFO] Streaming audio in {} ms chunks ({} samples each)...",
        args.chunk_ms, chunk_samples
    );

    // Stream chunks as they're decoded
    let mut chunks_sent = 0usize;
    while let Some(chunk) = audio_rx.recv().await {
        // Convert f32 samples to bytes (little-endian)
        let bytes: Vec<u8> = chunk.samples.iter().flat_map(|&s| s.to_le_bytes()).collect();

        write.send(Message::Binary(bytes.into())).await?;
        chunks_sent += 1;

        // Progress indicator
        if args.show_progress && chunks_sent % 50 == 0 {
            eprintln!("[PROGRESS] {} chunks streamed", chunks_sent);
        }

        // Simulate real-time streaming if enabled
        if args.realtime {
            tokio::time::sleep(tokio::time::Duration::from_millis(args.chunk_ms as u64)).await;
        }
    }

    // Wait for decoder to finish and check for errors
    decoder_handle.await??;

    eprintln!("[INFO] Streaming complete: {} chunks sent", chunks_sent);

    // Send stop message
    let stop_msg = ClientMessage::Stop {
        reason: Some("end_of_file".to_string()),
    };
    write
        .send(Message::Text(serde_json::to_string(&stop_msg)?))
        .await?;
    eprintln!("[INFO] Sent stop message, waiting for final results...");

    // Wait for receiver to finish with a timeout
    let timeout = tokio::time::Duration::from_secs(30);
    match tokio::time::timeout(timeout, receiver).await {
        Ok(result) => {
            result?;
        }
        Err(_) => {
            if session_stopped.load(Ordering::SeqCst) {
                eprintln!("[INFO] Session completed");
            } else {
                eprintln!("[WARN] Timeout waiting for server response");
            }
        }
    }

    eprintln!("[INFO] Done!");
    Ok(())
}

/// Probe an audio file to extract format information without decoding.
fn probe_audio_format(path: &PathBuf) -> Result<AudioFormat> {
    use symphonia::core::codecs::CODEC_TYPE_NULL;
    use symphonia::core::formats::FormatOptions;
    use symphonia::core::io::MediaSourceStream;
    use symphonia::core::meta::MetadataOptions;
    use symphonia::core::probe::Hint;

    let file = std::fs::File::open(path)?;
    let mss = MediaSourceStream::new(Box::new(file), Default::default());

    let mut hint = Hint::new();
    if let Some(ext) = path.extension().and_then(|e| e.to_str()) {
        hint.with_extension(ext);
    }

    let probed = symphonia::default::get_probe().format(
        &hint,
        mss,
        &FormatOptions::default(),
        &MetadataOptions::default(),
    )?;

    let track = probed
        .format
        .tracks()
        .iter()
        .find(|t| t.codec_params.codec != CODEC_TYPE_NULL)
        .ok_or_else(|| anyhow::anyhow!("No audio track found"))?;

    let sample_rate = track.codec_params.sample_rate.unwrap_or(16000);
    let channels = track
        .codec_params
        .channels
        .map(|c| c.count() as u16)
        .unwrap_or(1);

    Ok(AudioFormat {
        sample_rate,
        channels,
    })
}

/// Decode audio file and stream chunks through the channel.
/// This runs in a blocking thread to avoid blocking the async runtime.
fn decode_audio_streaming(
    path: &PathBuf,
    chunk_samples: usize,
    tx: mpsc::Sender<AudioChunk>,
) -> Result<()> {
    use symphonia::core::codecs::{DecoderOptions, CODEC_TYPE_NULL};
    use symphonia::core::formats::FormatOptions;
    use symphonia::core::io::MediaSourceStream;
    use symphonia::core::meta::MetadataOptions;
    use symphonia::core::probe::Hint;

    let file = std::fs::File::open(path)?;
    let mss = MediaSourceStream::new(Box::new(file), Default::default());

    let mut hint = Hint::new();
    if let Some(ext) = path.extension().and_then(|e| e.to_str()) {
        hint.with_extension(ext);
    }

    let probed = symphonia::default::get_probe().format(
        &hint,
        mss,
        &FormatOptions::default(),
        &MetadataOptions::default(),
    )?;

    let mut format = probed.format;
    let track = format
        .tracks()
        .iter()
        .find(|t| t.codec_params.codec != CODEC_TYPE_NULL)
        .ok_or_else(|| anyhow::anyhow!("No audio track found"))?;

    let track_id = track.id;
    let mut decoder =
        symphonia::default::get_codecs().make(&track.codec_params, &DecoderOptions::default())?;

    // Buffer for accumulating samples until we have a full chunk
    let mut sample_buffer: Vec<f32> = Vec::with_capacity(chunk_samples * 2);

    loop {
        let packet = match format.next_packet() {
            Ok(p) => p,
            Err(symphonia::core::errors::Error::IoError(ref e))
                if e.kind() == std::io::ErrorKind::UnexpectedEof =>
            {
                break;
            }
            Err(symphonia::core::errors::Error::ResetRequired) => {
                decoder.reset();
                continue;
            }
            Err(_) => break,
        };

        if packet.track_id() != track_id {
            continue;
        }

        let decoded = match decoder.decode(&packet) {
            Ok(d) => d,
            Err(symphonia::core::errors::Error::DecodeError(_)) => continue,
            Err(_) => continue,
        };

        // Append decoded samples to buffer
        append_samples(&decoded, &mut sample_buffer);

        // Send complete chunks as they become available
        while sample_buffer.len() >= chunk_samples {
            let chunk: Vec<f32> = sample_buffer.drain(..chunk_samples).collect();
            if tx.blocking_send(AudioChunk { samples: chunk }).is_err() {
                // Receiver dropped, stop decoding
                return Ok(());
            }
        }
    }

    // Send any remaining samples as a final partial chunk
    if !sample_buffer.is_empty() {
        let _ = tx.blocking_send(AudioChunk {
            samples: sample_buffer,
        });
    }

    Ok(())
}

/// Append decoded audio samples to the output buffer.
fn append_samples(buffer: &symphonia::core::audio::AudioBufferRef, out: &mut Vec<f32>) {
    use symphonia::core::audio::{AudioBufferRef, Signal};

    match buffer {
        AudioBufferRef::U8(buf) => {
            for frame in 0..buf.frames() {
                for plane in buf.planes().planes() {
                    out.push((plane[frame] as f32 - 128.0) / 128.0);
                }
            }
        }
        AudioBufferRef::U16(buf) => {
            for frame in 0..buf.frames() {
                for plane in buf.planes().planes() {
                    out.push((plane[frame] as f32 - 32768.0) / 32768.0);
                }
            }
        }
        AudioBufferRef::U24(buf) => {
            for frame in 0..buf.frames() {
                for plane in buf.planes().planes() {
                    out.push((plane[frame].inner() as f32 - 8388608.0) / 8388608.0);
                }
            }
        }
        AudioBufferRef::U32(buf) => {
            for frame in 0..buf.frames() {
                for plane in buf.planes().planes() {
                    out.push((plane[frame] as f64 - 2147483648.0) as f32 / 2147483648.0);
                }
            }
        }
        AudioBufferRef::S8(buf) => {
            for frame in 0..buf.frames() {
                for plane in buf.planes().planes() {
                    out.push(plane[frame] as f32 / 128.0);
                }
            }
        }
        AudioBufferRef::S16(buf) => {
            for frame in 0..buf.frames() {
                for plane in buf.planes().planes() {
                    out.push(plane[frame] as f32 / 32768.0);
                }
            }
        }
        AudioBufferRef::S24(buf) => {
            for frame in 0..buf.frames() {
                for plane in buf.planes().planes() {
                    out.push(plane[frame].inner() as f32 / 8388608.0);
                }
            }
        }
        AudioBufferRef::S32(buf) => {
            for frame in 0..buf.frames() {
                for plane in buf.planes().planes() {
                    out.push(plane[frame] as f32 / 2147483648.0);
                }
            }
        }
        AudioBufferRef::F32(buf) => {
            for frame in 0..buf.frames() {
                for plane in buf.planes().planes() {
                    out.push(plane[frame]);
                }
            }
        }
        AudioBufferRef::F64(buf) => {
            for frame in 0..buf.frames() {
                for plane in buf.planes().planes() {
                    out.push(plane[frame] as f32);
                }
            }
        }
    }
}