mod error;

use rubato::{FftFixedInOut, Resampler};
use symphonia::core::{
    audio::{AudioBuffer, Layout, SampleBuffer, Signal, SignalSpec},
    codecs::{Decoder, DecoderOptions, CODEC_TYPE_NULL},
    errors::Error as SymError,
    formats::{FormatOptions, FormatReader},
    io::MediaSourceStream,
    meta::MetadataOptions,
    probe::Hint
};
use tool_helper::{Error, Input};

#[derive(Copy, Clone, PartialEq)]
pub enum Orality {
    Stereo,
    Mono,
}

pub struct CDAudioSamples {
    samples:   Vec::<i16>,
    orality:   Orality,
}

impl CDAudioSamples {
    pub fn new(samples: Vec<i16>, channels: usize) -> Result<CDAudioSamples, Error> {
        let orality = match channels {
            0 => return Err(Error::from_str("Input file has no audio channels")),
            1 => Orality::Mono,
            2 => Orality::Stereo,
            _ => return Err(Error::from_str("Only Mono and Stereo input are supported")),
        };
        Ok(CDAudioSamples{samples, orality})
    }

    pub fn samples(&self) -> &Vec::<i16> {
        &self.samples
    }

    pub fn orality(&self) -> Orality {
        self.orality.clone()
    }
}

struct InternalAudioSamples {
    planar_samples: Vec<Vec<f32>>,
    frequency:      u32,
}

impl InternalAudioSamples {
    pub fn new(planar_samples: Vec<Vec<f32>>, frequency: u32) -> Result<InternalAudioSamples, Error> {
        if planar_samples.len() < 1 || planar_samples.len() > 2{
            Err(Error::from_str("Audio samples need to be either mono or stereo"))
        }

        else {
            Ok(InternalAudioSamples{planar_samples, frequency})
        }
    }

    pub fn into_audio_buffer(self) -> AudioBuffer<f32> {
        let     duration         = self.sample_len() as u64;
        let mut new_audio_buffer = AudioBuffer::new(duration, SignalSpec::new_with_layout(self.frequency, if self.planar_samples.len() == 1 {Layout::Mono} else {Layout::Stereo}));

        new_audio_buffer.render_silence(None);
        for (channel_idx, channel) in self.planar_samples.into_iter().enumerate() {
            let dst_channel = new_audio_buffer.chan_mut(channel_idx);
            for (sample_idx, sample) in channel.into_iter().enumerate() {
                dst_channel[sample_idx] = sample;
            }
        }
        new_audio_buffer
    }

    pub fn sample_len(&self) -> usize {
        self.planar_samples[0].len()
    }

    pub fn channels(&self) -> usize {
        self.planar_samples.len()
    }

    pub fn planar_slices(&self) -> Vec<&[f32]> {
        let mut planar_slices = Vec::new();

        for channel in &self.planar_samples {
            planar_slices.push(channel.as_slice());
        }
        planar_slices
    }
}

pub fn load_as_i16_audio(input: Input, target_frequency: u32) -> Result<CDAudioSamples, Error> {
    let raw_audio = load_raw_audio(input)?;
    let raw_audio = resample(raw_audio, target_frequency)?;

    down_sample_interleave(raw_audio)
}

fn load_raw_audio(input: Input) -> Result<InternalAudioSamples, Error> {
    let media_stream = MediaSourceStream::new(Box::new(load_to_ram(input)?), Default::default());
    let format       = symphonia::default::get_probe().format(&Hint::new(), media_stream, &FormatOptions::default(), &MetadataOptions::default()).map_err(error::probe)?.format;
    let track        = format.tracks().iter().find(|t| t.codec_params.codec != CODEC_TYPE_NULL).ok_or_else(error::find_track)?;
    
    // Create a decoder for the track.
    let decoder  = symphonia::default::get_codecs().make(&track.codec_params, &DecoderOptions::default()).map_err(error::decoder)?;
    let track_id = track.id;

    decode(format, decoder, track_id)
}

