pmomusic/pmoflac/src/encoder.rs

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

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

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

/// Extracted metadata values for FLAC encoding.
///
/// This is a simple struct containing the extracted values from TrackMetadata,
/// used to pass metadata into the blocking encoder task.
#[derive(Debug, Clone, Default)]
struct ExtractedMetadata {
    title: Option<String>,
    artist: Option<String>,
    album: Option<String>,
    year: Option<u32>,
    genre: Option<String>,
    track_number: Option<u32>,
    cover_pk: Option<String>,
    cover_url: Option<String>,
    server_base_url: Option<String>,
}

/// Options for configuring FLAC encoding.
#[derive(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>,

    /// Metadata to embed in the FLAC file (Vorbis Comments).
    /// Default: None (no metadata)
    pub metadata: Option<Arc<RwLock<dyn pmometadata::TrackMetadata>>>,

    /// Base URL of the server for constructing cover URLs.
    /// Default: None
    pub server_base_url: Option<String>,
}

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

impl std::fmt::Debug for EncoderOptions {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        f.debug_struct("EncoderOptions")
            .field("compression_level", &self.compression_level)
            .field("verify", &self.verify)
            .field("total_samples", &self.total_samples)
            .field("block_size", &self.block_size)
            .field("metadata", &self.metadata.as_ref().map(|_| "Some(...)"))
            .finish()
    }
}

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

    // Extract metadata before spawn_blocking (since TrackMetadata has async methods)
    let extracted_metadata = if let Some(metadata_lock) = &options.metadata {
        let metadata = metadata_lock.read().await;
        let title = metadata.get_title().await.ok().flatten();
        let artist = metadata.get_artist().await.ok().flatten();
        let album = metadata.get_album().await.ok().flatten();
        let year = metadata.get_year().await.ok().flatten();
        let cover_pk = metadata.get_cover_pk().await.ok().flatten();
        let cover_url = metadata.get_cover_url().await.ok().flatten();

        // Try to extract genre and track_number from extra fields
        let extra = metadata.get_extra().await.ok().flatten();
        let genre = extra.as_ref().and_then(|e| e.get("genre").cloned());
        let track_number = extra
            .as_ref()
            .and_then(|e| e.get("track_number").and_then(|s| s.parse::<u32>().ok()));

        Some(ExtractedMetadata {
            title,
            artist,
            album,
            year,
            genre,
            track_number,
            cover_pk,
            cover_url,
            server_base_url: options.server_base_url.clone(),
        })
    } else {
        None
    };

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

/// RAII guard for FLAC metadata block
struct MetadataGuard {
    metadata: *mut libflac_sys::FLAC__StreamMetadata,
}

impl Drop for MetadataGuard {
    fn drop(&mut self) {
        unsafe {
            if !self.metadata.is_null() {
                libflac_sys::FLAC__metadata_object_delete(self.metadata);
            }
        }
    }
}

/// Sets up Vorbis Comment metadata for the FLAC encoder
unsafe fn setup_metadata(
    encoder: *mut libflac_sys::FLAC__StreamEncoder,
    metadata: &ExtractedMetadata,
) -> Result<MetadataGuard, FlacError> {
    use libflac_sys::*;

    // Create a Vorbis Comment block
    let meta = FLAC__metadata_object_new(FLAC__METADATA_TYPE_VORBIS_COMMENT);
    if meta.is_null() {
        return Err(FlacError::LibFlacInit(
            "Failed to create metadata block".into(),
        ));
    }

    let guard = MetadataGuard { metadata: meta };

    // Helper to append a Vorbis comment
    let append_comment = |field_name: &str, value: &str| -> Result<(), FlacError> {
        let c_field_name = CString::new(field_name).map_err(|_| {
            FlacError::LibFlacInit("Failed to create CString for field name".into())
        })?;
        let c_value = CString::new(value).map_err(|_| {
            FlacError::LibFlacInit("Failed to create CString for field value".into())
        })?;

        let mut entry: FLAC__StreamMetadata_VorbisComment_Entry = std::mem::zeroed();
        let success = FLAC__metadata_object_vorbiscomment_entry_from_name_value_pair(
            &mut entry as *mut _,
            c_field_name.as_ptr(),
            c_value.as_ptr(),
        );

        if success == 0 {
            return Err(FlacError::LibFlacInit(format!(
                "Failed to create metadata entry for {}",
                field_name
            )));
        }

        let append_success =
            FLAC__metadata_object_vorbiscomment_append_comment(meta, entry, 0 /* copy */);

        if append_success == 0 {
            return Err(FlacError::LibFlacInit(format!(
                "Failed to append metadata entry for {}",
                field_name
            )));
        }

        Ok(())
    };

    // Add all available metadata fields
    if let Some(title) = &metadata.title {
        append_comment("TITLE", title)?;
    }
    if let Some(artist) = &metadata.artist {
        append_comment("ARTIST", artist)?;
    }
    if let Some(album) = &metadata.album {
        append_comment("ALBUM", album)?;
    }
    if let Some(year) = metadata.year {
        append_comment("DATE", &year.to_string())?;
    }
    if let Some(genre) = &metadata.genre {
        append_comment("GENRE", genre)?;
    }
    if let Some(track_number) = metadata.track_number {
        append_comment("TRACKNUMBER", &track_number.to_string())?;
    }
    // Construct cover URL: use cover_pk with server_base_url if available, fallback to cover_url
    if let (Some(pk), Some(base_url)) = (&metadata.cover_pk, &metadata.server_base_url) {
        let cover_url = format!("{}/covers/image/{}", base_url, pk);
        append_comment("COVERART", &cover_url)?;
    } else if let Some(cover_url) = &metadata.cover_url {
        append_comment("COVERART", cover_url)?;
    }

    // Set the metadata on the encoder
    let mut metadata_array = [meta];
    let set_success = FLAC__stream_encoder_set_metadata(encoder, metadata_array.as_mut_ptr(), 1);

    if set_success == 0 {
        return Err(FlacError::LibFlacInit(
            "Failed to set metadata on encoder".into(),
        ));
    }

    Ok(guard)
}

fn run_encoder(
    format: PcmFormat,
    options: EncoderOptions,
    metadata: Option<ExtractedMetadata>,
    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 {
            tracing::debug!(
                "FLAC encoder: setting total_samples_estimate = {} before init",
                total
            );
            ensure(
                FLAC__stream_encoder_set_total_samples_estimate(encoder, total),
                "set_total_samples_estimate failed",
            )?;
        } else {
            tracing::warn!(
                "FLAC encoder: total_samples is None, STREAMINFO will have total_samples=0"
            );
        }
        if let Some(block_size) = options.block_size {
            ensure(
                FLAC__stream_encoder_set_blocksize(encoder, block_size),
                "set_blocksize failed",
            )?;
        }

        // Setup metadata if provided
        let _metadata_guard = if let Some(meta) = metadata {
            Some(setup_metadata(encoder, &meta)?)
        } else {
            None
        };

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