pmomusic/pmoflac/src/encoder.rs

use std::{
    ffi::c_void,
    io,
    pin::Pin,
    task::{Context, Poll},
};

use tokio::{
    io::{AsyncRead, AsyncReadExt, AsyncWriteExt},
    sync::{mpsc, oneshot},
};

use crate::{
    error::FlacError,
    pcm::{PcmChunk, PcmFormat},
    stream::ManagedAsyncReader,
    util::le_bytes_to_interleaved_i32,
};

/// Channel capacity for async message passing between tasks.
const CHANNEL_CAPACITY: usize = 8;

/// Number of PCM frames to process per chunk (4096 frames = ~93ms at 44.1kHz).
const PCM_FRAMES_PER_CHUNK: usize = 4096;

/// An async stream that encodes PCM audio into FLAC format.
///
/// This struct implements `AsyncRead`, allowing you to read encoded FLAC data
/// as it becomes available. The encoding happens in a background task.
///
/// # Example
///
/// ```no_run
/// use pmoflac::{encode_flac_stream, EncoderOptions, PcmFormat};
/// use tokio::io::AsyncReadExt;
///
/// # #[tokio::main]
/// # async fn main() -> Result<(), Box<dyn std::error::Error>> {
/// let pcm_data: &[u8] = &[/* 16-bit stereo PCM */];
/// let format = PcmFormat {
///     sample_rate: 44_100,
///     channels: 2,
///     bits_per_sample: 16,
/// };
///
/// let mut stream = encode_flac_stream(pcm_data, format, EncoderOptions::default()).await?;
/// let mut flac_output = Vec::new();
/// stream.read_to_end(&mut flac_output).await?;
/// stream.wait().await?;
/// # Ok(())
/// # }
/// ```
pub struct FlacEncodedStream {
    format: PcmFormat,
    reader: ManagedAsyncReader<FlacError>,
}

impl FlacEncodedStream {
    /// Returns the PCM format used for encoding.
    pub fn format(&self) -> PcmFormat {
        self.format
    }

    /// Consumes the stream and returns its components.
    pub fn into_parts(self) -> (PcmFormat, ManagedAsyncReader<FlacError>) {
        (self.format, self.reader)
    }

    /// Waits for the background encoding task to complete.
    ///
    /// This should be called after reading all data to ensure proper cleanup
    /// and to catch any errors that occurred during encoding.
    pub async fn wait(self) -> Result<(), FlacError> {
        self.reader.wait().await
    }
}

impl tokio::io::AsyncRead for FlacEncodedStream {
    fn poll_read(
        mut self: Pin<&mut Self>,
        cx: &mut Context<'_>,
        buf: &mut tokio::io::ReadBuf<'_>,
    ) -> Poll<io::Result<()>> {
        Pin::new(&mut self.reader).poll_read(cx, buf)
    }
}

/// Options for configuring FLAC encoding.
#[derive(Debug, Clone)]
pub struct EncoderOptions {
    /// Compression level (0-12). Higher means better compression but slower.
    /// Default: 5 (balanced)
    pub compression_level: u32,

    /// Whether to verify the encoding by decoding in parallel.
    /// Default: false (disabled for performance)
    pub verify: bool,

    /// Total number of samples (optional). If known, improves seeking in output.
    pub total_samples: Option<u64>,

    /// Block size in samples (optional). If None, libFLAC chooses automatically.
    /// Typical values: 1152, 2304, 4096.
    pub block_size: Option<u32>,
}

impl Default for EncoderOptions {
    fn default() -> Self {
        Self {
            compression_level: 5,
            verify: false,
            total_samples: None,
            block_size: None,
        }
    }
}