fn decode(mut format: Box<dyn FormatReader>, mut decoder: Box<dyn Decoder>, track_id: u32) -> Result<InternalAudioSamples, Error> {
    let mut samples       = Vec::new();
    let mut channel_count = 0;
    let mut frequency     = 0;
    let mut read_buffer   = None;

    loop {
        // Get the next packet from the media format.
        let packet = match format.next_packet() {
            Ok(packet) => packet,
            Err(err) => {
                if let SymError::IoError(io_err) = &err {
                    if io_err.kind() == std::io::ErrorKind::UnexpectedEof {
                        return InternalAudioSamples::new(samples, frequency);
                    }                    
                }
                return Err(error::next_packet(err));
            }
        };

        // Consume any new metadata that has been read since the last packet.
        format.metadata().skip_to_latest();

        // If the packet does not belong to the selected track, skip over it.
        if packet.track_id() != track_id {
            continue;
        }

        // Decode the packet into audio samples.
        let packet = decoder.decode(&packet).map_err(error::decode)?;
        if read_buffer.is_none() {
            let duration = packet.capacity() as u64;
            let specs    = packet.spec();

            channel_count = specs.channels.count();
            frequency     = specs.rate;
            read_buffer   = Some(SampleBuffer::<f32>::new(duration, packet.spec().clone()));
            
            for _ in 0..channel_count {
                samples.push(Vec::new());
            }
        }

        if let Some(read_buffer) = &mut read_buffer {
            read_buffer.copy_planar_ref(packet);
            let     cur_samples = read_buffer.samples();
            let mut cur_samples = cur_samples.chunks(cur_samples.len()/channel_count);

            for dst_sample in &mut samples {
                dst_sample.extend(cur_samples.next().ok_or_else(|| Error::from_str("Not enough channels in input as expected"))?);
            }
        }        
    }
}

fn resample(input: InternalAudioSamples, target_frequency: u32) -> Result<InternalAudioSamples, Error> {
    const HIGH_QUALITY_CHUNKS:usize = (1024*10)*100;
    fn process_partial(input_option: Option<&[&[f32]]>, resampler: &mut FftFixedInOut<f32>, planar_output: &mut Vec<Vec<f32>>) -> Result<(), Error> {
        let new_samples = resampler.process_partial(input_option, None).map_err(error::resample)?;
        for (channel, channel_samples) in new_samples.into_iter().enumerate() {
            planar_output[channel].extend(channel_samples.iter());
        }
        Ok(())
    }

    let chunk_size         = HIGH_QUALITY_CHUNKS;
    let mut planar_input   = input.planar_slices();
    let mut resampler      = FftFixedInOut::<f32>::new(input.frequency as usize, target_frequency as usize, chunk_size, input.channels()).map_err(error::resampler_construction)?;
    let     delay          = resampler.output_delay();
    let mut sample_len     = input.sample_len();
    let     new_sample_len = (sample_len as f64*(target_frequency as f64/input.frequency as f64)) as usize;
    let mut planar_output  = {
        let mut planar_output = Vec::new();

        for _ in 0..planar_input.len() {
            planar_output.push(Vec::<f32>::new());
        }
        planar_output
    };

    loop {
        let next_input_frames = resampler.input_frames_next();
        if next_input_frames > sample_len {
            if sample_len > 0 {
                // Still frames left
                process_partial(Some(&planar_input), &mut resampler, &mut planar_output)?;
            }
            break;
        }

        let new_samples = resampler.process(&planar_input, None).map_err(error::resample)?;
        for (channel, slice) in planar_input.iter_mut().enumerate() {
            *slice = &slice[next_input_frames..];
            planar_output[channel].extend(new_samples[channel].iter());
        }
        sample_len -= next_input_frames;
    }

    if planar_output[0].len() < delay + new_sample_len {
        // Flush
        process_partial(None, &mut resampler, &mut planar_output)?;
    }

    for channel in &mut planar_output {
        let start = delay;
        let end   = start + new_sample_len;
        *channel = channel[start..end].into();
    }
    InternalAudioSamples::new(planar_output, target_frequency)
}

fn down_sample_interleave(input: InternalAudioSamples) -> Result<CDAudioSamples, Error> {
    let     channels      = input.channels();
    let     audio_buffer  = input.into_audio_buffer();
    let mut sample_buffer = SampleBuffer::<i16>::new(audio_buffer.capacity() as u64, audio_buffer.spec().clone());

    sample_buffer.copy_interleaved_typed::<f32>(&audio_buffer);
    CDAudioSamples::new(sample_buffer.samples().to_vec(), channels)
}

fn load_to_ram(mut input: Input) -> Result<std::io::Cursor<Vec<u8>>, Error> {
    let mut buffer = Vec::default();
    
    input.read_to_end(&mut buffer)?;
    Ok(std::io::Cursor::new(buffer))
}