pmomusic/pmoflac/tests/flac_streams.rs

use std::{
    future::Future,
    io,
    path::PathBuf,
    pin::Pin,
    sync::{
        atomic::{AtomicBool, AtomicUsize, Ordering},
        Arc,
    },
    task::{Context, Poll},
    time::Duration,
};

use tokio::io::{AsyncRead, AsyncReadExt, ReadBuf};

use pmoflac::{decode_flac_stream, encode_flac_stream, EncoderOptions, FlacError, PcmFormat};

fn fixture(name: &str) -> PathBuf {
    PathBuf::from(env!("CARGO_MANIFEST_DIR"))
        .join("test_data")
        .join(name)
}

struct VecAsyncReader {
    data: Vec<u8>,
    pos: usize,
}

impl VecAsyncReader {
    fn new(data: Vec<u8>) -> Self {
        Self { data, pos: 0 }
    }
}

impl AsyncRead for VecAsyncReader {
    fn poll_read(
        mut self: Pin<&mut Self>,
        _cx: &mut Context<'_>,
        buf: &mut ReadBuf<'_>,
    ) -> Poll<io::Result<()>> {
        if self.pos >= self.data.len() {
            return Poll::Ready(Ok(()));
        }
        let remaining = &self.data[self.pos..];
        let to_copy = remaining.len().min(buf.remaining());
        if to_copy == 0 {
            return Poll::Ready(Ok(()));
        }
        buf.put_slice(&remaining[..to_copy]);
        self.pos += to_copy;
        Poll::Ready(Ok(()))
    }
}

#[tokio::test]
async fn decode_stream_info_16_44() -> Result<(), FlacError> {
    let bytes = std::fs::read(fixture(
        "1abaa2c7fb4302e20ac570e79857b700.32bits-44.1Khz.flac",
    ))?;
    let mut stream = decode_flac_stream(VecAsyncReader::new(bytes)).await?;
    let info = stream.info().clone();

    assert_eq!(info.sample_rate, 44_100);
    assert_eq!(info.channels, 2);
    assert_eq!(info.bits_per_sample, 16);

    let mut pcm = Vec::new();
    stream.read_to_end(&mut pcm).await?;
    assert!(!pcm.is_empty());
    stream.wait().await?;

    Ok(())
}

#[tokio::test]
#[ignore] // Slow test with large 24-bit/192kHz file. Run with: cargo test -- --ignored
async fn roundtrip_encode_decode_24_192() -> Result<(), FlacError> {
    let bytes = std::fs::read(fixture("Yuri-Korzunov_Movement_24bit-192kHz.flac"))?;
    let mut decoder = decode_flac_stream(VecAsyncReader::new(bytes)).await?;
    let info = decoder.info().clone();

    let mut pcm_bytes = Vec::new();
    decoder.read_to_end(&mut pcm_bytes).await?;
    decoder.wait().await?;

    assert!(!pcm_bytes.is_empty());
    assert_eq!(
        pcm_bytes.len() % (info.bytes_per_sample() * info.channels as usize),
        0
    );

    let format = PcmFormat {
        sample_rate: info.sample_rate,
        channels: info.channels,
        bits_per_sample: info.bits_per_sample,
    };
    let options = EncoderOptions {
        total_samples: info.total_samples,
        ..Default::default()
    };

    let pcm_reader = VecAsyncReader::new(pcm_bytes.clone());
    let mut encoder_stream = encode_flac_stream(pcm_reader, format, options).await?;
    let mut encoded_bytes = Vec::new();
    encoder_stream.read_to_end(&mut encoded_bytes).await?;
    encoder_stream.wait().await?;

    assert!(!encoded_bytes.is_empty());

    let flac_reader = VecAsyncReader::new(encoded_bytes);
    let mut decoder_roundtrip = decode_flac_stream(flac_reader).await?;
    let mut pcm_roundtrip = Vec::new();
    decoder_roundtrip.read_to_end(&mut pcm_roundtrip).await?;
    decoder_roundtrip.wait().await?;

    assert_eq!(pcm_roundtrip, pcm_bytes);

    Ok(())
}

/// SlowReader simulates a slow stream by introducing delays between reads.
/// This helps verify that the encoder/decoder truly streams data rather than
/// buffering everything before producing output.
struct SlowReader {
    data: Vec<u8>,
    pos: usize,
    chunk_size: usize,
    delay: Duration,
    chunks_read: Arc<AtomicUsize>,
    sleep: Option<Pin<Box<tokio::time::Sleep>>>,
}

impl SlowReader {
    fn new(data: Vec<u8>, chunk_size: usize, delay: Duration) -> Self {
        Self {
            data,
            pos: 0,
            chunk_size,
            delay,
            chunks_read: Arc::new(AtomicUsize::new(0)),
            sleep: None,
        }
    }

    fn chunks_read(&self) -> Arc<AtomicUsize> {
        self.chunks_read.clone()
    }
}

impl AsyncRead for SlowReader {
    fn poll_read(
        mut self: Pin<&mut Self>,
        cx: &mut Context<'_>,
        buf: &mut ReadBuf<'_>,
    ) -> Poll<io::Result<()>> {
        // If we have a sleep in progress, poll it first
        if let Some(mut sleep) = self.sleep.take() {
            match sleep.as_mut().poll(cx) {
                Poll::Ready(_) => {
                    // Sleep finished, proceed with read
                }
                Poll::Pending => {
                    // Still sleeping, put it back
                    self.sleep = Some(sleep);
                    return Poll::Pending;
                }
            }
        }

        if self.pos >= self.data.len() {
            return Poll::Ready(Ok(()));
        }

        let remaining = &self.data[self.pos..];
        let to_copy = remaining.len().min(buf.remaining()).min(self.chunk_size);
        if to_copy == 0 {
            return Poll::Ready(Ok(()));
        }

        buf.put_slice(&remaining[..to_copy]);
        self.pos += to_copy;
        self.chunks_read.fetch_add(1, Ordering::SeqCst);

        // Start a new sleep for the next read
        self.sleep = Some(Box::pin(tokio::time::sleep(self.delay)));

        Poll::Ready(Ok(()))
    }
}

