//! Audio conversion pipeline: channel count and sample rate conversion for export.

use rubato::{Resampler, SincFixedIn, SincInterpolationParameters, SincInterpolationType, WindowFunction};

use super::ExportChannels;
use crate::error::CoreError;
use tracing::instrument;

/// Audio data after conversion, ready for encoding.
pub struct ConvertedAudio {
    pub samples: Vec<f32>,
    pub sample_rate: u32,
    pub channels: u16,
}

/// Convert channel count: mono mixdown, stereo upmix, or passthrough.
#[instrument(skip_all)]
pub fn convert_channels(
    samples: &[f32],
    src_channels: u16,
    target: &ExportChannels,
) -> (Vec<f32>, u16) {
    if src_channels == 0 || samples.is_empty() {
        return (Vec::new(), src_channels.max(1));
    }
    match target {
        ExportChannels::Original => (samples.to_vec(), src_channels),
        ExportChannels::Mono => {
            if src_channels == 1 {
                return (samples.to_vec(), 1);
            }
            let ch = src_channels as usize;
            let num_frames = samples.len() / ch;
            let mut mono = Vec::with_capacity(num_frames);
            for frame in 0..num_frames {
                let mut sum = 0.0f32;
                for c in 0..ch {
                    sum += samples[frame * ch + c];
                }
                mono.push(sum / ch as f32);
            }
            (mono, 1)
        }
        ExportChannels::Stereo => {
            if src_channels == 2 {
                return (samples.to_vec(), 2);
            }
            if src_channels == 1 {
                // Mono -> stereo: duplicate each sample
                let mut stereo = Vec::with_capacity(samples.len() * 2);
                for &s in samples {
                    stereo.push(s);
                    stereo.push(s);
                }
                return (stereo, 2);
            }
            // Multi-channel -> stereo: take first two channels
            let ch = src_channels as usize;
            let num_frames = samples.len() / ch;
            let mut stereo = Vec::with_capacity(num_frames * 2);
            for frame in 0..num_frames {
                let base = frame * ch;
                stereo.push(samples[base]);
                stereo.push(samples.get(base + 1).copied().unwrap_or(0.0));
            }
            (stereo, 2)
        }
    }
}

/// Resample interleaved audio from src_rate to dst_rate using rubato.
/// Returns samples unchanged if rates match.
#[instrument(skip_all)]
pub fn resample(
    samples: &[f32],
    channels: u16,
    src_rate: u32,
    dst_rate: u32,
) -> Result<Vec<f32>, CoreError> {
    if src_rate == dst_rate {
        return Ok(samples.to_vec());
    }
    if channels == 0 || src_rate == 0 || dst_rate == 0 {
        return Err(CoreError::Export(format!(
            "invalid resample params: channels={channels}, src_rate={src_rate}, dst_rate={dst_rate}"
        )));
    }

    let ch = channels as usize;
    let num_frames = samples.len() / ch;

    // De-interleave into per-channel vectors
    let mut channel_bufs: Vec<Vec<f32>> = vec![Vec::with_capacity(num_frames); ch];
    for frame in 0..num_frames {
        for c in 0..ch {
            channel_bufs[c].push(samples[frame * ch + c]);
        }
    }

    let params = SincInterpolationParameters {
        sinc_len: 256,
        f_cutoff: 0.95,
        interpolation: SincInterpolationType::Linear,
        oversampling_factor: 256,
        window: WindowFunction::BlackmanHarris2,
    };

    let ratio = dst_rate as f64 / src_rate as f64;
    let chunk_size = 1024;

    let mut resampler = SincFixedIn::<f32>::new(ratio, 2.0, params, chunk_size, ch)
        .map_err(|e| CoreError::Export(format!("resampler init: {e}")))?;

    let mut output_channels: Vec<Vec<f32>> = vec![Vec::new(); ch];
    let mut pos = 0;

    while pos < num_frames {
        let end = (pos + chunk_size).min(num_frames);
        let actual_len = end - pos;

        let input_chunk: Vec<Vec<f32>> = channel_bufs
            .iter()
            .map(|buf| {
                let mut chunk = buf[pos..end].to_vec();
                // Pad last chunk to chunk_size if needed
                chunk.resize(chunk_size, 0.0);
                chunk
            })
            .collect();

        let output_chunk = resampler
            .process(&input_chunk, None)
            .map_err(|e| CoreError::Export(format!("resample: {e}")))?;

        // For the last chunk, we may have processed padding — calculate expected output
        let expected_output = if actual_len < chunk_size {
            (actual_len as f64 * ratio).ceil() as usize
        } else {
            output_chunk[0].len()
        };

        for (c, chunk) in output_chunk.into_iter().enumerate() {
            let chunk_len = chunk.len();
            let take = expected_output.min(chunk_len);
            output_channels[c].extend_from_slice(&chunk[..take]);
        }

        pos += chunk_size;
    }

    // Re-interleave
    let out_frames = output_channels[0].len();
    let mut interleaved = Vec::with_capacity(out_frames * ch);
    for frame in 0..out_frames {
        for channel in &output_channels {
            interleaved.push(channel[frame]);
        }
    }

    Ok(interleaved)
}

