symphonia-wem 0.1.0

//! Header packet generation.

use std::io::{Cursor, Write as _};

use bitvec::{array::BitArray, order::Lsb0, vec::BitVec, view::BitView as _};
use byteorder::{LittleEndian, WriteBytesExt as _};
use symphonia_core::{
    errors::{Error, Result},
    io::{MediaSourceStream, ReadBytes as _},
};

use crate::{
    bits::{read, read_bool, read_write, read_write_bool, write},
    chunk::Fmt,
    codebook::CodebookLibrary,
};

/// Identification header bytes.
pub(crate) struct IdentHeader([u8; 30]);

impl IdentHeader {
    /// Generate a new identification header.
    pub(crate) fn new(format: &Fmt) -> Result<Self> {
        let mut bytes = Cursor::new([0; 30]);

        // The packet type (ident header)
        bytes.write_u8(1)?;

        // Magic
        bytes.write_all(b"vorbis")?;

        // Vorbis version
        bytes.write_u32::<LittleEndian>(0)?;

        // Audio channels
        bytes.write_u8(
            u8::try_from(format.channels.get())
                .map_err(|_| Error::LimitError("too many channels"))?,
        )?;
        // Audio sample rate
        bytes.write_u32::<LittleEndian>(format.sample_rate.get())?;

        // Maximum bitrate
        bytes.write_i32::<LittleEndian>(0)?;
        // Nominal bitrate
        bytes.write_i32::<LittleEndian>(format.avg_bytes_per_second.cast_signed() * 8)?;
        // Minimum bitrate
        bytes.write_i32::<LittleEndian>(0)?;

        // Blocksizes
        bytes.write_u8(format.block_size_0 | (format.block_size_1 << 4))?;

        // Framing
        bytes.write_u8(1)?;

        Ok(Self(bytes.into_inner()))
    }

    /// Get the bytes.
    pub(crate) const fn into_inner(self) -> [u8; 30] {
        self.0
    }
}

/// Setup header bytes.
pub(crate) struct SetupHeader {
    /// Packet bytes.
    bytes: Vec<u8>,
    /// Mode blockflag.
    pub(crate) mode_block_flags: BitArray<u64, Lsb0>,
    /// Mode bits.
    pub(crate) mode_bits: u32,
}

