sharkyflac 0.2.0

use std::io::{self, Read, Seek};

use crate::bit_io::BitReader;
use crate::frame::subframe::{self};
use crate::frame::{self, Frame};
use crate::metadata::{self, Application, Cuesheet, Picture, SeekTable, Streaminfo, VorbisComment};

#[derive(Debug, thiserror::Error)]
pub enum DecodeError {
    #[error("parse: {0}")]
    Parse(#[from] metadata::Error),

    #[error("frame: {0}")]
    Frame(#[from] frame::Error),

    #[error("IO: {0}")]
    Io(#[from] io::Error),

    #[error(
        "channel count mismatch: frame declares {frame} channels, but STREAMINFO declares {streaminfo}"
    )]
    ChannelCountMismatch {
        frame:      usize,
        streaminfo: usize,
    },

    #[error("There were multiple STREAMINFO blocks")]
    MultipleStreaminfo,

    #[error("There were multiple seektables")]
    MultipleSeektable,

    #[error("There were multiple Vorbis comment blocks")]
    MultipleVorbisComment,
}

/// This is a high-level API. If you need support for changing variable
/// frame headers (e.g. channel counts, bits per sample, etc.), then use the
/// more flexible, low-level `Frame::read_streaming` interface.
///
/// It's highly recommended to use a [`BufReader`][std::io::BufReader] if you're
/// reading from a filesystem or the network.
pub struct Flac<R: Read + Seek> {
    pub streaminfo:   Streaminfo,
    pub applications: Vec<Application>,
    pub comment:      Option<VorbisComment>,
    pub cuesheets:    Vec<Cuesheet>,
    pub pictures:     Vec<Picture>,
    pub seektable:    Option<SeekTable>,

    /// Reader for parsing the frames
    reader: BitReader<R>,
}

impl<R: Read + Seek> Flac<R> {
    pub fn new(inner: R) -> Result<Self, DecodeError> {
        let mut reader = BitReader::new(inner);

        let mut magic = [0u8; 4];
        reader.read_exact(&mut magic)?;
        if magic != crate::MAGIC {
            return Err(io::Error::new(io::ErrorKind::InvalidData, "not a FLAC file").into());
        }

        let mut streaminfo = None;
        let mut applications = Vec::new();
        let mut comment = None;
        let mut cuesheets = Vec::new();
        let mut pictures = Vec::new();
        let mut seektable = None;

        for block in metadata::BlockIter::new(&mut reader) {
            let block = block?;
            match block.data {
                metadata::BlockData::Streaminfo(d) => streaminfo = Some(d),
                metadata::BlockData::Application(d) => applications.push(d),
                metadata::BlockData::VorbisComment(d) => comment
                    .replace(d)
                    .is_none_or_err(DecodeError::MultipleVorbisComment)?,
                metadata::BlockData::Cuesheet(d) => cuesheets.push(d),
                metadata::BlockData::Picture(d) => pictures.push(d),
                metadata::BlockData::SeekTable(d) => seektable
                    .replace(d)
                    .is_none_or_err(DecodeError::MultipleSeektable)?,
                metadata::BlockData::Padding => {}
            }
        }

        let streaminfo = streaminfo.ok_or_else(|| {
            io::Error::new(io::ErrorKind::InvalidData, "missing STREAMINFO block")
        })?;

        Ok(Self {
            streaminfo,
            applications,
            comment,
            cuesheets,
            pictures,
            seektable,
            reader,
        })
    }

    /// Get the inner [`BitReader`]
    pub fn inner(self) -> BitReader<R> {
        self.reader
    }

    /// Returns an iterator over decoded frames
    pub fn frames(self) -> FrameIter<R> {
        FrameIter {
            total_channels: self.streaminfo.num_channels(),
            inner:          self,
            blocks:         Vec::new(),
        }
    }

    /// Parse all the samples and return them interleaved.
    pub fn interleaved_samples(self) -> Result<Vec<i32>, DecodeError> {
        let mut samples = Vec::new();

        for frame in self.frames() {
            let frame = frame?;
            samples.extend(frame.samples);
        }

        Ok(samples)
    }
}

/// A single decoded frame. Samples are stored as *interleaved* PCM. The number
/// of channels is equal to `streaminfo.channel_count()`.
#[derive(Debug, Clone)]
pub struct DecodedFrame {
    /// The frame header
    pub header: frame::Header,

    /// Interleaved samples: `[ch0, ch1, ch0, ch1, ...]`. The total sample count
    /// (length) is equal to `header.block_size.samples() *
    /// streaminfo.channel_count()`.
    pub samples: Vec<i32>,

    /// Number of *per channel* samples in this frame.
    pub block_size: usize,
}

/// Iterator returned by `Flac::frames()`.
pub struct FrameIter<R: Read + Seek> {
    inner:          Flac<R>,
    total_channels: usize,

    blocks: Vec<subframe::Block>,
}

impl<R: Read + Seek> Iterator for FrameIter<R> {
    type Item = Result<DecodedFrame, DecodeError>;

    fn next(&mut self) -> Option<Self::Item> {
        let (header, _footer) = match Frame::read_streaming(
            &mut self.inner.reader,
            &self.inner.streaminfo,
            |_, block| {
                let mut block = block?;
                block.scale(self.inner.streaminfo.bits_per_sample());
                self.blocks.push(block);
                Ok(())
            },
        ) {
            Ok(x) => x,
            Err(frame::Error::Io(e)) if e.kind() == io::ErrorKind::UnexpectedEof => return None,
            Err(e) => return Some(Err(e.into())),
        };

        if header.channels.count() != self.total_channels {
            return Some(Err(DecodeError::ChannelCountMismatch {
                frame:      header.channels.count(),
                streaminfo: self.total_channels,
            }));
        }

        // Decorrelate stereo
        if let [ch0, ch1] = self.blocks.as_mut_slice() {
            subframe::Block::decorrelate(ch0, ch1, header.channels);
        }

        let block_size = header.block_size.samples();
        let mut samples = Vec::with_capacity(block_size * self.total_channels);

        for sample_idx in 0..block_size {
            for channel in &self.blocks {
                samples.push(channel[sample_idx]);
            }
        }

        self.blocks.clear();

        Some(Ok(DecodedFrame {
            header,
            samples,
            block_size,
        }))
    }
}

trait OptionToResult {
    fn is_none_or_err<E>(&self, err: E) -> Result<(), E>;
}

impl<T> OptionToResult for Option<T> {
    #[inline]
    fn is_none_or_err<E>(&self, err: E) -> Result<(), E> {
        if self.is_none() { Ok(()) } else { Err(err) }
    }
}