1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
367
368
369
370
371
372
373
374
375
376
377
378
379
380
381
382
383
384
385
386
387
388
389
390
391
392
393
394
395
396
397
398
399
400
401
402
403
404
405
406
407
408
409
410
411
412
413
414
415
416
417
418
419
420
421
422
423
424
425
426
427
428
429
430
431
432
433
434
435
436
437
438
439
440
441
442
443
444
445
446
447
448
449
450
451
452
453
use super::r#ref::IlstRef;
use crate::error::{FileEncodingError, Result};
use crate::file::FileType;
use crate::macros::{err, try_vec};
use crate::mp4::atom_info::{AtomIdent, AtomInfo};
use crate::mp4::ilst::r#ref::AtomRef;
use crate::mp4::moov::Moov;
use crate::mp4::read::{atom_tree, meta_is_full, nested_atom, verify_mp4, AtomReader};
use crate::mp4::AtomData;
use crate::picture::{MimeType, Picture};

use std::fs::File;
use std::io::{Cursor, Read, Seek, SeekFrom, Write};

use byteorder::{BigEndian, WriteBytesExt};

pub(crate) fn write_to<'a, I: 'a>(data: &mut File, tag: &mut IlstRef<'a, I>) -> Result<()>
where
	I: IntoIterator<Item = &'a AtomData>,
{
	let mut reader = AtomReader::new(data)?;

	verify_mp4(&mut reader)?;

	let moov = Moov::find(&mut reader)?;
	let pos = reader.position()?;

	reader.rewind()?;

	let mut file_bytes = Vec::new();
	reader.read_to_end(&mut file_bytes)?;

	let mut cursor = Cursor::new(file_bytes);
	cursor.seek(SeekFrom::Start(pos))?;

	let ilst = build_ilst(&mut tag.atoms)?;
	let remove_tag = ilst.is_empty();

	let udta = nested_atom(&mut cursor, moov.len, b"udta")?;

	// Nothing to do
	if remove_tag && udta.is_none() {
		return Ok(());
	}

	// Total size of new atoms
	let mut new_udta_size;
	// Size of the existing udta atom
	let mut existing_udta_size = 0;

	// ilst is nested in udta.meta, so we need to check what atoms actually exist
	if let Some(udta) = udta {
		existing_udta_size = udta.len;
		new_udta_size = existing_udta_size;

		let meta = nested_atom(&mut cursor, udta.len, b"meta")?;
		match meta {
			Some(meta) => {
				// We may encounter a non-full `meta` atom
				meta_is_full(&mut cursor)?;

				// We can use the existing `udta` and `meta` atoms
				save_to_existing(
					&mut cursor,
					(meta, udta),
					&mut new_udta_size,
					ilst,
					remove_tag,
				)?
			},
			// Nothing to do
			None if remove_tag => return Ok(()),
			// We have to create the `meta` atom
			None => {
				existing_udta_size = udta.len;

				// `meta` (12) + `ilst`
				let capacity = 12 + ilst.len();
				let buf = Vec::with_capacity(capacity);

				let mut bytes = Cursor::new(buf);
				create_meta(&mut bytes, &ilst)?;

				let bytes = bytes.into_inner();

				new_udta_size = udta.len + bytes.len() as u64;

				cursor.seek(SeekFrom::Start(udta.start))?;
				write_size(udta.start, new_udta_size, udta.extended, &mut cursor)?;

				cursor
					.get_mut()
					.splice(udta.start as usize..udta.start as usize, bytes);
			},
		}
	} else {
		// We have to create the `udta` atom

		// `udta` + `meta` (12) + `hdlr` (33) + `ilst`
		let capacity = 53 + ilst.len();
		let buf = Vec::with_capacity(capacity);

		let mut bytes = Cursor::new(buf);
		bytes.write_all(&[0, 0, 0, 0, b'u', b'd', b't', b'a'])?;

		create_meta(&mut bytes, &ilst)?;

		// udta size
		bytes.rewind()?;
		write_size(0, ilst.len() as u64 + 8, false, &mut bytes)?;

		let bytes = bytes.into_inner();

		new_udta_size = bytes.len() as u64;

		let udta_pos = (moov.start + 8) as usize;
		cursor.get_mut().splice(udta_pos..udta_pos, bytes);
	}

	cursor.seek(SeekFrom::Start(moov.start))?;

	// Change the size of the moov atom
	write_size(
		moov.start,
		(moov.len - existing_udta_size) + new_udta_size,
		moov.extended,
		&mut cursor,
	)?;

	let data = reader.into_inner();

	data.rewind()?;
	data.set_len(0)?;
	data.write_all(&cursor.into_inner())?;

	Ok(())
}

fn save_to_existing(
	cursor: &mut Cursor<Vec<u8>>,
	(meta, udta): (AtomInfo, AtomInfo),
	new_udta_size: &mut u64,
	ilst: Vec<u8>,
	remove_tag: bool,
) -> Result<()> {
	let replacement;
	let range;

	let (ilst_idx, tree) = atom_tree(cursor, meta.len - 8, b"ilst")?;

	if tree.is_empty() {
		// Nothing to do
		if remove_tag {
			return Ok(());
		}

		let meta_end = (meta.start + meta.len) as usize;

		replacement = ilst;
		range = meta_end..meta_end;
	} else {
		let existing_ilst = &tree[ilst_idx];
		let existing_ilst_size = existing_ilst.len;

		let mut range_start = existing_ilst.start;
		let range_end = existing_ilst.start + existing_ilst_size;

		if remove_tag {
			// We just need to strip out the `ilst` atom

			replacement = Vec::new();
			range = range_start as usize..range_end as usize;
		} else {
			// Check for some padding atoms we can utilize
			let mut available_space = existing_ilst_size;

			// Check for one directly before the `ilst` atom
			if ilst_idx > 0 {
				let previous_atom = &tree[ilst_idx - 1];

				if previous_atom.ident == AtomIdent::Fourcc(*b"free") {
					range_start = previous_atom.start;
					available_space += previous_atom.len;
				}
			}

			// And after
			if ilst_idx != tree.len() - 1 {
				let next_atom = &tree[ilst_idx + 1];

				if next_atom.ident == AtomIdent::Fourcc(*b"free") {
					available_space += next_atom.len;
				}
			}

			let ilst_len = ilst.len() as u64;

			// Check if we have enough padding to fit the `ilst` atom and a new `free` atom
			if available_space > ilst_len && available_space - ilst_len > 8 {
				// We have enough space to make use of the padding

				let remaining_space = available_space - ilst_len;
				if remaining_space > u64::from(u32::MAX) {
					err!(TooMuchData);
				}

				let remaining_space = remaining_space as u32;

				cursor.seek(SeekFrom::Start(range_start))?;
				cursor.write_all(&ilst)?;

				// Write the remaining padding
				cursor.write_u32::<BigEndian>(remaining_space)?;
				cursor.write_all(b"free")?;
				cursor.write_all(&try_vec![1; (remaining_space - 8) as usize])?;

				return Ok(());
			}

			replacement = ilst;
			range = range_start as usize..range_end as usize;
		}
	}

	let new_meta_size = (meta.len - range.len() as u64) + replacement.len() as u64;

	// Replace the `ilst` atom
	cursor.get_mut().splice(range, replacement);

	if new_meta_size != meta.len {
		// We need to change the `meta` and `udta` atom sizes

		*new_udta_size = (udta.len - meta.len) + new_meta_size;

		cursor.seek(SeekFrom::Start(meta.start))?;
		write_size(meta.start, new_meta_size, meta.extended, cursor)?;

		cursor.seek(SeekFrom::Start(udta.start))?;
		write_size(udta.start, *new_udta_size, udta.extended, cursor)?;
	}

	Ok(())
}

fn create_meta(cursor: &mut Cursor<Vec<u8>>, ilst: &[u8]) -> Result<()> {
	const HDLR_SIZE: u64 = 33;

	let start = cursor.stream_position()?;
	// meta atom
	cursor.write_all(&[0, 0, 0, 0, b'm', b'e', b't', b'a', 0, 0, 0, 0])?;

	// hdlr atom
	cursor.write_u32::<BigEndian>(0)?;
	cursor.write_all(b"hdlr")?;
	cursor.write_u64::<BigEndian>(0)?;
	cursor.write_all(b"mdirappl")?;
	cursor.write_all(&[0, 0, 0, 0, 0, 0, 0, 0, 0])?;

	cursor.seek(SeekFrom::Start(start))?;

	let meta_size = 4 + HDLR_SIZE + ilst.len() as u64;
	write_size(start, meta_size, false, cursor)?;

	// Seek to `hdlr` size
	let hdlr_size_pos = cursor.seek(SeekFrom::Current(4))?;
	write_size(hdlr_size_pos, HDLR_SIZE, false, cursor)?;

	cursor.seek(SeekFrom::End(0))?;
	cursor.write_all(ilst)?;

	Ok(())
}

fn write_size(start: u64, size: u64, extended: bool, writer: &mut Cursor<Vec<u8>>) -> Result<()> {
	if size > u64::from(u32::MAX) {
		// 0001 (identifier) ????????
		writer.write_u32::<BigEndian>(1)?;
		// Skip identifier
		writer.seek(SeekFrom::Current(4))?;

		let extended_size = size.to_be_bytes();
		let inner = writer.get_mut();

		if extended {
			// Overwrite existing extended size
			writer.write_u64::<BigEndian>(size)?;
		} else {
			for i in extended_size {
				inner.insert((start + 8 + u64::from(i)) as usize, i);
			}

			writer.seek(SeekFrom::Current(8))?;
		}
	} else {
		// ???? (identifier)
		writer.write_u32::<BigEndian>(size as u32)?;
		writer.seek(SeekFrom::Current(4))?;
	}

	Ok(())
}

pub(super) fn build_ilst<'a, I: 'a>(
	atoms: &mut dyn Iterator<Item = AtomRef<'a, I>>,
) -> Result<Vec<u8>>
where
	I: IntoIterator<Item = &'a AtomData>,
{
	let mut peek = atoms.peekable();

	if peek.peek().is_none() {
		return Ok(Vec::new());
	}

	let mut writer = Cursor::new(vec![0, 0, 0, 0, b'i', b'l', b's', b't']);
	writer.seek(SeekFrom::End(0))?;

	for atom in peek {
		let start = writer.stream_position()?;

		// Empty size, we get it later
		writer.write_all(&[0; 4])?;

		match atom.ident {
			AtomIdent::Fourcc(ref fourcc) => writer.write_all(fourcc)?,
			AtomIdent::Freeform { mean, name } => write_freeform(&mean, &name, &mut writer)?,
		}

		write_atom_data(atom.data, &mut writer)?;

		let end = writer.stream_position()?;

		let size = end - start;

		writer.seek(SeekFrom::Start(start))?;

		write_size(start, size, false, &mut writer)?;

		writer.seek(SeekFrom::Start(end))?;
	}

	let size = writer.get_ref().len();

	writer.rewind()?;

	write_size(0, size as u64, false, &mut writer)?;

	Ok(writer.into_inner())
}

fn write_freeform(mean: &str, name: &str, writer: &mut Cursor<Vec<u8>>) -> Result<()> {
	// ---- : ???? : ????

	// ----
	writer.write_all(b"----")?;

	// .... MEAN 0000 ????
	writer.write_u32::<BigEndian>((12 + mean.len()) as u32)?;
	writer.write_all(&[b'm', b'e', b'a', b'n', 0, 0, 0, 0])?;
	writer.write_all(mean.as_bytes())?;

	// .... NAME 0000 ????
	writer.write_u32::<BigEndian>((12 + name.len()) as u32)?;
	writer.write_all(&[b'n', b'a', b'm', b'e', 0, 0, 0, 0])?;
	writer.write_all(name.as_bytes())?;

	Ok(())
}

fn write_atom_data<'a, I: 'a>(data: I, writer: &mut Cursor<Vec<u8>>) -> Result<()>
where
	I: IntoIterator<Item = &'a AtomData>,
{
	for value in data {
		match value {
			AtomData::UTF8(text) => write_data(1, text.as_bytes(), writer)?,
			AtomData::UTF16(text) => write_data(2, text.as_bytes(), writer)?,
			AtomData::Picture(ref pic) => write_picture(pic, writer)?,
			AtomData::SignedInteger(int) => write_signed_int(*int, writer)?,
			AtomData::UnsignedInteger(uint) => write_unsigned_int(*uint, writer)?,
			AtomData::Bool(b) => write_signed_int(i32::from(*b), writer)?,
			AtomData::Unknown { code, ref data } => write_data(*code, data, writer)?,
		};
	}

	Ok(())
}

fn write_signed_int(int: i32, writer: &mut Cursor<Vec<u8>>) -> Result<()> {
	let start_pos = int.trailing_zeros() as usize;
	write_int(21, int.to_be_bytes(), start_pos, writer)
}

fn write_unsigned_int(uint: u32, writer: &mut Cursor<Vec<u8>>) -> Result<()> {
	let start_pos = uint.trailing_zeros() as usize;
	write_int(22, uint.to_be_bytes(), start_pos, writer)
}

fn write_int(
	flags: u32,
	bytes: [u8; 4],
	mut start_pos: usize,
	writer: &mut Cursor<Vec<u8>>,
) -> Result<()> {
	if start_pos == 1 || start_pos == 4 {
		start_pos = 0;
	}
	write_data(flags, &bytes[start_pos..], writer)
}

fn write_picture(picture: &Picture, writer: &mut Cursor<Vec<u8>>) -> Result<()> {
	match picture.mime_type {
		// GIF is deprecated
		MimeType::Gif => write_data(12, &picture.data, writer),
		MimeType::Jpeg => write_data(13, &picture.data, writer),
		MimeType::Png => write_data(14, &picture.data, writer),
		MimeType::Bmp => write_data(27, &picture.data, writer),
		// We'll assume implicit (0) was the intended type
		MimeType::None => write_data(0, &picture.data, writer),
		_ => Err(FileEncodingError::new(
			FileType::MP4,
			"Attempted to write an unsupported picture format",
		)
		.into()),
	}
}

fn write_data(flags: u32, data: &[u8], writer: &mut Cursor<Vec<u8>>) -> Result<()> {
	if flags > 16_777_215 {
		return Err(FileEncodingError::new(
			FileType::MP4,
			"Attempted to write a code that cannot fit in 24 bits",
		)
		.into());
	}

	// .... DATA (version = 0) (flags) (locale = 0000) (data)
	let size = 16_u64 + data.len() as u64;

	writer.write_all(&[0, 0, 0, 0, b'd', b'a', b't', b'a'])?;
	write_size(writer.seek(SeekFrom::Current(-8))?, size, false, writer)?;

	// Version
	writer.write_u8(0)?;

	writer.write_uint::<BigEndian>(u64::from(flags), 3)?;

	// Locale
	writer.write_all(&[0; 4])?;
	writer.write_all(data)?;

	Ok(())
}