shinkai-translator 0.1.3

use std::collections::BTreeMap;

use crate::domain::{RenderPlan, SubtitleCue, SubtitleDocument, SubtitleFormat};
use crate::error::TranslatorError;

use super::{normalize_newlines, parse_arrow_timing_line, push_terminal_newline, split_blocks};

pub fn parse(source: &str) -> Result<SubtitleDocument, TranslatorError> {
    let normalized = normalize_newlines(source);
    let blocks = split_blocks(&normalized);
    let mut cues = Vec::new();

    for block in blocks {
        let lines: Vec<&str> = block.lines().collect();
        if lines.is_empty() {
            continue;
        }

        let (identifier, timing_index) = if lines[0].contains("-->") {
            (None, 0)
        } else {
            (Some(lines[0].trim().to_owned()), 1)
        };

        let timing_line = lines.get(timing_index).ok_or_else(|| {
            TranslatorError::Parse(format!("missing timing line in SRT block: {block}"))
        })?;
        let (start, end, settings) = parse_arrow_timing_line(timing_line)?;
        let text = if lines.len() > timing_index + 1 {
            lines[(timing_index + 1)..].join("\n")
        } else {
            String::new()
        };

        cues.push(SubtitleCue::new(
            format!("cue-{}", cues.len() + 1),
            identifier,
            start,
            end,
            settings,
            text,
            BTreeMap::new(),
        ));
    }

    Ok(SubtitleDocument::from_parts(
        SubtitleFormat::Srt,
        cues,
        RenderPlan::Srt,
    ))
}

pub fn render(document: &SubtitleDocument) -> Result<String, TranslatorError> {
    if document.format() != SubtitleFormat::Srt {
        return Err(TranslatorError::UnsupportedFormat(
            "document is not SRT".to_owned(),
        ));
    }

    let blocks = document
        .cues()
        .iter()
        .enumerate()
        .map(|(index, cue)| {
            let identifier = cue
                .identifier()
                .map(ToOwned::to_owned)
                .unwrap_or_else(|| (index + 1).to_string());
            let timing = match cue.settings() {
                Some(settings) => format!("{} --> {} {settings}", cue.start(), cue.end()),
                None => format!("{} --> {}", cue.start(), cue.end()),
            };

            if cue.text().is_empty() {
                format!("{identifier}\n{timing}")
            } else {
                format!("{identifier}\n{timing}\n{}", cue.text())
            }
        })
        .collect::<Vec<_>>();

    Ok(push_terminal_newline(blocks.join("\n\n")))
}

#[cfg(test)]
mod tests {
    use super::{parse, render};

    #[test]
    fn parses_and_renders_multiline_srt() {
        let source = "1\n00:00:01,000 --> 00:00:03,000\nhello\nworld\n\n2\n00:00:04,000 --> 00:00:05,000\nbye\n";
        let document = parse(source).expect("parse should succeed");

        assert_eq!(document.cue_count(), 2);
        assert_eq!(document.cues()[0].text(), "hello\nworld");

        let rendered = render(&document).expect("render should succeed");
        assert_eq!(rendered, source);
    }
}