/// Encodes PCM audio data into a FLAC stream.
///
/// This function spawns background tasks to perform the encoding asynchronously.
/// The returned `FlacEncodedStream` implements `AsyncRead` for streaming the FLAC output.
///
/// # Threading Model
///
/// - A Tokio task reads PCM chunks and converts them to i32 samples
/// - A blocking task (via `spawn_blocking`) runs the libFLAC encoder
/// - The encoder's write callback sends encoded data via a channel
/// - Another Tokio task writes FLAC data to an internal duplex stream
///
/// This architecture ensures true streaming: FLAC frames are produced as soon as
/// enough PCM data is available, without waiting for the entire input.
///
/// # Arguments
///
/// * `reader` - Any async reader containing PCM audio in little-endian interleaved format
/// * `format` - Describes the PCM format (sample rate, channels, bit depth)
/// * `options` - Encoding options (compression level, verify, etc.)
///
/// # PCM Input Format
///
/// The PCM data must be:
/// - **Little-endian** byte order
/// - **Interleaved** channels (L, R, L, R for stereo)
/// - **Signed integers** with bit depth matching `format.bits_per_sample`
///
/// # Returns
///
/// A `FlacEncodedStream` that can be read to obtain FLAC-encoded data.
///
/// # Errors
///
/// Returns an error if:
/// - The PCM format is invalid (e.g., unsupported bit depth)
/// - The input stream has incomplete sample data
/// - libFLAC initialization fails
/// - An I/O error occurs
///
/// # Example
///
/// ```no_run
/// use pmoflac::{encode_flac_stream, EncoderOptions, PcmFormat};
/// use tokio::io::AsyncReadExt;
///
/// # #[tokio::main]
/// # async fn main() -> Result<(), Box<dyn std::error::Error>> {
/// // Generate 1 second of silence at 44.1kHz stereo 16-bit
/// let sample_rate = 44_100u32;
/// let channels = 2u8;
/// let pcm_len = sample_rate as usize * channels as usize * 2;
///
/// let format = PcmFormat {
///     sample_rate,
///     channels,
///     bits_per_sample: 16,
/// };
///
/// let options = EncoderOptions {
///     compression_level: 8,
///     total_samples: Some(sample_rate as u64),
///     ..Default::default()
/// };
///
/// let mut stream = encode_flac_stream(
///     tokio::io::repeat(0).take(pcm_len as u64),
///     format,
///     options,
/// ).await?;
/// let mut flac_data = Vec::new();
/// stream.read_to_end(&mut flac_data).await?;
/// stream.wait().await?;
///
/// println!("Encoded {} bytes of PCM to {} bytes of FLAC",
///     pcm_len, flac_data.len());
/// # Ok(())
/// # }
/// ```
pub async fn encode_flac_stream<R>(
    reader: R,
    format: PcmFormat,
    options: EncoderOptions,
) -> Result<FlacEncodedStream, FlacError>
where
    R: AsyncRead + Unpin + Send + 'static,
{
    format
        .validate()
        .map_err(|msg| FlacError::Unsupported(format!("invalid PCM format: {msg}")))?;

    if options.compression_level > 12 {
        return Err(FlacError::Unsupported(
            "compression level must be in 0..=12".into(),
        ));
    }

    let (pcm_tx, pcm_rx) = mpsc::channel::<Result<PcmChunk, FlacError>>(CHANNEL_CAPACITY);
    let format_for_reader = format;
    tokio::spawn(async move {
        let _ = feed_pcm_chunks(reader, format_for_reader, pcm_tx).await;
    });

    let (flac_reader, mut flac_writer) = tokio::io::duplex(256 * 1024);
    let (flac_tx, mut flac_rx) = mpsc::channel::<Result<Vec<u8>, FlacError>>(CHANNEL_CAPACITY);
    let (init_tx, init_rx) = oneshot::channel::<Result<(), FlacError>>();

    let format_for_encoder = format;
    let options_for_encoder = options;
    let blocking_handle = tokio::task::spawn_blocking(move || {
        run_encoder(
            format_for_encoder,
            options_for_encoder,
            pcm_rx,
            flac_tx,
            init_tx,
        )
    });

    let writer_handle = tokio::spawn(async move {
        while let Some(chunk) = flac_rx.recv().await {
            let bytes = chunk?;
            if bytes.is_empty() {
                continue;
            }
            flac_writer.write_all(&bytes).await?;
        }
        flac_writer.shutdown().await?;
        match blocking_handle.await {
            Ok(res) => res,
            Err(err) => Err(FlacError::TaskJoin {
                role: "flac-encode",
                details: err.to_string(),
            }),
        }
    });

    init_rx.await.map_err(|_| FlacError::ChannelClosed)??;

    let reader = ManagedAsyncReader::new("flac-encode-writer", flac_reader, writer_handle);
    Ok(FlacEncodedStream { format, reader })
}

