mlua-pulse 0.1.0

use crate::composition::PulseSong;
use crate::error::{PulseError, PulseResult};
use std::path::Path;
use tunes::engine::AudioEngine;
use tunes::track::{AudioEvent, Mixer};

const DEFAULT_NORMALIZED_SAMPLE_RATE: u32 = 44_100;
const DEFAULT_NORMALIZED_FLAC_BITS_PER_SAMPLE: u32 = 16;
pub(crate) const PCM_CHANNELS: usize = 2;

/// Offline audio normalization mode.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum NormalizeMode {
    /// Scale the final render so its sample peak reaches the target level.
    Peak,
    /// Scale the final render to a target RMS level, capped by a peak ceiling.
    Rms,
}

/// Options for final audio normalization during WAV/FLAC export.
#[derive(Debug, Clone, Copy, PartialEq)]
pub struct NormalizeOptions {
    mode: NormalizeMode,
    target_db: f32,
    peak_ceiling_db: f32,
}

impl NormalizeOptions {
    /// Creates peak normalization options.
    pub fn peak(target_db: f32) -> PulseResult<Self> {
        validate_normalize_db("target_db", target_db)?;
        Ok(Self {
            mode: NormalizeMode::Peak,
            target_db,
            peak_ceiling_db: target_db,
        })
    }

    /// Creates RMS normalization options with a peak ceiling.
    pub fn rms(target_db: f32, peak_ceiling_db: f32) -> PulseResult<Self> {
        validate_normalize_db("target_db", target_db)?;
        validate_normalize_db("true_peak_db", peak_ceiling_db)?;
        Ok(Self {
            mode: NormalizeMode::Rms,
            target_db,
            peak_ceiling_db,
        })
    }

    /// Returns the normalization mode.
    pub fn mode(self) -> NormalizeMode {
        self.mode
    }

    /// Returns the target level in dBFS.
    pub fn target_db(self) -> f32 {
        self.target_db
    }

    /// Returns the peak ceiling in dBFS.
    pub fn peak_ceiling_db(self) -> f32 {
        self.peak_ceiling_db
    }
}

/// Options that control Standard MIDI File playback loudness.
#[derive(Debug, Clone, Copy, PartialEq)]
pub struct MidiExportOptions {
    velocity_gain: f32,
    min_velocity: Option<f32>,
}

impl Default for MidiExportOptions {
    fn default() -> Self {
        Self {
            velocity_gain: 1.0,
            min_velocity: None,
        }
    }
}

/// Options that control offline audio export.
#[derive(Debug, Clone, Copy, Default, PartialEq)]
pub struct ExportOptions {
    gpu: bool,
    sample_rate: Option<u32>,
    normalize: Option<NormalizeOptions>,
    flac_bits_per_sample: Option<u32>,
    midi: MidiExportOptions,
}

impl ExportOptions {
    /// Creates default export options.
    pub fn new() -> Self {
        Self::default()
    }

    /// Enables or disables the `tunes` GPU export path.
    #[must_use]
    pub fn with_gpu(mut self, gpu: bool) -> Self {
        self.gpu = gpu;
        self
    }

    /// Sets an explicit output sample rate for WAV/FLAC exports.
    pub fn with_sample_rate(mut self, sample_rate: u32) -> PulseResult<Self> {
        validate_export_sample_rate(sample_rate)?;
        self.sample_rate = Some(sample_rate);
        Ok(self)
    }

    /// Enables peak normalization for WAV/FLAC exports.
    pub fn with_peak_normalization(mut self, target_db: f32) -> PulseResult<Self> {
        self.normalize = Some(NormalizeOptions::peak(target_db)?);
        Ok(self)
    }

    /// Enables RMS normalization for WAV/FLAC exports.
    pub fn with_rms_normalization(
        mut self,
        target_db: f32,
        peak_ceiling_db: f32,
    ) -> PulseResult<Self> {
        self.normalize = Some(NormalizeOptions::rms(target_db, peak_ceiling_db)?);
        Ok(self)
    }

