pmomusic/pmoaudio/src/nodes/file_source.rs

use crate::{
    nodes::{AudioError, MultiSubscriberNode, TypedAudioNode, DEFAULT_CHUNK_DURATION_MS},
    type_constraints::TypeRequirement,
    AudioChunk, AudioChunkData, AudioSegment, I24,
};
use pmoflac::{decode_audio_stream, AudioFileMetadata, StreamInfo};
use pmometadata::{MemoryTrackMetadata, TrackMetadata};
use std::{path::PathBuf, sync::Arc, time::Duration};
use tokio::{fs::File, io::AsyncReadExt, sync::mpsc};

/// FileSource - Lit un fichier audio et publie des `AudioSegment`
///
/// Cette source utilise `pmoflac` pour décoder le fichier (FLAC/MP3/OGG/WAV/AIFF)
/// puis transforme les échantillons PCM en `AudioSegment` stéréo avec le type approprié
/// (I16, I24, ou I32) selon la profondeur de bit du fichier source.
///
/// Le node émet trois types de syncmarkers :
/// - `TopZeroSync` au début du flux
/// - `TrackBoundary` avec les métadonnées du fichier
/// - `EndOfStream` à la fin du flux
pub struct FileSource {
    path: PathBuf,
    chunk_frames: usize,
    subscribers: MultiSubscriberNode,
}

impl FileSource {
    /// Crée une nouvelle source de fichier avec calcul automatique de la taille des chunks.
    ///
    /// La taille des chunks sera calculée automatiquement pour obtenir environ 50ms
    /// de latence par chunk, en fonction du sample rate du fichier.
    ///
    /// * `path` - chemin du fichier audio à lire
    pub fn new<P: Into<PathBuf>>(path: P) -> Self {
        Self::with_chunk_size(path, 0) // 0 = auto-calculer
    }

    /// Crée une nouvelle source de fichier avec une taille de chunk spécifique.
    ///
    /// * `path` - chemin du fichier audio à lire
    /// * `chunk_frames` - nombre d'échantillons par canal par chunk (0 = auto)
    pub fn with_chunk_size<P: Into<PathBuf>>(path: P, chunk_frames: usize) -> Self {
        Self {
            path: path.into(),
            chunk_frames,
            subscribers: MultiSubscriberNode::new(),
        }
    }

    /// Ajoute un abonné qui recevra les segments audio.
    pub fn add_subscriber(&mut self, tx: mpsc::Sender<Arc<AudioSegment>>) {
        self.subscribers.add_subscriber(tx);
    }