/// Apply channel conversion then resampling.
pub fn apply_conversion(
    samples: &[f32],
    src_channels: u16,
    src_rate: u32,
    target_channels: &ExportChannels,
    target_rate: Option<u32>,
) -> Result<ConvertedAudio, CoreError> {
    let (converted, out_channels) = convert_channels(samples, src_channels, target_channels);
    let dst_rate = target_rate.unwrap_or(src_rate);
    let resampled = resample(&converted, out_channels, src_rate, dst_rate)?;

    Ok(ConvertedAudio {
        samples: resampled,
        sample_rate: dst_rate,
        channels: out_channels,
    })
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn mono_mixdown() {
        // Stereo: L=0.4, R=0.6 -> mono: 0.5
        let samples = vec![0.4, 0.6, 0.2, 0.8];
        let (out, ch) = convert_channels(&samples, 2, &ExportChannels::Mono);
        assert_eq!(ch, 1);
        assert_eq!(out.len(), 2);
        assert!((out[0] - 0.5).abs() < 1e-6);
        assert!((out[1] - 0.5).abs() < 1e-6);
    }

    #[test]
    fn stereo_upmix() {
        let samples = vec![0.3, -0.3, 0.7];
        let (out, ch) = convert_channels(&samples, 1, &ExportChannels::Stereo);
        assert_eq!(ch, 2);
        assert_eq!(out, vec![0.3, 0.3, -0.3, -0.3, 0.7, 0.7]);
    }

    #[test]
    fn passthrough() {
        let samples = vec![0.1, 0.2, 0.3, 0.4];
        let (out, ch) = convert_channels(&samples, 2, &ExportChannels::Original);
        assert_eq!(ch, 2);
        assert_eq!(out, samples);
    }

    #[test]
    fn resample_noop() {
        let samples = vec![0.1, 0.2, 0.3, 0.4];
        let out = resample(&samples, 1, 44100, 44100).unwrap();
        assert_eq!(out, samples);
    }

    #[test]
    fn resample_changes_length() {
        // Use multiple full chunks to avoid edge effects. 4096 mono samples at 44100 -> 48000.
        let num_samples = 4096;
        let samples: Vec<f32> = (0..num_samples).map(|i| i as f32 / num_samples as f32).collect();
        let out = resample(&samples, 1, 44100, 48000).unwrap();
        // Output should be longer than input since 48000 > 44100
        assert!(
            out.len() > num_samples,
            "expected output longer than input ({num_samples}), got {}",
            out.len()
        );
        // And roughly in the right ratio (within 15% to account for resampler latency/padding)
        let expected = (num_samples as f64 * 48000.0 / 44100.0) as usize;
        assert!(
            out.len() > expected / 2,
            "output too short: expected ~{expected}, got {}",
            out.len()
        );
    }
}