1use std::cmp::Reverse;
19use std::collections::{BTreeSet, BinaryHeap};
20
21use crate::{
22 CaptionCue, MAX_CLIPS_PER_TRACK, MAX_PROJECT_JSON_BYTES, MAX_TEXT_CONTENT_BYTES, MAX_TRACKS,
23};
24use alpha2_time::{MediaTime, TICKS_PER_SECOND, TimeRange};
25
26pub const MAX_CAPTION_FILE_BYTES: usize = MAX_PROJECT_JSON_BYTES;
30
31pub const MAX_CAPTION_FILE_CUES: usize = 65_536;
35
36#[derive(Clone, Debug, Default, PartialEq)]
38pub struct ImportedCaptions {
39 pub cues: Vec<CaptionCue>,
42 pub skipped: usize,
44}
45
46#[derive(Clone, Copy, Debug, Eq, PartialEq)]
48#[non_exhaustive]
49pub enum CaptionImportError {
50 MissingWebVttHeader,
52 FileTooLarge {
54 bytes: usize,
56 maximum: usize,
58 },
59 TooManyCues {
61 maximum: usize,
63 },
64 TooManyRows {
68 rows: usize,
70 maximum: usize,
72 },
73}
74
75impl std::fmt::Display for CaptionImportError {
76 fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
77 match self {
78 Self::MissingWebVttHeader => formatter.write_str("a WebVTT file begins with WEBVTT"),
79 Self::FileTooLarge { bytes, maximum } => write!(
80 formatter,
81 "the subtitle file is {bytes} bytes; at most {maximum} are read"
82 ),
83 Self::TooManyCues { maximum } => write!(
84 formatter,
85 "the subtitle file holds more than {maximum} cues"
86 ),
87 Self::TooManyRows { rows, maximum } => write!(
88 formatter,
89 "the captions need {rows} caption tracks; a project holds at most {maximum} tracks"
90 ),
91 }
92 }
93}
94
95impl std::error::Error for CaptionImportError {}
96
97const TICKS_PER_MILLISECOND: i64 = TICKS_PER_SECOND / 1_000;
98
99pub fn parse_srt(text: &str) -> Result<ImportedCaptions, CaptionImportError> {
105 check_file_size(text)?;
106 let mut imported = ImportedCaptions::default();
107 for block in blocks(text) {
108 let mut lines = block.iter().copied();
109 let Some(mut first) = lines.next() else {
110 continue;
111 };
112 if !first.contains("-->") {
114 match lines.next() {
115 Some(line) if first.trim().chars().all(|c| c.is_ascii_digit()) => first = line,
116 _ => {
117 imported.skipped += 1;
118 continue;
119 }
120 }
121 }
122 let text = lines.map(strip_srt_markup).collect::<Vec<_>>().join("\n");
123 imported.read(timing(first), &text)?;
124 }
125 imported.cues.sort_by_key(|cue| cue.timeline_range.start());
126 Ok(imported)
127}
128
129pub fn parse_webvtt(text: &str) -> Result<ImportedCaptions, CaptionImportError> {
137 check_file_size(text)?;
138 let text = text.strip_prefix('\u{feff}').unwrap_or(text);
139 let header_ok = text
140 .strip_prefix("WEBVTT")
141 .is_some_and(|rest| rest.is_empty() || rest.starts_with([' ', '\t', '\n', '\r']));
142 if !header_ok {
143 return Err(CaptionImportError::MissingWebVttHeader);
144 }
145 let mut imported = ImportedCaptions::default();
146 let mut blocks = blocks(text).into_iter();
147 let header = blocks.next().unwrap_or_default();
150 let cue_under_header = header
151 .iter()
152 .position(|line| line.contains("-->"))
153 .map(|start| header[start..].to_vec());
154 for block in cue_under_header.into_iter().chain(blocks) {
155 let first = block[0];
156 if ["NOTE", "STYLE", "REGION"].iter().any(|keyword| {
157 first
158 .strip_prefix(keyword)
159 .is_some_and(|rest| rest.is_empty() || rest.starts_with([' ', '\t']))
160 }) {
161 continue;
162 }
163 let (timing_line, body) = if first.contains("-->") {
165 (first, &block[1..])
166 } else if block.len() > 1 && block[1].contains("-->") {
167 (block[1], &block[2..])
168 } else {
169 imported.skipped += 1;
170 continue;
171 };
172 let text = body
173 .iter()
174 .map(|line| decode_references(&strip_tags(line)))
175 .collect::<Vec<_>>()
176 .join("\n");
177 imported.read(timing(timing_line), &text)?;
178 }
179 imported.cues.sort_by_key(|cue| cue.timeline_range.start());
180 Ok(imported)
181}
182
183impl ImportedCaptions {
184 fn read(&mut self, range: Option<TimeRange>, text: &str) -> Result<(), CaptionImportError> {
187 match (range, cue_text(text)) {
188 (Some(range), Some(text)) => {
189 if self.cues.len() == MAX_CAPTION_FILE_CUES {
190 return Err(CaptionImportError::TooManyCues {
191 maximum: MAX_CAPTION_FILE_CUES,
192 });
193 }
194 self.cues.push(CaptionCue {
195 timeline_range: range,
196 text,
197 });
198 }
199 _ => self.skipped += 1,
200 }
201 Ok(())
202 }
203}
204
205fn check_file_size(text: &str) -> Result<(), CaptionImportError> {
206 if text.len() > MAX_CAPTION_FILE_BYTES {
207 return Err(CaptionImportError::FileTooLarge {
208 bytes: text.len(),
209 maximum: MAX_CAPTION_FILE_BYTES,
210 });
211 }
212 Ok(())
213}
214
215pub fn caption_rows(cues: &[CaptionCue]) -> Result<Vec<Vec<CaptionCue>>, CaptionImportError> {
225 let mut sorted = cues.to_vec();
226 sorted.sort_by_key(|cue| cue.timeline_range.start());
227 let mut rows: Vec<Vec<CaptionCue>> = Vec::new();
228 let mut ended = BTreeSet::<usize>::new();
231 let mut ending = BinaryHeap::<Reverse<(MediaTime, usize)>>::new();
232 for cue in sorted {
233 let start = cue.timeline_range.start();
234 while let Some(&Reverse((end, row))) = ending.peek()
235 && end <= start
236 {
237 ending.pop();
238 ended.insert(row);
239 }
240 let row = ended.pop_first().unwrap_or_else(|| {
241 rows.push(Vec::new());
242 rows.len() - 1
243 });
244 let end = cue.timeline_range.end();
245 rows[row].push(cue);
246 if rows[row].len() < MAX_CLIPS_PER_TRACK {
247 ending.push(Reverse((end, row)));
248 }
249 }
250 if rows.len() > MAX_TRACKS {
251 return Err(CaptionImportError::TooManyRows {
252 rows: rows.len(),
253 maximum: MAX_TRACKS,
254 });
255 }
256 Ok(rows)
257}
258
259fn blocks(text: &str) -> Vec<Vec<&str>> {
261 let text = text.strip_prefix('\u{feff}').unwrap_or(text);
262 let mut blocks = Vec::new();
263 let mut current = Vec::new();
264 for line in lines(text) {
265 if line.trim().is_empty() {
266 if !current.is_empty() {
267 blocks.push(std::mem::take(&mut current));
268 }
269 } else {
270 current.push(line);
271 }
272 }
273 if !current.is_empty() {
274 blocks.push(current);
275 }
276 blocks
277}
278
279fn lines(text: &str) -> Vec<&str> {
283 let bytes = text.as_bytes();
284 let mut lines = Vec::new();
285 let (mut start, mut at) = (0, 0);
286 while at < bytes.len() {
287 match bytes[at] {
288 b'\n' => {
289 lines.push(&text[start..at]);
290 at += 1;
291 start = at;
292 }
293 b'\r' => {
294 let returns = bytes[at..]
295 .iter()
296 .take_while(|&&byte| byte == b'\r')
297 .count();
298 lines.push(&text[start..at]);
299 at += returns;
300 if bytes.get(at) == Some(&b'\n') {
301 at += 1;
302 } else {
303 lines.extend(std::iter::repeat_n("", returns - 1));
304 }
305 start = at;
306 }
307 _ => at += 1,
308 }
309 }
310 lines.push(&text[start..]);
311 lines
312}
313
314fn timing(line: &str) -> Option<TimeRange> {
316 let (start, rest) = line.split_once("-->")?;
317 let end = rest.split_whitespace().next()?;
318 let (start, end) = (timestamp(start.trim())?, timestamp(end)?);
319 TimeRange::new(start, end)
320 .ok()
321 .filter(|range| !range.is_empty())
322}
323
324fn timestamp(text: &str) -> Option<MediaTime> {
328 let (clock, millis) = text.split_once([',', '.'])?;
329 if millis.len() != 3 || !millis.bytes().all(|byte| byte.is_ascii_digit()) {
330 return None;
331 }
332 let fields = clock.split(':').collect::<Vec<_>>();
333 let number = |field: &str| {
334 (!field.is_empty() && field.bytes().all(|byte| byte.is_ascii_digit()))
335 .then(|| field.parse::<i64>().ok())
336 .flatten()
337 };
338 let (hours, minutes, seconds) = match fields.as_slice() {
339 [hours, minutes, seconds] => (number(hours)?, number(minutes)?, number(seconds)?),
340 [minutes, seconds] => (0, number(minutes)?, number(seconds)?),
341 _ => return None,
342 };
343 if minutes > 59 || seconds > 59 {
344 return None;
345 }
346 let milliseconds = ((hours.checked_mul(60)?.checked_add(minutes)?)
347 .checked_mul(60)?
348 .checked_add(seconds)?)
349 .checked_mul(1_000)?
350 .checked_add(millis.parse::<i64>().ok()?)?;
351 let ticks = milliseconds.checked_mul(TICKS_PER_MILLISECOND)?;
352 (ticks <= alpha2_time::MAX_SERIALIZED_TICKS).then(|| MediaTime::from_ticks(ticks))
353}
354
355fn cue_text(text: &str) -> Option<String> {
358 let text = text.trim_matches('\n').to_owned();
359 (!text.trim().is_empty() && text.len() <= MAX_TEXT_CONTENT_BYTES).then_some(text)
360}
361
362fn strip_srt_markup(line: &str) -> String {
364 let mut output = String::with_capacity(line.len());
365 let mut rest = line;
366 while let Some(start) = rest.find("{\\") {
367 output.push_str(&rest[..start]);
368 match rest[start..].find('}') {
369 Some(end) => rest = &rest[start + end + 1..],
370 None => {
371 rest = "";
372 }
373 }
374 }
375 output.push_str(rest);
376 strip_srt_tags(&output)
377}
378
379fn strip_srt_tags(line: &str) -> String {
383 let mut output = String::with_capacity(line.len());
384 let mut rest = line;
385 while let Some(start) = rest.find('<') {
386 output.push_str(&rest[..start]);
387 let tag = &rest[start + 1..];
388 let name = tag.strip_prefix('/').unwrap_or(tag);
389 let end = tag.find(['<', '>']);
392 match end {
393 Some(end)
394 if tag.as_bytes()[end] == b'>'
395 && name.starts_with(|character: char| character.is_ascii_alphabetic()) =>
396 {
397 rest = &tag[end + 1..];
398 }
399 _ => {
400 output.push('<');
401 rest = tag;
402 }
403 }
404 }
405 output.push_str(rest);
406 output
407}
408
409fn strip_tags(line: &str) -> String {
411 let mut output = String::with_capacity(line.len());
412 let mut depth = false;
413 for character in line.chars() {
414 match character {
415 '<' => depth = true,
416 '>' if depth => depth = false,
417 _ if !depth => output.push(character),
418 _ => {}
419 }
420 }
421 output
422}
423
424fn decode_references(text: &str) -> String {
431 let mut output = String::with_capacity(text.len());
432 let mut rest = text;
433 while let Some(start) = rest.find('&') {
434 output.push_str(&rest[..start]);
435 rest = &rest[start + 1..];
436 let length = rest
437 .bytes()
438 .take_while(|byte| byte.is_ascii_alphanumeric() || matches!(byte, b'#' | b'+'))
439 .count();
440 let decoded = (rest.as_bytes().get(length) == Some(&b';'))
441 .then(|| reference(&rest[..length]))
442 .flatten();
443 match decoded {
444 Some(character) => {
445 output.push(character);
446 rest = &rest[length + 1..];
447 }
448 None => output.push('&'),
449 }
450 }
451 output.push_str(rest);
452 output
453}
454
455fn reference(name: &str) -> Option<char> {
457 match name {
458 "amp" => Some('&'),
459 "lt" => Some('<'),
460 "gt" => Some('>'),
461 "nbsp" => Some('\u{a0}'),
462 "lrm" => Some('\u{200e}'),
463 "rlm" => Some('\u{200f}'),
464 _ => name
465 .strip_prefix("#x")
466 .or_else(|| name.strip_prefix("#X"))
467 .and_then(|hex| u32::from_str_radix(hex, 16).ok())
468 .or_else(|| name.strip_prefix('#').and_then(|dec| dec.parse().ok()))
469 .and_then(char::from_u32),
470 }
471}
472
473#[cfg(test)]
474#[path = "caption_import_tests.rs"]
475mod tests;