    /// Lance la lecture du fichier et diffuse les segments audio.
    pub async fn run(self) -> Result<(), AudioError> {
        // Ouvrir le fichier
        let file = File::open(&self.path).await.map_err(|e| {
            AudioError::ProcessingError(format!("Failed to open {:?}: {}", self.path, e))
        })?;

        // Décoder le flux audio
        let mut stream = decode_audio_stream(file)
            .await
            .map_err(|e| AudioError::ProcessingError(format!("Decode error: {}", e)))?;
        let stream_info = stream.info().clone();

        validate_stream(&stream_info)?;

        // Calculer la taille des chunks si non spécifiée (0 = auto)
        let chunk_frames = if self.chunk_frames == 0 {
            // Calculer pour obtenir DEFAULT_CHUNK_DURATION_MS millisecondes
            let frames =
                (stream_info.sample_rate as f64 * DEFAULT_CHUNK_DURATION_MS / 1000.0) as usize;
            // Arrondir à la puissance de 2 la plus proche pour optimiser les buffers
            frames.next_power_of_two().max(256)
        } else {
            self.chunk_frames.max(1)
        };

        // Émettre TopZeroSync
        let top_zero = AudioSegment::new_top_zero_sync();
        self.subscribers.push(top_zero).await?;

        // Extraire et émettre les métadonnées du fichier
        match AudioFileMetadata::from_file(&self.path) {
            Ok(file_metadata) => {
                let mut metadata = MemoryTrackMetadata::new();

                // Convertir AudioFileMetadata vers MemoryTrackMetadata
                if let Some(title) = file_metadata.title {
                    let _ = metadata.set_title(Some(title)).await;
                }
                if let Some(artist) = file_metadata.artist {
                    let _ = metadata.set_artist(Some(artist)).await;
                }
                if let Some(album) = file_metadata.album {
                    let _ = metadata.set_album(Some(album)).await;
                }
                if let Some(year) = file_metadata.year {
                    let _ = metadata.set_year(Some(year)).await;
                }
                if let Some(duration_secs) = file_metadata.duration_secs {
                    let _ = metadata
                        .set_duration(Some(Duration::from_secs(duration_secs)))
                        .await;
                }

                // Émettre TrackBoundary
                let track_boundary = AudioSegment::new_track_boundary(0, 0.0, Arc::new(tokio::sync::RwLock::new(metadata)));
                self.subscribers.push(track_boundary).await?;
            }
            Err(e) => {
                eprintln!(
                    "Warning: Failed to extract metadata from {:?}: {}",
                    self.path, e
                );
                // Continuer sans métadonnées
            }
        }

        // Préparer la lecture des chunks audio
        let frame_bytes = stream_info.bytes_per_sample() * stream_info.channels as usize;
        let chunk_byte_len = chunk_frames * frame_bytes;
        let mut pending = Vec::new();
        let mut read_buf = vec![0u8; frame_bytes * 512.max(chunk_frames)];
        let mut chunk_index = 0u64;
        let mut total_frames = 0u64;

        // Lire et émettre les chunks audio
        loop {
            // Remplir le buffer
            if pending.len() < chunk_byte_len {
                let read = stream.read(&mut read_buf).await.map_err(|e| {
                    AudioError::ProcessingError(format!("I/O error while decoding: {}", e))
                })?;
                if read == 0 {
                    break;
                }
                pending.extend_from_slice(&read_buf[..read]);
            }

            if pending.is_empty() {
                break;
            }

            // Extraire un chunk
            let frames_in_pending = pending.len() / frame_bytes;
            let frames_to_emit = frames_in_pending.min(chunk_frames);
            let take_bytes = frames_to_emit * frame_bytes;
            let chunk_bytes = pending.drain(..take_bytes).collect::<Vec<u8>>();

            // Calculer le timestamp
            let timestamp_sec = total_frames as f64 / stream_info.sample_rate as f64;

            // Créer le segment audio
            let segment = bytes_to_segment(
                &chunk_bytes,
                &stream_info,
                frames_to_emit,
                chunk_index,
                timestamp_sec,
            )?;
            self.subscribers.push(segment).await?;

            chunk_index += 1;
            total_frames += frames_to_emit as u64;
        }

        // Traiter le reste éventuel (moins qu'un chunk complet)
        if !pending.is_empty() {
            let frames = pending.len() / frame_bytes;
            if frames > 0 {
                let timestamp_sec = total_frames as f64 / stream_info.sample_rate as f64;
                let segment =
                    bytes_to_segment(&pending, &stream_info, frames, chunk_index, timestamp_sec)?;
                self.subscribers.push(segment).await?;
                total_frames += frames as u64;
                chunk_index += 1;
            }
        }

        // Émettre EndOfStream
        let final_timestamp = total_frames as f64 / stream_info.sample_rate as f64;
        let eos = AudioSegment::new_end_of_stream(chunk_index, final_timestamp);
        self.subscribers.push(eos).await?;

        // Attendre la fin du décodage
        stream
            .wait()
            .await
            .map_err(|e| AudioError::ProcessingError(format!("Decode task failed: {}", e)))?;

        Ok(())
    }
}

fn validate_stream(info: &StreamInfo) -> Result<(), AudioError> {
    if !(1..=2).contains(&info.channels) {
        return Err(AudioError::ProcessingError(format!(
            "Unsupported channel count: {}",
            info.channels
        )));
    }
    match info.bits_per_sample {
        8 | 16 | 24 | 32 => Ok(()),
        other => Err(AudioError::ProcessingError(format!(
            "Unsupported bit depth: {}",
            other
        ))),
    }
}

/// Convertit des bytes PCM en AudioSegment avec le type approprié
fn bytes_to_segment(
    chunk_bytes: &[u8],
    info: &StreamInfo,
    frames: usize,
    order: u64,
    timestamp_sec: f64,
) -> Result<Arc<AudioSegment>, AudioError> {
    let bytes_per_sample = info.bytes_per_sample();
    let channels = info.channels as usize;
    let frame_bytes = bytes_per_sample * channels;

    // Créer le chunk du bon type selon la profondeur de bit
    let chunk = match info.bits_per_sample {
        16 => {
            // Type I16
            let mut stereo = Vec::with_capacity(frames);
            for frame_idx in 0..frames {
                let base = frame_idx * frame_bytes;
                let l = i16::from_le_bytes(
                    chunk_bytes[base..base + bytes_per_sample]
                        .try_into()
                        .unwrap(),
                );
                let r = if channels == 1 {
                    l
                } else {
                    i16::from_le_bytes(
                        chunk_bytes[base + bytes_per_sample..base + 2 * bytes_per_sample]
                            .try_into()
                            .unwrap(),
                    )
                };
                stereo.push([l, r]);
            }
            let chunk_data = AudioChunkData::new(stereo, info.sample_rate, 0.0);
            AudioChunk::I16(chunk_data)
        }
        24 => {
            // Type I24
            let mut stereo = Vec::with_capacity(frames);
            for frame_idx in 0..frames {
                let base = frame_idx * frame_bytes;
                let l_i32 = {
                    let mut buf = [0u8; 4];
                    buf[..3].copy_from_slice(&chunk_bytes[base..base + 3]);
                    // Sign extend
                    if chunk_bytes[base + 2] & 0x80 != 0 {
                        buf[3] = 0xFF;
                    }
                    i32::from_le_bytes(buf)
                };
                let l = I24::new(l_i32).ok_or_else(|| {
                    AudioError::ProcessingError(format!("Invalid I24 value: {}", l_i32))
                })?;

                let r = if channels == 1 {
                    l
                } else {
                    let r_i32 = {
                        let mut buf = [0u8; 4];
                        buf[..3].copy_from_slice(
                            &chunk_bytes[base + bytes_per_sample..base + bytes_per_sample + 3],
                        );
                        // Sign extend
                        if chunk_bytes[base + bytes_per_sample + 2] & 0x80 != 0 {
                            buf[3] = 0xFF;
                        }
                        i32::from_le_bytes(buf)
                    };
                    I24::new(r_i32).ok_or_else(|| {
                        AudioError::ProcessingError(format!("Invalid I24 value: {}", r_i32))
                    })?
                };
                stereo.push([l, r]);
            }
            let chunk_data = AudioChunkData::new(stereo, info.sample_rate, 0.0);
            AudioChunk::I24(chunk_data)
        }
        32 => {
            // Type I32
            let mut stereo = Vec::with_capacity(frames);
            for frame_idx in 0..frames {
                let base = frame_idx * frame_bytes;
                let l = i32::from_le_bytes(
                    chunk_bytes[base..base + bytes_per_sample]
                        .try_into()
                        .unwrap(),
                );
                let r = if channels == 1 {
                    l
                } else {
                    i32::from_le_bytes(
                        chunk_bytes[base + bytes_per_sample..base + 2 * bytes_per_sample]
                            .try_into()
                            .unwrap(),
                    )
                };
                stereo.push([l, r]);
            }
            let chunk_data = AudioChunkData::new(stereo, info.sample_rate, 0.0);
            AudioChunk::I32(chunk_data)
        }
        other => {
            return Err(AudioError::ProcessingError(format!(
                "Unsupported bit depth: {}",
                other
            )))
        }
    };

    // Créer le segment audio
    Ok(Arc::new(AudioSegment {
        order,
        timestamp_sec,
        segment: crate::_AudioSegment::Chunk(Arc::new(chunk)),
    }))
}