impl SetupHeader {
    /// Create a fake vorbis setup header packet.
    #[expect(clippy::too_many_lines)]
    pub(crate) fn new(fmt: &Fmt, mss: &mut MediaSourceStream) -> Result<Self> {
        let mut bytes = BitVec::<_, Lsb0>::new();

        // The packet type (setup header)
        bytes.write_u8(5)?;

        // Magic
        bytes.write_all(b"vorbis")?;

        // Skip to the packet and skipping the size
        let setup_start = fmt.setup_packet_offset + 2;
        mss.ignore_bytes(u64::from(setup_start))?;

        // Read until the first audio packet
        let buf =
            mss.read_boxed_slice_exact((fmt.first_audio_packet_offset - setup_start) as usize)?;

        // From now on we read individual bits
        let buf = buf.view_bits::<Lsb0>();

        // Get the amount of codebooks
        let (mut buf, codebook_count_minus_one): (_, u16) = read_write(buf, &mut bytes, 8);
        let codebook_count = codebook_count_minus_one + 1;

        // Rewrite the codebooks
        let codebook_lib = CodebookLibrary::from_aotuv();
        for _ in 0..codebook_count {
            // Get the codebook index
            let id: u16;
            (buf, id) = read(buf, 10);

            // Rewrite the codebook
            let new_bytes = codebook_lib.rebuild(id as usize)?;
            bytes.extend(new_bytes);
        }

        // Time domain transforms placeholder

        // Time count minus one
        write(0_u8, &mut bytes, 6);
        // Dummy time value
        write(0_u16, &mut bytes, 16);

        // Rebuild floors
        let (mut buf, floor_count_minus_one): (_, u8) = read_write(buf, &mut bytes, 6);
        let floor_count = floor_count_minus_one + 1;

        for _ in 0..floor_count {
            // Floor type 1
            write(1_u16, &mut bytes, 16);

            let floor_partitions: usize;
            (buf, floor_partitions) = read_write(buf, &mut bytes, 5);

            // Build the class list
            let mut floor_partition_class_list = Vec::with_capacity(floor_partitions);
            let mut maximum_class = 0;
            for _ in 0..floor_partitions {
                let floor_partition_class: u8;
                (buf, floor_partition_class) = read_write(buf, &mut bytes, 4);

                floor_partition_class_list.push(floor_partition_class);
                maximum_class = maximum_class.max(floor_partition_class);
            }

            let floor_class_dimensions_list = (0..=maximum_class)
                .map(|_| {
                    let class_dimensions_minus_one: u8;
                    (buf, class_dimensions_minus_one) = read_write(buf, &mut bytes, 3);

                    let class_subclasses: u8;
                    (buf, class_subclasses) = read_write(buf, &mut bytes, 2);

                    if class_subclasses != 0 {
                        let masterbook: u8;
                        (buf, masterbook) = read_write(buf, &mut bytes, 8);

                        if u16::from(masterbook) >= codebook_count {
                            return symphonia_core::errors::decode_error(
                                "wem: invalid floor 1 master book",
                            );
                        }
                    }

                    for _ in 0..(1 << u32::from(class_subclasses)) {
                        let subclass_book_plus_one: u8;
                        (buf, subclass_book_plus_one) = read_write(buf, &mut bytes, 8);

                        let subclass_book = i16::from(subclass_book_plus_one) - 1;
                        if subclass_book >= 0 && subclass_book >= codebook_count.cast_signed() {
                            return symphonia_core::errors::decode_error(
                                "wem: invalid floor 1 subclass book",
                            );
                        }
                    }

                    Ok(class_dimensions_minus_one + 1)
                })
                .collect::<Result<Vec<_>>>()?;

            let _floor_multiplier_minus_one: u8;
            (buf, _floor_multiplier_minus_one) = read_write(buf, &mut bytes, 2);

            let range_bits: usize;
            (buf, range_bits) = read_write(buf, &mut bytes, 4);

            for current_class_number in floor_partition_class_list {
                for _ in 0..floor_class_dimensions_list[current_class_number as usize] {
                    let _x: u16;
                    (buf, _x) = read_write(buf, &mut bytes, range_bits);
                }
            }
        }

        // Residues
        let (mut buf, residue_count_minus_one): (_, u8) = read_write(buf, &mut bytes, 6);
        let residue_count = residue_count_minus_one + 1;

        for _ in 0..residue_count {
            let residue_type: u16;
            (buf, residue_type) = read(buf, 2);
            write(residue_type, &mut bytes, 16);

            if residue_type > 2 {
                return symphonia_core::errors::decode_error("wem: corrupt invalid residue type");
            }

            let _residue_begin: u32;
            (buf, _residue_begin) = read_write(buf, &mut bytes, 24);

            let _residue_end: u32;
            (buf, _residue_end) = read_write(buf, &mut bytes, 24);

            let _residue_partition_size_minus_one: u32;
            (buf, _residue_partition_size_minus_one) = read_write(buf, &mut bytes, 24);

            let residue_classifications_minus_one: u8;
            (buf, residue_classifications_minus_one) = read_write(buf, &mut bytes, 6);
            let residue_classifications = residue_classifications_minus_one + 1;

            let residue_classbook: u8;
            (buf, residue_classbook) = read_write(buf, &mut bytes, 8);

            if u16::from(residue_classbook) >= codebook_count {
                return symphonia_core::errors::decode_error("wem: corrupt residue classbook");
            }

            let residue_cascade = (0..residue_classifications)
                .map(|_| {
                    let low_bits: u8;
                    (buf, low_bits) = read_write(buf, &mut bytes, 3);

                    let bit_flag;
                    (buf, bit_flag) = read_bool(buf);
                    bytes.push(bit_flag);
                    let high_bits = if bit_flag {
                        let high_bits: u8;
                        (buf, high_bits) = read_write(buf, &mut bytes, 5);

                        high_bits
                    } else {
                        0
                    };

                    u32::from(high_bits) * 8 + u32::from(low_bits)
                })
                .collect::<Vec<_>>();

            residue_cascade
                .into_iter()
                .try_for_each(|residue_cascade| {
                    for offset in 0..8 {
                        if (residue_cascade & (1 << offset)) > 0 {
                            let residue_book: u8;
                            (buf, residue_book) = read_write(buf, &mut bytes, 8);

                            if u16::from(residue_book) >= codebook_count {
                                return symphonia_core::errors::decode_error(
                                    "wem: corrupt residue book",
                                );
                            }
                        }
                    }

                    Ok(())
                })?;
        }

        // Mapping
        let (mut buf, mapping_count_minus_one): (_, u8) = read_write(buf, &mut bytes, 6);
        let mapping_count = mapping_count_minus_one + 1;

        for _ in 0..mapping_count {
            // Mapping type 0
            write(0_u16, &mut bytes, 16);

            let submaps_flag;
            (buf, submaps_flag) = read_write_bool(buf, &mut bytes);
            let submaps = if submaps_flag {
                let submaps_minus_one: u8;
                (buf, submaps_minus_one) = read_write(buf, &mut bytes, 4);

                submaps_minus_one + 1
            } else {
                1
            };

            let square_polar_flag;
            (buf, square_polar_flag) = read_write_bool(buf, &mut bytes);
            if square_polar_flag {
                let coupling_steps_minus_one: u16;
                (buf, coupling_steps_minus_one) = read_write(buf, &mut bytes, 8);
                let coupling_steps = coupling_steps_minus_one + 1;

                for _ in 0..coupling_steps {
                    let magnitude: u32;
                    (buf, magnitude) = read_write(
                        buf,
                        &mut bytes,
                        crate::math::log2(u32::from(fmt.channels.get()) - 1) as usize,
                    );

                    let angle: u32;
                    (buf, angle) = read_write(
                        buf,
                        &mut bytes,
                        crate::math::log2(u32::from(fmt.channels.get()) - 1) as usize,
                    );

                    if angle == magnitude
                        || magnitude >= u32::from(fmt.channels.get())
                        || angle >= u32::from(fmt.channels.get())
                    {
                        return symphonia_core::errors::decode_error("wem: corrupt coupling");
                    }
                }
            }

            let mapping_reserved: u8;
            (buf, mapping_reserved) = read_write(buf, &mut bytes, 2);
            if mapping_reserved != 0 {
                return symphonia_core::errors::decode_error(
                    "wem: corrupt mapping reserved field nonzero",
                );
            }

            if submaps > 1 {
                for _ in 0..fmt.channels.get() {
                    let mapping_mux: u8;
                    (buf, mapping_mux) = read_write(buf, &mut bytes, 4);

                    if mapping_mux >= submaps {
                        return symphonia_core::errors::decode_error(
                            "wem: corrupt mapping mux >= submaps",
                        );
                    }
                }
            }

            for _ in 0..submaps {
                let _time_config: u8;
                (buf, _time_config) = read_write(buf, &mut bytes, 8);

                let floor_number: u8;
                (buf, floor_number) = read_write(buf, &mut bytes, 8);
                if floor_number >= floor_count {
                    return symphonia_core::errors::decode_error("wem: corrupt floor mapping");
                }

                let residue_number: u8;
                (buf, residue_number) = read_write(buf, &mut bytes, 8);
                if residue_number >= residue_count {
                    return symphonia_core::errors::decode_error("wem: corrupt residue mapping");
                }
            }
        }

        // Mode count
        let (mut buf, mode_count_minus_one): (_, u8) = read_write(buf, &mut bytes, 6);
        let mode_count = mode_count_minus_one + 1;

        // Set all the mode flags
        let mut mode_block_flags = BitArray::new(0);
        for index in 0..mode_count {
            let block_flag;
            (buf, block_flag) = read_write_bool(buf, &mut bytes);

            // Push the block flag
            mode_block_flags.set(index as usize, block_flag);

            // Window type
            write(0_u16, &mut bytes, 16);
            // Transform type
            write(0_u16, &mut bytes, 16);

            // Handle mapping error
            let mapping: u8;
            (buf, mapping) = read_write(buf, &mut bytes, 8);
            if mapping >= mapping_count {
                return symphonia_core::errors::decode_error("wem: corrupt invalid mode mapping");
            }
        }

        let mode_bits = crate::math::log2(u32::from(mode_count_minus_one));

        // Framing
        write(1_u8, &mut bytes, 1);

        // TODO: verify size
        Ok(Self {
            bytes: bytes.into_vec(),
            mode_block_flags,
            mode_bits,
        })
    }

    /// Get the bytes.
    pub(crate) fn into_inner(self) -> Vec<u8> {
        self.bytes
    }
}