path: root/remuxer/src/demuxers/flac.rs

/*
    This file is part of jellything (https://codeberg.org/metamuffin/jellything)
    which is licensed under the GNU Affero General Public License (version 3); see /COPYING.
    Copyright (C) 2025 metamuffin <metamuffin.org>
*/

use crate::demuxers::{Demuxer, DemuxerNew, ReadSeek};
use anyhow::{Result, anyhow, bail};
use std::io::{BufReader, Read, Seek, SeekFrom};
use winter_matroska::{Audio, Cluster, Info, TrackEntry, TrackType, Tracks};

pub struct FlacDemuxer {
    reader: BufReader<Box<dyn ReadSeek>>,
    metadata: Option<Vec<MetadataBlock>>,
    first_frame_offset: u64,
}
impl DemuxerNew for FlacDemuxer {
    fn new(reader: Box<dyn ReadSeek>) -> Self {
        Self {
            reader: BufReader::new(reader),
            metadata: None,
            first_frame_offset: 0,
        }
    }
}

struct MetadataBlock {
    r#type: u8,
    data: Vec<u8>,
}
#[allow(unused)]
impl MetadataBlock {
    const TY_STREAMINFO: u8 = 0;
    const TY_PADDING: u8 = 1;
    const TY_APPLICATION: u8 = 2;
    const TY_SEEK_TABLE: u8 = 3;
    const TY_VORBIS_COMMENT: u8 = 4;
    const TY_CUESHEET: u8 = 5;
    const TY_PICTURE: u8 = 6;
}

#[allow(unused)]
struct StreamInfo {
    min_block_size: u16,
    max_block_size: u16,
    min_frame_size: u32,
    max_frame_size: u32,
    sample_rate: u32,
    num_channels: u8,
    bit_depth: u8,
}
impl StreamInfo {
    pub fn parse(d: &[u8; 22]) -> Self {
        let k = u64::from_be_bytes([d[10], d[11], d[12], d[13], d[14], d[15], d[16], d[17]]);
        Self {
            min_block_size: u16::from_be_bytes([d[0], d[1]]),
            max_block_size: u16::from_be_bytes([d[2], d[3]]),
            min_frame_size: u32::from_be_bytes([0, d[4], d[5], d[6]]),
            max_frame_size: u32::from_be_bytes([0, d[7], d[8], d[9]]),
            sample_rate: (k >> (64 - 20)) as u32,
            num_channels: ((k >> (64 - 20 - 3)) & 0b111) as u8 + 1,
            bit_depth: ((k >> (64 - 20 - 3 - 5)) & 0b11111) as u8 + 1,
        }
    }
}

