mp4-atom 0.16.1

use crate::*;

// ItemInformationBox. ISO/IEC 14496-12:2022 Section 8.11.3
// This is used to work out where the items are

ext! {
    name: Iloc,
    versions: [0, 1, 2],
    flags: {}
}

#[derive(Debug, Clone, PartialEq, Eq)]
#[cfg_attr(feature = "serde", derive(serde::Serialize, serde::Deserialize))]
pub struct ItemLocationExtent {
    pub item_reference_index: u64,
    pub offset: u64,
    pub length: u64,
}

#[derive(Debug, Clone, PartialEq, Eq)]
#[cfg_attr(feature = "serde", derive(serde::Serialize, serde::Deserialize))]
pub struct ItemLocation {
    pub item_id: u32,
    pub construction_method: u8, // enum?
    pub data_reference_index: u16,
    pub base_offset: u64,
    pub extents: Vec<ItemLocationExtent>,
}

#[derive(Debug, Clone, PartialEq, Eq)]
#[cfg_attr(feature = "serde", derive(serde::Serialize, serde::Deserialize))]
pub struct Iloc {
    pub item_locations: Vec<ItemLocation>,
}

impl AtomExt for Iloc {
    type Ext = IlocExt;

    const KIND_EXT: FourCC = FourCC::new(b"iloc");

    fn decode_body_ext<B: Buf>(buf: &mut B, ext: IlocExt) -> Result<Self> {
        let sizes0 = u8::decode(buf)?;
        let offset_size = sizes0 >> 4;
        let length_size = sizes0 & 0x0F;
        let sizes1 = u8::decode(buf)?;
        let base_offset_size = sizes1 >> 4;
        let index_size: u8 = if ext.version == IlocVersion::V1 || ext.version == IlocVersion::V2 {
            sizes1 & 0x0F
        } else {
            0
        };

        let item_count = if ext.version == IlocVersion::V0 || ext.version == IlocVersion::V1 {
            u16::decode(buf)? as usize
        } else {
            u32::decode(buf)? as usize
        };
        let mut item_locations = vec![];
        for _i in 0..item_count {
            let item_id = if ext.version == IlocVersion::V0 || ext.version == IlocVersion::V1 {
                u16::decode(buf)? as u32
            } else {
                u32::decode(buf)?
            };
            let construction_method: u8 =
                if ext.version == IlocVersion::V1 || ext.version == IlocVersion::V2 {
                    let construction_method_packed = u16::decode(buf)?;
                    (construction_method_packed & 0x0f) as u8
                } else {
                    0
                };
            let data_reference_index = u16::decode(buf)?;
            let base_offset = match base_offset_size {
                0 => 0u64,
                4 => u32::decode(buf)? as u64,
                8 => u64::decode(buf)?,
                _ => return Err(Error::Reserved),
            };
            let extent_count = u16::decode(buf)?;
            let mut extents = vec![];
            for _j in 0..extent_count {
                let item_reference_index: u64 =
                    if ext.version == IlocVersion::V1 || ext.version == IlocVersion::V2 {
                        match index_size {
                            0 => 0,
                            4 => u32::decode(buf)? as u64,
                            8 => u64::decode(buf)?,
                            _ => return Err(Error::Reserved),
                        }
                    } else {
                        0
                    };
                let extent_offset = match offset_size {
                    0 => 0u64,
                    4 => u32::decode(buf)? as u64,
                    8 => u64::decode(buf)?,
                    _ => return Err(Error::Reserved),
                };
                let extent_length = match length_size {
                    0 => 0u64,
                    4 => u32::decode(buf)? as u64,
                    8 => u64::decode(buf)?,
                    _ => return Err(Error::Reserved),
                };
                extents.push(ItemLocationExtent {
                    item_reference_index,
                    offset: extent_offset,
                    length: extent_length,
                });
            }
            item_locations.push(ItemLocation {
                item_id,
                construction_method,
                data_reference_index,
                base_offset,
                extents,
            })
        }
        Ok(Iloc { item_locations })
    }