    /// Sets an already parsed normalization option.
    #[must_use]
    pub fn with_normalization(mut self, normalize: NormalizeOptions) -> Self {
        self.normalize = Some(normalize);
        self
    }

    /// Sets the FLAC bit depth used by the normalized/manual FLAC encoder path.
    ///
    /// `tunes` handles the ordinary FLAC export path. When final loudness
    /// normalization or an explicit FLAC bit depth is requested, `mlua-pulse`
    /// renders the mixer into a stereo buffer and writes FLAC directly. This
    /// option selects that direct encoder bit depth.
    pub fn with_flac_bits_per_sample(mut self, bits_per_sample: u32) -> PulseResult<Self> {
        validate_flac_bits_per_sample(bits_per_sample)?;
        self.flac_bits_per_sample = Some(bits_per_sample);
        Ok(self)
    }

    /// Scales melodic MIDI note velocities during MIDI export.
    pub fn with_midi_velocity_gain(mut self, gain: f32) -> PulseResult<Self> {
        validate_midi_velocity_gain(gain)?;
        self.midi.velocity_gain = gain;
        Ok(self)
    }

    /// Sets a minimum non-zero MIDI note velocity during MIDI export.
    pub fn with_midi_min_velocity(mut self, min_velocity: f32) -> PulseResult<Self> {
        validate_midi_min_velocity(min_velocity)?;
        self.midi.min_velocity = Some(min_velocity);
        Ok(self)
    }

    /// Returns whether GPU export was requested.
    pub fn gpu(self) -> bool {
        self.gpu
    }

    /// Returns the requested output sample rate, if any.
    pub fn sample_rate(self) -> Option<u32> {
        self.sample_rate
    }

    /// Returns audio normalization options, if enabled.
    pub fn normalize(self) -> Option<NormalizeOptions> {
        self.normalize
    }

    /// Returns the requested normalized FLAC bit depth, or the default.
    pub fn flac_bits_per_sample(self) -> u32 {
        self.flac_bits_per_sample
            .unwrap_or(DEFAULT_NORMALIZED_FLAC_BITS_PER_SAMPLE)
    }

    /// Returns whether the FLAC bit depth was explicitly selected.
    pub fn has_explicit_flac_bits_per_sample(self) -> bool {
        self.flac_bits_per_sample.is_some()
    }

    /// Returns the MIDI velocity gain.
    pub fn midi_velocity_gain(self) -> f32 {
        self.midi.velocity_gain
    }

    /// Returns the minimum non-zero MIDI velocity, if enabled.
    pub fn midi_min_velocity(self) -> Option<f32> {
        self.midi.min_velocity
    }
}

fn validate_export_sample_rate(sample_rate: u32) -> PulseResult<()> {
    if sample_rate == 0 {
        return Err(PulseError::InvalidExportSampleRate { sample_rate });
    }

    Ok(())
}

fn validate_normalize_db(option: &str, value: f32) -> PulseResult<()> {
    if value.is_finite() && (-60.0..=0.0).contains(&value) {
        Ok(())
    } else {
        Err(PulseError::InvalidExportNormalizeOption {
            option: option.to_string(),
            value: value.to_string(),
        })
    }
}

fn validate_flac_bits_per_sample(bits_per_sample: u32) -> PulseResult<()> {
    if matches!(bits_per_sample, 16 | 24) {
        Ok(())
    } else {
        Err(PulseError::InvalidExportFlacBitsPerSample { bits_per_sample })
    }
}

fn validate_midi_velocity_gain(gain: f32) -> PulseResult<()> {
    if gain.is_finite() && gain > 0.0 && gain <= 16.0 {
        Ok(())
    } else {
        Err(PulseError::InvalidExportMidiOption {
            option: "midi_velocity_gain".to_string(),
            value: gain.to_string(),
        })
    }
}