async fn feed_pcm_chunks<R>(
    reader: R,
    format: PcmFormat,
    tx: mpsc::Sender<Result<PcmChunk, FlacError>>,
) -> Result<(), FlacError>
where
    R: AsyncRead + Unpin,
{
    let bytes_per_frame = format.bytes_per_sample() * format.channels as usize;
    let chunk_bytes = PCM_FRAMES_PER_CHUNK * bytes_per_frame;
    let mut pending = Vec::with_capacity(chunk_bytes * 2);
    let mut reader = tokio::io::BufReader::new(reader);

    loop {
        while pending.len() >= chunk_bytes {
            let samples =
                le_bytes_to_interleaved_i32(&pending[..chunk_bytes], format.bits_per_sample)
                    .map_err(|msg| FlacError::Encode(msg))?;
            pending.drain(..chunk_bytes);
            let frames = (samples.len() / format.channels as usize) as u32;
            let chunk = PcmChunk::new(samples, frames, format.channels);
            if tx.send(Ok(chunk)).await.is_err() {
                return Ok(());
            }
        }

        let read = reader.read_buf(&mut pending).await?;
        if read == 0 {
            break;
        }
    }

    if !pending.is_empty() {
        if pending.len() % bytes_per_frame != 0 {
            let msg = "PCM stream ended with a partial frame (incomplete sample data)".to_string();
            let _ = tx.send(Err(FlacError::Encode(msg.clone()))).await;
            return Err(FlacError::Encode(msg));
        }
        let samples = le_bytes_to_interleaved_i32(&pending, format.bits_per_sample)
            .map_err(|msg| FlacError::Encode(msg))?;
        let frames = (samples.len() / format.channels as usize) as u32;
        let chunk = PcmChunk::new(samples, frames, format.channels);
        let _ = tx.send(Ok(chunk)).await;
    }

    Ok(())
}