impl FlacDemuxer {
    fn read_metadata(&mut self) -> Result<&Vec<MetadataBlock>> {
        if self.metadata.is_some() {
            return Ok(self.metadata.as_ref().unwrap());
        }
        self.reader.seek(SeekFrom::Start(0))?;

        let mut magic = [0u8; 4];
        self.reader.read_exact(&mut magic)?;
        if magic != *b"fLaC" {
            bail!("incorrect magic bytes")
        }

        let mut blocks = Vec::new();
        loop {
            let mut header = [0u8; 4];
            self.reader.read_exact(&mut header)?;
            let last = header[0] & 0x80 != 0;
            let r#type = header[0] & 0x7f;
            let size = u32::from_be_bytes(header) & 0x00FFFFFF;

            let mut data = vec![0u8; size as usize];
            self.reader.read_exact(&mut data)?;
            blocks.push(MetadataBlock { r#type, data });

            if last {
                break;
            }
        }

        self.first_frame_offset = self.reader.stream_position()?;

        self.metadata = Some(blocks);
        return Ok(self.metadata.as_ref().unwrap());
    }
    fn stream_info(&mut self) -> Result<StreamInfo> {
        let m = self.read_metadata()?;
        if m[0].r#type == MetadataBlock::TY_STREAMINFO {
            Ok(StreamInfo::parse(m[0].data.as_slice().try_into().map_err(
                |_| anyhow!("Streaminfo block is not 22 bytes"),
            )?))
        } else {
            bail!("first metadata block is not Streaminfo")
        }
    }
}
impl Demuxer for FlacDemuxer {
    fn info(&mut self) -> Result<Info> {
        Ok(Info {
            duration: Some(120000.), // TODO
            timestamp_scale: 1_000_000,
            ..Default::default()
        })
    }
    fn tracks(&mut self) -> Result<Option<Tracks>> {
        let si = self.stream_info()?;
        let mut buf = Vec::new();
        buf.extend(b"fLaC");
        let meta = self.read_metadata()?;
        for (i, mb) in meta.iter().enumerate() {
            buf.push(if i == meta.len() - 1 { 0x80 } else { 0 } | mb.r#type);
            buf.extend(&u32::to_be_bytes(0)[1..]);
            buf.extend(&mb.data);
        }
        let te = TrackEntry {
            codec_id: "A_FLAC".to_string(),
            codec_private: Some(buf),
            track_number: 1,
            track_type: TrackType::Audio,
            audio: Some(Audio {
                bit_depth: Some(si.bit_depth as u64),
                channels: si.num_channels as u64,
                sampling_frequency: si.sample_rate as f64,
                ..Default::default()
            }),
            ..Default::default()
        };
        Ok(Some(Tracks { entries: vec![te] }))
    }

    fn seek_cluster(&mut self, position: Option<u64>) -> Result<()> {
        if let Some(position) = position {
            self.reader.seek(SeekFrom::Start(position))?;
        } else {
            if self.first_frame_offset == 0 {
                self.read_metadata()?;
            }
            self.reader.seek(SeekFrom::Start(self.first_frame_offset))?;
        }
        Ok(())
    }
    fn read_cluster(&mut self) -> Result<Option<(u64, Cluster)>> {
        let mut header = [0u8; 5];
        self.reader.read_exact(&mut header)?;

        let sync_and_blocking = u16::from_be_bytes([header[0], header[1]]);
        let fixed_blocking = match sync_and_blocking {
            0xfff8 => true,
            0xfff9 => false,
            _ => bail!("invalid frame sync code"),
        };

        let block_size_bits = header[2] >> 4;
        let sample_rate_bits = header[2] & 0x0f;

        let channel_count = match header[3] >> 4 {
            x @ 0..8 => x + 1,
            8..11 => 2,
            _ => bail!("reserved channel bits used"),
        };
        let bit_depth = match (header[3] >> 1) & 0b111 {
            0b000 => 0, // TODO streaminfo
            0b001 => 8,
            0b010 => 12,
            0b011 => bail!("reserved bit depth used"),
            0b100 => 16,
            0b101 => 20,
            0b110 => 24,
            0b111 => 32,
            _ => unreachable!(),
        };
        if header[3] & 1 != 0 {
            bail!("reserveed bit set")
        }

        let coded_num_length = match header[4].leading_ones() {
            0 => 0,
            1 => bail!("invalid coded number vint length (loc=1)"),
            x @ 2..8 => x - 1,
            8 => bail!("invalid coded number vint length (loc=8)"),
            _ => unreachable!(),
        };
        let mut coded_num_buf = [0u8; 6];
        self.reader.read_exact(&mut coded_num_buf)?;

        let block_size = match block_size_bits {
            0b0000 => bail!("reserved block size used"),
            0b0001 => 192,
            x @ 0b0010..=0b0101 => 144 * 2u32.pow(x as u32),
            0b0110 => {
                let mut buf = [0u8; 1];
                self.reader.read_exact(&mut buf)?;
                buf[0] as u32 + 1
            }
            0b0111 => {
                let mut buf = [0u8; 2];
                self.reader.read_exact(&mut buf)?;
                u16::from_be_bytes(buf) as u32 + 1
            }
            x @ 0b1000..=0b1111 => 2u32.pow(x as u32),
            _ => unreachable!(),
        };

        let sample_rate = match sample_rate_bits {
            0b0000 => 0, // TODO streaminfo,
            0b0001 => 88200,
            0b0010 => 176400,
            0b0011 => 192000,
            0b0100 => 8000,
            0b0101 => 16000,
            0b0110 => 22050,
            0b0111 => 24000,
            0b1000 => 32000,
            0b1001 => 44100,
            0b1010 => 48000,
            0b1011 => 96000,
            0b1100 => {
                let mut buf = [0u8; 1];
                self.reader.read_exact(&mut buf)?;
                buf[0] as u32 * 1000
            }
            0b1101 => {
                let mut buf = [0u8; 2];
                self.reader.read_exact(&mut buf)?;
                u16::from_be_bytes(buf) as u32
            }
            0b1110 => {
                let mut buf = [0u8; 2];
                self.reader.read_exact(&mut buf)?;
                u16::from_be_bytes(buf) as u32 * 10
            }
            0b1111 => bail!("forbidden sample rate bits used"),
            _ => unreachable!(),
        };

        let mut crc_buf = [0u8; 1];
        self.reader.read_exact(&mut crc_buf)?;

        Ok(None)
    }
}