fn validate_midi_min_velocity(min_velocity: f32) -> PulseResult<()> {
    if min_velocity.is_finite() && (0.0..=1.0).contains(&min_velocity) {
        Ok(())
    } else {
        Err(PulseError::InvalidExportMidiOption {
            option: "midi_min_velocity".to_string(),
            value: min_velocity.to_string(),
        })
    }
}

/// Converts a filesystem path to UTF-8 text for `tunes` APIs.
pub(crate) fn path_to_string(path: &Path) -> PulseResult<String> {
    path.to_str()
        .map(str::to_string)
        .ok_or_else(|| PulseError::InvalidExportPath {
            path: path.to_path_buf(),
        })
}

/// Creates the parent directory for an output path and returns UTF-8 path text.
pub(crate) fn prepare_output_path(path: &Path) -> PulseResult<String> {
    if let Some(parent) = path.parent() {
        if !parent.as_os_str().is_empty() {
            std::fs::create_dir_all(parent).map_err(|error| PulseError::ExportFailed {
                message: error.to_string(),
            })?;
        }
    }

    path_to_string(path)
}

fn audio_engine(options: ExportOptions) -> PulseResult<AudioEngine> {
    if options.gpu() {
        return gpu_audio_engine();
    }

    AudioEngine::new().map_err(|error| PulseError::ExportFailed {
        message: error.to_string(),
    })
}

#[cfg(feature = "gpu")]
fn gpu_audio_engine() -> PulseResult<AudioEngine> {
    AudioEngine::new_with_gpu().map_err(|error| PulseError::ExportFailed {
        message: error.to_string(),
    })
}

#[cfg(not(feature = "gpu"))]
fn gpu_audio_engine() -> PulseResult<AudioEngine> {
    Err(PulseError::GpuFeatureNotEnabled)
}

fn prepare_mixer_for_export(mixer: &mut Mixer, options: ExportOptions) -> PulseResult<()> {
    if options.gpu() {
        enable_mixer_gpu(mixer)?;
    }

    Ok(())
}

fn normalized_sample_rate(options: ExportOptions) -> u32 {
    options
        .sample_rate()
        .unwrap_or(DEFAULT_NORMALIZED_SAMPLE_RATE)
}

fn render_normalized_buffer(
    mixer: &mut Mixer,
    options: ExportOptions,
) -> PulseResult<(Vec<f32>, u32)> {
    let sample_rate = normalized_sample_rate(options);
    validate_export_sample_rate(sample_rate)?;
    prepare_mixer_for_export(mixer, options)?;
    let mut buffer = mixer.render_to_buffer(sample_rate as f32);
    if let Some(normalize) = options.normalize() {
        normalize_buffer(&mut buffer, normalize);
    }
    Ok((buffer, sample_rate))
}

fn normalize_buffer(buffer: &mut [f32], options: NormalizeOptions) {
    if buffer.is_empty() {
        return;
    }

    let current_peak = peak(buffer);
    if current_peak <= f32::EPSILON {
        return;
    }

    let gain = match options.mode() {
        NormalizeMode::Peak => db_to_linear(options.target_db()) / current_peak,
        NormalizeMode::Rms => {
            let current_rms = rms(buffer);
            if current_rms <= f32::EPSILON {
                return;
            }
            db_to_linear(options.target_db()) / current_rms
        }
    };

    if !gain.is_finite() {
        return;
    }

    for sample in &mut *buffer {
        *sample = (*sample * gain).clamp(-1.0, 1.0);
    }

    if matches!(options.mode(), NormalizeMode::Rms) {
        limit_buffer_peaks(buffer, db_to_linear(options.peak_ceiling_db()));
    }
}

fn peak(buffer: &[f32]) -> f32 {
    buffer
        .iter()
        .fold(0.0_f32, |current, sample| current.max(sample.abs()))
}

fn rms(buffer: &[f32]) -> f32 {
    let energy = buffer.iter().map(|sample| sample * sample).sum::<f32>() / buffer.len() as f32;
    energy.sqrt()
}