fn run_encoder(
    format: PcmFormat,
    options: EncoderOptions,
    mut rx: mpsc::Receiver<Result<PcmChunk, FlacError>>,
    tx: mpsc::Sender<Result<Vec<u8>, FlacError>>,
    init_tx: oneshot::Sender<Result<(), FlacError>>,
) -> Result<(), FlacError> {
    use libflac_sys::*;

    unsafe {
        let encoder = FLAC__stream_encoder_new();
        if encoder.is_null() {
            let _ = init_tx.send(Err(FlacError::LibFlacInit(
                "FLAC__stream_encoder_new returned null".into(),
            )));
            return Err(FlacError::LibFlacInit(
                "FLAC__stream_encoder_new returned null".into(),
            ));
        }

        let _encoder_guard = EncoderHandle { ptr: encoder };

        let mut state = EncoderClientState::new(tx);

        let ensure = |ok: FLAC__bool, msg: &str| {
            if ok == 0 {
                Err(FlacError::LibFlacInit(msg.into()))
            } else {
                Ok(())
            }
        };

        ensure(
            FLAC__stream_encoder_set_channels(encoder, format.channels as u32),
            "set_channels failed",
        )?;
        ensure(
            FLAC__stream_encoder_set_bits_per_sample(encoder, format.bits_per_sample as u32),
            "set_bits_per_sample failed",
        )?;
        ensure(
            FLAC__stream_encoder_set_sample_rate(encoder, format.sample_rate),
            "set_sample_rate failed",
        )?;
        ensure(
            FLAC__stream_encoder_set_compression_level(encoder, options.compression_level),
            "set_compression_level failed",
        )?;
        ensure(
            FLAC__stream_encoder_set_streamable_subset(encoder, 1),
            "set_streamable_subset failed",
        )?;
        ensure(
            FLAC__stream_encoder_set_verify(encoder, if options.verify { 1 } else { 0 }),
            "set_verify failed",
        )?;
        if let Some(total) = options.total_samples {
            ensure(
                FLAC__stream_encoder_set_total_samples_estimate(encoder, total),
                "set_total_samples_estimate failed",
            )?;
        }
        if let Some(block_size) = options.block_size {
            ensure(
                FLAC__stream_encoder_set_blocksize(encoder, block_size),
                "set_blocksize failed",
            )?;
        }

        let init_status = FLAC__stream_encoder_init_stream(
            encoder,
            Some(write_callback),
            None,
            None,
            None,
            &mut state as *mut EncoderClientState as *mut c_void,
        );
        if init_status != libflac_sys::FLAC__STREAM_ENCODER_INIT_STATUS_OK {
            let msg = format!("init_stream failed: status {init_status}");
            let _ = init_tx.send(Err(FlacError::LibFlacInit(msg.clone())));
            return Err(FlacError::LibFlacInit(msg));
        }

        let _ = init_tx.send(Ok(()));

        while let Some(chunk_result) = rx.blocking_recv() {
            let chunk = match chunk_result {
                Ok(chunk) => chunk,
                Err(err) => {
                    FLAC__stream_encoder_finish(encoder);
                    return Err(err);
                }
            };
            if chunk.frames == 0 {
                continue;
            }

            let success = FLAC__stream_encoder_process_interleaved(
                encoder,
                chunk.data.as_ptr(),
                chunk.frames,
            );
            if success == 0 {
                if let Some(err) = state.error.take() {
                    FLAC__stream_encoder_finish(encoder);
                    return Err(err);
                }
                FLAC__stream_encoder_finish(encoder);
                return Err(FlacError::Encode("libFLAC reported encode failure".into()));
            }
        }

        let finish_ok = FLAC__stream_encoder_finish(encoder);
        if finish_ok == 0 {
            if let Some(err) = state.error.take() {
                return Err(err);
            }
            return Err(FlacError::Encode(
                "libFLAC failed to finalize stream".into(),
            ));
        }

        if let Some(err) = state.error.take() {
            return Err(err);
        }
    }

    Ok(())
}

struct EncoderClientState {
    tx: mpsc::Sender<Result<Vec<u8>, FlacError>>,
    error: Option<FlacError>,
}

impl EncoderClientState {
    fn new(tx: mpsc::Sender<Result<Vec<u8>, FlacError>>) -> Self {
        Self { tx, error: None }
    }
}

struct EncoderHandle {
    ptr: *mut libflac_sys::FLAC__StreamEncoder,
}

impl Drop for EncoderHandle {
    fn drop(&mut self) {
        unsafe {
            if !self.ptr.is_null() {
                libflac_sys::FLAC__stream_encoder_delete(self.ptr);
            }
        }
    }
}

unsafe extern "C" fn write_callback(
    _encoder: *const libflac_sys::FLAC__StreamEncoder,
    buffer: *const libflac_sys::FLAC__byte,
    bytes: usize,
    _samples: u32,
    _current_frame: u32,
    client_data: *mut c_void,
) -> libflac_sys::FLAC__StreamEncoderWriteStatus {
    let state = &mut *(client_data as *mut EncoderClientState);
    let slice = std::slice::from_raw_parts(buffer, bytes);
    match state.tx.blocking_send(Ok(slice.to_vec())) {
        Ok(_) => libflac_sys::FLAC__STREAM_ENCODER_WRITE_STATUS_OK,
        Err(_) => {
            state.error = Some(FlacError::LibFlacWrite(
                "failed to send encoded data (receiver dropped)".into(),
            ));
            libflac_sys::FLAC__STREAM_ENCODER_WRITE_STATUS_FATAL_ERROR
        }
    }
}