mp4-edit 0.1.1

use bon::bon;
use derive_more::{Deref, DerefMut};
use std::fmt;

use crate::{
    atom::{
        util::{DebugList, DebugUpperHex},
        FourCC,
    },
    parser::ParseAtomData,
    writer::SerializeAtom,
    ParseError,
};

#[cfg(feature = "experimental-trim")]
use {crate::atom::stsz::RemovedSampleSizes, anyhow::anyhow, std::ops::Range};

pub const STCO: FourCC = FourCC::new(b"stco");
pub const CO64: FourCC = FourCC::new(b"co64");

#[derive(Default, Clone, Deref, DerefMut, PartialEq, Eq)]
pub struct ChunkOffsets(Vec<u64>);

impl ChunkOffsets {
    pub fn into_inner(self) -> Vec<u64> {
        self.0
    }

    pub fn inner(&self) -> &[u64] {
        &self.0
    }
}

impl From<Vec<u64>> for ChunkOffsets {
    fn from(value: Vec<u64>) -> Self {
        Self(value)
    }
}

impl FromIterator<u64> for ChunkOffsets {
    fn from_iter<T: IntoIterator<Item = u64>>(iter: T) -> Self {
        Self(Vec::from_iter(iter))
    }
}

impl fmt::Debug for ChunkOffsets {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        fmt::Debug::fmt(&DebugList::new(self.0.iter().map(DebugUpperHex), 10), f)
    }
}

/// Chunk Offset Atom - contains file offsets of chunks
#[derive(Default, Debug, Clone)]
pub struct ChunkOffsetAtom {
    /// Version of the stco atom format (0)
    pub version: u8,
    /// Flags for the stco atom (usually all zeros)
    pub flags: [u8; 3],
    /// List of chunk offsets
    pub chunk_offsets: ChunkOffsets,
    /// Whether this uses 64-bit offsets (co64) or 32-bit (stco)
    pub is_64bit: bool,
}

#[cfg(feature = "experimental-trim")]
#[derive(Debug)]
pub(crate) enum ChunkOffsetOperationUnresolved {
    Remove(Range<usize>),
    Insert {
        /// unadjusted chunk index used to get reference offset
        chunk_index_unadjusted: usize,
        /// chunk index that takes into account all previous operations
        chunk_index: usize,
        /// sample indices used to calculate delta from old offset
        sample_indices: Range<usize>,
    },
    ShiftRight {
        /// unadjusted chunk index used to get reference offset
        chunk_index_unadjusted: usize,
        /// chunk index that takes into account all previous operations
        chunk_index: usize,
        /// sample indices used to calculate delta from old offset
        sample_indices: Range<usize>,
    },
}

#[cfg(feature = "experimental-trim")]
impl ChunkOffsetOperationUnresolved {
    pub fn resolve(
        self,
        chunk_offsets: &ChunkOffsets,
        removed_sample_sizes: &RemovedSampleSizes,
    ) -> anyhow::Result<ChunkOffsetOperation> {
        let derive_new_offset = |chunk_index: usize, sample_indices: Range<usize>| {
            let prev_offset = *chunk_offsets
                .get(chunk_index)
                .ok_or_else(|| anyhow!("chunk index {chunk_index} not found"))?;

            let delta = removed_sample_sizes
                .get_sizes(sample_indices.clone())
                .ok_or_else(|| anyhow!("sample indices {sample_indices:?} not found"))?
                .iter()
                .map(|s| *s as u64)
                .sum::<u64>();

            Ok::<u64, anyhow::Error>(prev_offset + delta)
        };

        Ok(match self {
            Self::Remove(chunk_offsets) => ChunkOffsetOperation::Remove(chunk_offsets),
            Self::Insert {
                chunk_index_unadjusted,
                chunk_index,
                sample_indices,
            } => {
                let offset = derive_new_offset(chunk_index_unadjusted - 1, sample_indices)?;
                ChunkOffsetOperation::Insert(chunk_index, offset)
            }
            Self::ShiftRight {
                chunk_index_unadjusted,
                chunk_index,
                sample_indices,
            } => {
                let new_offset = derive_new_offset(chunk_index_unadjusted, sample_indices)?;
                ChunkOffsetOperation::Replace(chunk_index, new_offset)
            }
        })
    }
}

#[cfg(feature = "experimental-trim")]
#[derive(Debug)]
pub(crate) enum ChunkOffsetOperation {
    Remove(Range<usize>),
    Insert(usize, u64),
    Replace(usize, u64),
}

#[bon]
impl ChunkOffsetAtom {
    #[builder]
    pub fn new(
        #[builder(default = 0)] version: u8,
        #[builder(default = [0u8; 3])] flags: [u8; 3],
        #[builder(with = FromIterator::from_iter)] chunk_offsets: Vec<u64>,
        #[builder(default = false)] is_64bit: bool,
    ) -> Self {
        Self {
            version,
            flags,
            chunk_offsets: chunk_offsets.into(),
            is_64bit,
        }
    }

    /// Returns the total number of chunks
    pub fn chunk_count(&self) -> usize {
        self.chunk_offsets.len()
    }