fn db_to_linear(db: f32) -> f32 {
    10.0_f32.powf(db / 20.0)
}

fn limit_buffer_peaks(buffer: &mut [f32], ceiling: f32) {
    if !ceiling.is_finite() || ceiling <= 0.0 {
        return;
    }

    let ceiling = ceiling.min(1.0);
    for sample in buffer {
        let magnitude = sample.abs();
        if magnitude > ceiling {
            *sample = sample.signum() * ceiling;
        }
    }
}

fn write_wav_buffer(path: &str, sample_rate: u32, buffer: &[f32]) -> PulseResult<()> {
    let spec = hound::WavSpec {
        channels: 2,
        sample_rate,
        bits_per_sample: 16,
        sample_format: hound::SampleFormat::Int,
    };
    let mut writer =
        hound::WavWriter::create(path, spec).map_err(|error| PulseError::ExportFailed {
            message: error.to_string(),
        })?;

    for &sample in buffer {
        let sample = (sample.clamp(-1.0, 1.0) * 32767.0).round() as i16;
        writer
            .write_sample(sample)
            .map_err(|error| PulseError::ExportFailed {
                message: error.to_string(),
            })?;
    }

    writer.finalize().map_err(|error| PulseError::ExportFailed {
        message: error.to_string(),
    })
}

/// Renders a mixer to interleaved signed 16-bit little-endian PCM samples.
pub(crate) fn render_mixer_pcm_i16(
    mixer: &mut Mixer,
    options: ExportOptions,
) -> PulseResult<(Vec<i16>, u32)> {
    let (buffer, sample_rate) = render_normalized_buffer(mixer, options)?;
    let samples = buffer
        .iter()
        .map(|sample| float_sample_to_i16(*sample))
        .collect();
    Ok((samples, sample_rate))
}

fn float_sample_to_i16(sample: f32) -> i16 {
    (sample.clamp(-1.0, 1.0) * 32767.0).round() as i16
}

fn write_flac_buffer(
    path: &str,
    sample_rate: u32,
    bits_per_sample: u32,
    buffer: &[f32],
) -> PulseResult<()> {
    use flacenc::component::BitRepr;
    use flacenc::error::Verify;
    use flacenc::source::MemSource;

    validate_flac_bits_per_sample(bits_per_sample)?;
    let scale = ((1_i64 << (bits_per_sample - 1)) - 1) as f32;

    let samples_i32 = buffer
        .iter()
        .map(|sample| (sample.clamp(-1.0, 1.0) * scale).round() as i32)
        .collect::<Vec<_>>();
    let config = flacenc::config::Encoder::default()
        .into_verified()
        .map_err(|(_, error)| PulseError::ExportFailed {
            message: format!("flac config error: {error:?}"),
        })?;
    let source = MemSource::from_samples(
        &samples_i32,
        2,
        bits_per_sample as usize,
        sample_rate as usize,
    );
    let flac_stream = flacenc::encode_with_fixed_block_size(&config, source, config.block_size)
        .map_err(|error| PulseError::ExportFailed {
            message: format!("flac encoding failed: {error:?}"),
        })?;
    let mut sink = flacenc::bitsink::ByteSink::new();
    flac_stream
        .write(&mut sink)
        .map_err(|error| PulseError::ExportFailed {
            message: format!("flac write failed: {error:?}"),
        })?;

    std::fs::write(path, sink.as_slice()).map_err(|error| PulseError::ExportFailed {
        message: error.to_string(),
    })
}