impl TypedAudioNode for FileSource {
    fn input_type(&self) -> Option<TypeRequirement> {
        // FileSource est une source, elle ne consomme pas d'audio
        None
    }

    fn output_type(&self) -> Option<TypeRequirement> {
        // FileSource peut produire n'importe quel type entier (I16, I24, I32)
        // selon la profondeur de bit du fichier source
        Some(TypeRequirement::any_integer())
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use pmoflac::{encode_flac_stream, EncoderOptions, PcmFormat};
    use std::io::Cursor;
    use tokio::io::AsyncWriteExt;
    use tokio::sync::mpsc;

    #[tokio::test]
    async fn test_file_source_decodes_flac() {
        let temp_dir = tempfile::tempdir().unwrap();
        let flac_path = temp_dir.path().join("test.flac");

        let sample_rate = 48_000;
        let frames = 256;
        let mut pcm = Vec::with_capacity(frames * 4);
        for i in 0..frames {
            let sample = ((i % 32) as f32 / 31.0 * 2.0 - 1.0) * 0.5; // simple ramp
            let sample_i16 = (sample * 32767.0) as i16;
            pcm.extend_from_slice(&sample_i16.to_le_bytes());
            pcm.extend_from_slice(&sample_i16.to_le_bytes());
        }

        let format = PcmFormat {
            sample_rate,
            channels: 2,
            bits_per_sample: 16,
        };

        let mut flac_stream =
            encode_flac_stream(Cursor::new(pcm.clone()), format, EncoderOptions::default())
                .await
                .unwrap();

        let mut file = File::create(&flac_path).await.expect("create flac file");
        tokio::io::copy(&mut flac_stream, &mut file)
            .await
            .expect("write flac");
        file.flush().await.expect("flush file");
        flac_stream.wait().await.unwrap();

        let mut source = FileSource::with_chunk_size(&flac_path, 64);
        let (tx, mut rx) = mpsc::channel(16);
        source.add_subscriber(tx);

        tokio::spawn(async move {
            source.run().await.unwrap();
        });

        let mut received_frames = 0usize;
        let mut received_syncmarkers = 0usize;
        let mut seen_top_zero = false;
        let mut seen_eos = false;

        while let Some(segment) = rx.recv().await {
            if segment.is_audio_chunk() {
                if let Some(chunk) = segment.as_chunk() {
                    received_frames += chunk.len();
                    assert_eq!(chunk.sample_rate(), sample_rate);
                }
            } else {
                received_syncmarkers += 1;
                if let Some(marker) = segment.as_sync_marker() {
                    match **marker {
                        crate::SyncMarker::TopZeroSync => seen_top_zero = true,
                        crate::SyncMarker::EndOfStream => seen_eos = true,
                        crate::SyncMarker::TrackBoundary { .. } => {}
                        _ => {}
                    }
                }
            }
        }

        // Vérifier que tous les frames ont été reçus
        assert_eq!(received_frames, frames);
        // Vérifier qu'on a bien reçu des syncmarkers
        assert!(received_syncmarkers >= 2); // Au moins TopZeroSync et EndOfStream
        assert!(seen_top_zero, "Should have received TopZeroSync");
        assert!(seen_eos, "Should have received EndOfStream");
    }
}