    /// Applies a list of operations
    #[cfg(feature = "experimental-trim")]
    pub(crate) fn apply_operations(&mut self, ops: Vec<ChunkOffsetOperation>) {
        for op in ops {
            match op {
                ChunkOffsetOperation::Remove(chunk_indices_to_remove) => {
                    self.chunk_offsets.drain(chunk_indices_to_remove);
                }
                ChunkOffsetOperation::Insert(chunk_index, offset) => {
                    self.chunk_offsets.insert(chunk_index, offset);
                }
                ChunkOffsetOperation::Replace(chunk_index, new_offset) => {
                    let chunk = self
                        .chunk_offsets
                        .get_mut(chunk_index)
                        .expect("chunk offset must exist");
                    *chunk = new_offset;
                }
            }
        }
    }
}

impl<S: chunk_offset_atom_builder::State> ChunkOffsetAtomBuilder<S> {
    pub fn chunk_offset(
        self,
        chunk_offset: impl Into<u64>,
    ) -> ChunkOffsetAtomBuilder<chunk_offset_atom_builder::SetChunkOffsets<S>>
    where
        S::ChunkOffsets: chunk_offset_atom_builder::IsUnset,
    {
        self.chunk_offsets(vec![chunk_offset.into()])
    }
}

impl ParseAtomData for ChunkOffsetAtom {
    fn parse_atom_data(atom_type: FourCC, input: &[u8]) -> Result<Self, ParseError> {
        crate::atom::util::parser::assert_atom_type!(atom_type, STCO, CO64);
        use crate::atom::util::parser::stream;
        use winnow::Parser;
        Ok(match atom_type {
            STCO => parser::parse_stco_data.parse(stream(input))?,
            CO64 => parser::parse_co64_data.parse(stream(input))?,
            _ => unreachable!(),
        })
    }
}

impl SerializeAtom for ChunkOffsetAtom {
    fn atom_type(&self) -> FourCC {
        // Use the appropriate atom type based on is_64bit
        if self.is_64bit {
            CO64
        } else {
            STCO
        }
    }

    fn into_body_bytes(self) -> Vec<u8> {
        serializer::serialize_stco_co64_data(self)
    }
}

mod serializer {
    use crate::atom::{util::serializer::be_u32, ChunkOffsetAtom};

    pub fn serialize_stco_co64_data(atom: ChunkOffsetAtom) -> Vec<u8> {
        let mut data = Vec::new();

        data.push(atom.version);
        data.extend(atom.flags);
        data.extend(be_u32(
            atom.chunk_offsets
                .len()
                .try_into()
                .expect("chunk offsets length must fit in u32"),
        ));

        atom.chunk_offsets.0.into_iter().for_each(|offset| {
            if atom.is_64bit {
                data.extend(offset.to_be_bytes());
            } else {
                data.extend(be_u32(
                    offset.try_into().expect("chunk offset must fit in u32"),
                ))
            }
        });

        data
    }
}

mod parser {
    use winnow::{
        binary::{be_u32, be_u64},
        combinator::{empty, repeat, seq, trace},
        error::{ContextError, ErrMode, StrContext},
        ModalResult, Parser,
    };

    use super::{ChunkOffsetAtom, ChunkOffsets};
    use crate::atom::util::parser::{byte_array, version, Stream};

    pub fn parse_stco_data(input: &mut Stream<'_>) -> ModalResult<ChunkOffsetAtom> {
        parse_stco_co64_data_inner(false).parse_next(input)
    }

    pub fn parse_co64_data(input: &mut Stream<'_>) -> ModalResult<ChunkOffsetAtom> {
        parse_stco_co64_data_inner(true).parse_next(input)
    }

    fn parse_stco_co64_data_inner<'i>(
        is_64bit: bool,
    ) -> impl Parser<Stream<'i>, ChunkOffsetAtom, ErrMode<ContextError>> {
        trace(
            if is_64bit { "co64" } else { "stco" },
            move |input: &mut Stream<'_>| {
                seq!(ChunkOffsetAtom {
                    version: version,
                    flags: byte_array.context(StrContext::Label("flags")),
                    chunk_offsets: chunk_offsets(is_64bit)
                        .map(ChunkOffsets)
                        .context(StrContext::Label("chunk_offsets")),
                    is_64bit: empty.value(is_64bit),
                })
                .parse_next(input)
            },
        )
    }

    fn chunk_offsets<'i>(
        is_64bit: bool,
    ) -> impl Parser<Stream<'i>, Vec<u64>, ErrMode<ContextError>> {
        trace("chunk_offsets", move |input: &mut Stream<'_>| {
            let entry_count = be_u32.parse_next(input)?;
            repeat(entry_count as usize, chunk_offset(is_64bit)).parse_next(input)
        })
    }

    fn chunk_offset<'i>(is_64bit: bool) -> impl Parser<Stream<'i>, u64, ErrMode<ContextError>> {
        trace("chunk_offset", move |input: &mut Stream<'_>| {
            if is_64bit {
                be_u64.parse_next(input)
            } else {
                be_u32.map(|v| v as u64).parse_next(input)
            }
        })
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::atom::test_utils::test_atom_roundtrip;

    /// Test round-trip for all available stco/co64 test data files
    #[test]
    fn test_stco_co64_roundtrip() {
        test_atom_roundtrip::<ChunkOffsetAtom>(STCO);
        test_atom_roundtrip::<ChunkOffsetAtom>(CO64);
    }
}