/// Applies MIDI-only loudness options to a mixer before Standard MIDI File export.
pub(crate) fn apply_midi_export_options(mixer: &mut Mixer, options: ExportOptions) {
    let gain = options.midi_velocity_gain();
    let min_velocity = options.midi_min_velocity();
    if (gain - 1.0).abs() < f32::EPSILON && min_velocity.is_none() {
        return;
    }

    for track in mixer.all_tracks_mut() {
        let track_volume = track.volume;
        for event in &mut track.events {
            match event {
                AudioEvent::Note(note) => {
                    let combined = (note.velocity * track_volume).clamp(0.0, 1.0);
                    note.velocity = lifted_velocity(combined, gain, min_velocity);
                }
                AudioEvent::Drum(drum) => {
                    let combined = (drum.velocity * track_volume).clamp(0.0, 1.0);
                    drum.velocity = lifted_velocity(combined, gain, min_velocity);
                }
                _ => {}
            }
        }
        track.volume = 1.0;
    }
}

fn lifted_velocity(value: f32, gain: f32, min_velocity: Option<f32>) -> f32 {
    if value <= 0.0 {
        return 0.0;
    }

    let boosted = (value * gain).clamp(0.0, 1.0);
    min_velocity
        .map(|minimum| boosted.max(minimum))
        .unwrap_or(boosted)
}

#[cfg(feature = "gpu")]
fn enable_mixer_gpu(mixer: &mut Mixer) -> PulseResult<()> {
    mixer.enable_gpu();
    Ok(())
}

#[cfg(not(feature = "gpu"))]
fn enable_mixer_gpu(_mixer: &mut Mixer) -> PulseResult<()> {
    Err(PulseError::GpuFeatureNotEnabled)
}

/// Exports an existing `tunes` mixer to WAV with explicit export options.
pub(crate) fn export_mixer_wav_with_options(
    mixer: &mut Mixer,
    path: impl AsRef<Path>,
    options: ExportOptions,
) -> PulseResult<()> {
    let path_text = prepare_output_path(path.as_ref())?;
    if options.normalize().is_some() {
        let (buffer, sample_rate) = render_normalized_buffer(mixer, options)?;
        return write_wav_buffer(&path_text, sample_rate, &buffer);
    }

    if let Some(sample_rate) = options.sample_rate() {
        prepare_mixer_for_export(mixer, options)?;
        return mixer.export_wav(&path_text, sample_rate).map_err(|error| {
            PulseError::ExportFailed {
                message: error.to_string(),
            }
        });
    }

    let engine = audio_engine(options)?;

    engine
        .export_wav(mixer, &path_text)
        .map_err(|error| PulseError::ExportFailed {
            message: error.to_string(),
        })
}

/// Exports an existing `tunes` mixer to FLAC with explicit export options.
pub(crate) fn export_mixer_flac_with_options(
    mixer: &mut Mixer,
    path: impl AsRef<Path>,
    options: ExportOptions,
) -> PulseResult<()> {
    let path_text = prepare_output_path(path.as_ref())?;
    if options.normalize().is_some() || options.has_explicit_flac_bits_per_sample() {
        let (buffer, sample_rate) = render_normalized_buffer(mixer, options)?;
        return write_flac_buffer(
            &path_text,
            sample_rate,
            options.flac_bits_per_sample(),
            &buffer,
        );
    }

    if let Some(sample_rate) = options.sample_rate() {
        prepare_mixer_for_export(mixer, options)?;
        return mixer.export_flac(&path_text, sample_rate).map_err(|error| {
            PulseError::ExportFailed {
                message: error.to_string(),
            }
        });
    }

    let engine = audio_engine(options)?;

    engine
        .export_flac(mixer, &path_text)
        .map_err(|error| PulseError::ExportFailed {
            message: error.to_string(),
        })
}

/// Exports a song to WAV through `tunes::engine::AudioEngine::export_wav`.
pub fn export_wav(song: &PulseSong, path: impl AsRef<Path>) -> PulseResult<()> {
    export_wav_with_options(song, path, ExportOptions::new())
}

/// Exports a song to WAV through `tunes::engine::AudioEngine::export_wav` with options.
pub fn export_wav_with_options(
    song: &PulseSong,
    path: impl AsRef<Path>,
    options: ExportOptions,
) -> PulseResult<()> {
    let mut mixer = song.to_mixer()?;
    export_mixer_wav_with_options(&mut mixer, path, options)
}