#[tokio::test]
async fn encoder_streams_without_buffering_all_input() -> Result<(), FlacError> {
    // Generate PCM data: 1 second of 16-bit stereo at 44.1kHz
    let sample_rate = 44_100;
    let channels = 2u8;
    let bits_per_sample = 16u8;
    let duration_secs = 1;
    let total_samples = sample_rate * channels as u32 * duration_secs;
    let bytes_per_sample = 2;
    let total_bytes = total_samples as usize * bytes_per_sample;

    // Generate sine wave PCM data
    let mut pcm_data = Vec::with_capacity(total_bytes);
    for i in 0..total_samples / channels as u32 {
        let t = i as f32 / sample_rate as f32;
        let sample = (t * 440.0 * 2.0 * std::f32::consts::PI).sin();
        let sample_i16 = (sample * 16384.0) as i16;
        let bytes = sample_i16.to_le_bytes();
        // Stereo: same for both channels
        pcm_data.extend_from_slice(&bytes);
        pcm_data.extend_from_slice(&bytes);
    }

    let format = PcmFormat {
        sample_rate,
        channels,
        bits_per_sample,
    };

    // Create a slow reader that delivers 8KB chunks with 10ms delay
    let chunk_size = 8 * 1024;
    let delay = Duration::from_millis(10);
    let slow_reader = SlowReader::new(pcm_data.clone(), chunk_size, delay);
    let chunks_read_counter = slow_reader.chunks_read();

    // Start encoding
    let mut encoder_stream =
        encode_flac_stream(slow_reader, format, EncoderOptions::default()).await?;

    // Try to read some FLAC data before all PCM data has been consumed
    let mut first_chunk = vec![0u8; 4096];
    let output_started = Arc::new(AtomicBool::new(false));
    let output_started_clone = output_started.clone();

    // Spawn a task to check when we get first output
    let read_handle = tokio::spawn(async move {
        match encoder_stream.read(&mut first_chunk).await {
            Ok(n) if n > 0 => {
                output_started_clone.store(true, Ordering::SeqCst);
                Ok((n, encoder_stream))
            }
            Ok(_) => Err(FlacError::Encode("No data read".into())),
            Err(e) => Err(FlacError::Io(e)),
        }
    });

    // Wait a bit to let the encoder start processing
    tokio::time::sleep(Duration::from_millis(200)).await;

    // Check that we've started getting output
    let (first_read, mut encoder_stream) =
        read_handle.await.map_err(|e| FlacError::TaskJoin {
            role: "read-test",
            details: e.to_string(),
        })??;

    assert!(first_read > 0, "Should have received some FLAC data");
    assert!(
        output_started.load(Ordering::SeqCst),
        "Output should have started"
    );

    let chunks_read_so_far = chunks_read_counter.load(Ordering::SeqCst);
    let total_chunks = (pcm_data.len() + chunk_size - 1) / chunk_size;

    // Verify streaming behavior: we should get output before reading everything
    // With delays of 50ms per chunk, if we're truly streaming, we should see output
    // before the slowreader has been fully consumed.
    // Note: this is a heuristic test. In practice, the encoder needs enough data
    // to fill at least one block before it can output anything.
    println!(
        "Streaming check: read {}/{} chunks when first output arrived",
        chunks_read_so_far, total_chunks
    );

    // More lenient check: just verify we got SOME output
    assert!(
        first_read > 0,
        "Should have received FLAC output (got {} bytes)",
        first_read
    );

    // Read the rest to completion
    let mut rest = Vec::new();
    encoder_stream.read_to_end(&mut rest).await?;
    encoder_stream.wait().await?;

    Ok(())
}

#[tokio::test]
async fn decoder_streams_without_buffering_all_input() -> Result<(), FlacError> {
    // Load a FLAC file
    let bytes = std::fs::read(fixture(
        "1abaa2c7fb4302e20ac570e79857b700.32bits-44.1Khz.flac",
    ))?;

    // Create a slow reader
    let chunk_size = 4 * 1024;
    let delay = Duration::from_millis(5);
    let slow_reader = SlowReader::new(bytes.clone(), chunk_size, delay);
    let chunks_read_counter = slow_reader.chunks_read();

    // Start decoding
    let mut decoder_stream = decode_flac_stream(slow_reader).await?;

    // Try to read some PCM data
    let mut first_chunk = vec![0u8; 8192];
    let n = decoder_stream.read(&mut first_chunk).await?;

    assert!(n > 0, "Should have received some PCM data");

    let chunks_read_so_far = chunks_read_counter.load(Ordering::SeqCst);
    let total_chunks = (bytes.len() + chunk_size - 1) / chunk_size;

    // Verify streaming behavior
    assert!(
        chunks_read_so_far < (total_chunks * 8 / 10),
        "Decoder should produce output before consuming all input. Read {}/{} chunks",
        chunks_read_so_far,
        total_chunks
    );

    // Read the rest
    let mut rest = Vec::new();
    decoder_stream.read_to_end(&mut rest).await?;
    decoder_stream.wait().await?;

    Ok(())
}