    fn encode_body_ext<B: BufMut>(&self, buf: &mut B) -> Result<IlocExt> {
        let item_count: u32 = self
            .item_locations
            .len()
            .try_into()
            .map_err(|_| Error::TooLarge(Self::KIND))?;

        let mut base_offset_size = 0u8;
        let mut offset_size = 4u8;
        let mut length_size = 4u8;
        let mut index_size = 0u8;
        let mut needs_version_1 = false;
        let mut needs_version_2 = self.item_locations.len() > u16::MAX as usize;

        for item_location in &self.item_locations {
            if item_location.construction_method > 2 {
                return Err(Error::InvalidSize);
            }
            if item_location.extents.len() > u16::MAX as usize {
                return Err(Error::TooLarge(Self::KIND));
            }

            needs_version_1 |= item_location.construction_method != 0;
            needs_version_2 |= item_location.item_id > u16::MAX as u32;

            if item_location.base_offset > 0 {
                if item_location.base_offset > u32::MAX as u64 {
                    base_offset_size = 8u8;
                } else if base_offset_size != 8 {
                    base_offset_size = 4u8;
                }
            }

            for extent in &item_location.extents {
                if extent.item_reference_index > 0 {
                    needs_version_1 = true;
                    if extent.item_reference_index > u32::MAX as u64 {
                        index_size = 8;
                    } else if index_size != 8 {
                        index_size = 4;
                    }
                }
                if extent.offset > u32::MAX as u64 {
                    offset_size = 8;
                }
                if extent.length > u32::MAX as u64 {
                    length_size = 8;
                }
            }
        }

        let version = if needs_version_2 {
            IlocVersion::V2
        } else if needs_version_1 {
            IlocVersion::V1
        } else {
            IlocVersion::V0
        };

        let size0 = (offset_size << 4) | length_size;
        let size1 = (base_offset_size << 4) | index_size;
        size0.encode(buf)?;
        size1.encode(buf)?;
        if version == IlocVersion::V0 || version == IlocVersion::V1 {
            u16::try_from(item_count)
                .map_err(|_| Error::TooLarge(Self::KIND))?
                .encode(buf)?;
        } else {
            item_count.encode(buf)?;
        }
        for item_location in &self.item_locations {
            if version == IlocVersion::V0 || version == IlocVersion::V1 {
                u16::try_from(item_location.item_id)
                    .map_err(|_| Error::TooLarge(Self::KIND))?
                    .encode(buf)?;
            } else {
                item_location.item_id.encode(buf)?;
            }
            if version == IlocVersion::V1 || version == IlocVersion::V2 {
                (item_location.construction_method as u16).encode(buf)?
            }
            item_location.data_reference_index.encode(buf)?;
            match base_offset_size {
                0 => {}
                4 => u32::try_from(item_location.base_offset)
                    .map_err(|_| Error::TooLarge(Self::KIND))?
                    .encode(buf)?,
                8 => item_location.base_offset.encode(buf)?,
                _ => unreachable!("iloc base_offset_size must be in [0,4,8]"),
            }
            u16::try_from(item_location.extents.len())
                .map_err(|_| Error::TooLarge(Self::KIND))?
                .encode(buf)?;
            for extent in &item_location.extents {
                match index_size {
                    0 => {}
                    4 => u32::try_from(extent.item_reference_index)
                        .map_err(|_| Error::TooLarge(Self::KIND))?
                        .encode(buf)?,
                    8 => extent.item_reference_index.encode(buf)?,
                    _ => unreachable!("iloc index_size must be in [0,4,8]"),
                }
                match offset_size {
                    0 => {}
                    4 => u32::try_from(extent.offset)
                        .map_err(|_| Error::TooLarge(Self::KIND))?
                        .encode(buf)?,
                    8 => extent.offset.encode(buf)?,
                    _ => unreachable!("iloc offset_size must be in [0,4,8]"),
                }
                match length_size {
                    0 => {}
                    4 => u32::try_from(extent.length)
                        .map_err(|_| Error::TooLarge(Self::KIND))?
                        .encode(buf)?,
                    8 => extent.length.encode(buf)?,
                    _ => unreachable!("iloc length_size must be in [0,4,8]"),
                }
            }
        }
        Ok(IlocExt { version })
    }
}

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

    const ENCODED_ILOC_LIBAVIF: &[u8] = &[
        0x00, 0x00, 0x00, 0x1e, 0x69, 0x6c, 0x6f, 0x63, 0x00, 0x00, 0x00, 0x00, 0x44, 0x00, 0x00,
        0x01, 0x00, 0x01, 0x00, 0x00, 0x00, 0x01, 0x00, 0x00, 0x01, 0x38, 0x00, 0x00, 0x00, 0x1a,
    ];

    #[test]
    fn test_iloc_libavif_decode() {
        let buf: &mut std::io::Cursor<&&[u8]> = &mut std::io::Cursor::new(&ENCODED_ILOC_LIBAVIF);

        let iloc: Iloc = Iloc {
            item_locations: vec![ItemLocation {
                item_id: 1,
                construction_method: 0,
                data_reference_index: 0,
                base_offset: 0,
                extents: vec![ItemLocationExtent {
                    item_reference_index: 0,
                    offset: 312,
                    length: 26,
                }],
            }],
        };
        let decoded = Iloc::decode(buf).unwrap();
        assert_eq!(decoded, iloc);
    }

    // Regression for issue #158: out-of-range *_size nibbles (anything
    // outside {0, 4, 8} per ISO/IEC 14496-12 §8.11.3.3) must return
    // Err, not panic.

    #[test]
    fn test_iloc_reserved_base_offset_size() {
        // version=0, flags=0, sizes0=0x00 (offset_size=0, length_size=0),
        // sizes1=0xF0 (base_offset_size=15), item_count=1,
        // item_id=0, data_reference_index=0, then base_offset_size match.
        let body: &[u8] = &[0x00, 0xF0, 0x00, 0x01, 0x00, 0x00, 0x00, 0x00];
        let ext = IlocExt {
            version: IlocVersion::V0,
        };
        assert!(matches!(
            Iloc::decode_body_ext(&mut std::io::Cursor::new(body), ext),
            Err(Error::Reserved)
        ));
    }

    #[test]
    fn test_iloc_reserved_index_size() {
        // V1: sizes1 low nibble is index_size. base_offset_size=0,
        // index_size=15. Reach the index_size match by giving an extent.
        let body: &[u8] = &[
            0x00, 0x0F, // sizes0=0, sizes1=0x0F
            0x00, 0x01, // item_count=1
            0x00, 0x00, // item_id
            0x00, 0x00, // construction_method
            0x00, 0x00, // data_reference_index
            // base_offset_size=0 -> no base_offset
            0x00, 0x01, // extent_count=1
        ];
        let ext = IlocExt {
            version: IlocVersion::V1,
        };
        assert!(matches!(
            Iloc::decode_body_ext(&mut std::io::Cursor::new(body), ext),
            Err(Error::Reserved)
        ));
    }

    #[test]
    fn test_iloc_reserved_offset_size() {
        // V0: sizes0 high nibble is offset_size=0xF.
        let body: &[u8] = &[
            0xF0, 0x00, // sizes0=0xF0 (offset_size=15), sizes1=0
            0x00, 0x01, // item_count=1
            0x00, 0x00, // item_id
            0x00, 0x00, // data_reference_index
            0x00, 0x01, // extent_count=1
        ];
        let ext = IlocExt {
            version: IlocVersion::V0,
        };
        assert!(matches!(
            Iloc::decode_body_ext(&mut std::io::Cursor::new(body), ext),
            Err(Error::Reserved)
        ));
    }

    #[test]
    fn test_iloc_reserved_length_size() {
        // V0: sizes0 low nibble is length_size=0xF, offset_size=0
        // so the offset_size arm passes through 0 and reaches length_size.
        let body: &[u8] = &[
            0x0F, 0x00, // sizes0=0x0F (length_size=15), sizes1=0
            0x00, 0x01, // item_count=1
            0x00, 0x00, // item_id
            0x00, 0x00, // data_reference_index
            0x00, 0x01, // extent_count=1
        ];
        let ext = IlocExt {
            version: IlocVersion::V0,
        };
        assert!(matches!(
            Iloc::decode_body_ext(&mut std::io::Cursor::new(body), ext),
            Err(Error::Reserved)
        ));
    }

    #[test]
    fn test_iloc_avif_encode() {
        let iloc: Iloc = Iloc {
            item_locations: vec![ItemLocation {
                item_id: 1,
                construction_method: 0,
                data_reference_index: 0,
                base_offset: 0,
                extents: vec![ItemLocationExtent {
                    item_reference_index: 0,
                    offset: 312,
                    length: 26,
                }],
            }],
        };
        let mut buf = Vec::new();
        iloc.encode(&mut buf).unwrap();

        assert_eq!(buf.as_slice(), ENCODED_ILOC_LIBAVIF);
    }

    #[test]
    fn test_iloc_version_1_fields_round_trip() {
        let iloc = Iloc {
            item_locations: vec![ItemLocation {
                item_id: 1,
                construction_method: 1,
                data_reference_index: 0,
                base_offset: 0,
                extents: vec![ItemLocationExtent {
                    item_reference_index: u32::MAX as u64 + 1,
                    offset: u32::MAX as u64 + 1,
                    length: u32::MAX as u64 + 2,
                }],
            }],
        };
        let mut buf = Vec::new();
        iloc.encode(&mut buf).unwrap();

        assert_eq!(buf[8], 1); // version 1
        assert_eq!(buf[12], 0x88); // 64-bit extent offset and length
        assert_eq!(buf[13], 0x08); // no base offset, 64-bit extent index
        assert_eq!(Iloc::decode(&mut buf.as_slice()).unwrap(), iloc);
    }

    #[test]
    fn test_iloc_version_2_item_id_round_trip() {
        let iloc = Iloc {
            item_locations: vec![ItemLocation {
                item_id: u16::MAX as u32 + 1,
                construction_method: 0,
                data_reference_index: 0,
                base_offset: 0,
                extents: vec![],
            }],
        };
        let mut buf = Vec::new();
        iloc.encode(&mut buf).unwrap();

        assert_eq!(buf[8], 2);
        assert_eq!(Iloc::decode(&mut buf.as_slice()).unwrap(), iloc);
    }

    #[test]
    fn test_iloc_version_2_item_count_round_trip() {
        let item_location = ItemLocation {
            item_id: 0,
            construction_method: 0,
            data_reference_index: 0,
            base_offset: 0,
            extents: vec![],
        };
        let iloc = Iloc {
            item_locations: vec![item_location; u16::MAX as usize + 1],
        };
        let mut buf = Vec::new();
        iloc.encode(&mut buf).unwrap();

        assert_eq!(buf[8], 2);
        assert_eq!(Iloc::decode(&mut buf.as_slice()).unwrap(), iloc);
    }

    #[test]
    fn test_iloc_rejects_unrepresentable_construction_method() {
        let iloc = Iloc {
            item_locations: vec![ItemLocation {
                item_id: 0,
                construction_method: 3,
                data_reference_index: 0,
                base_offset: 0,
                extents: vec![],
            }],
        };

        assert!(matches!(
            iloc.encode(&mut Vec::new()),
            Err(Error::InvalidSize)
        ));
    }

    #[test]
    fn test_iloc_rejects_too_many_extents() {
        let extent = ItemLocationExtent {
            item_reference_index: 0,
            offset: 0,
            length: 0,
        };
        let iloc = Iloc {
            item_locations: vec![ItemLocation {
                item_id: 0,
                construction_method: 0,
                data_reference_index: 0,
                base_offset: 0,
                extents: vec![extent; u16::MAX as usize + 1],
            }],
        };

        assert!(matches!(
            iloc.encode(&mut Vec::new()),
            Err(Error::TooLarge(kind)) if kind == Iloc::KIND
        ));
    }
}