crosssubtitle-ai/src-tauri/src/subtitle.rs

use anyhow::{anyhow, Result};

use crate::models::SubtitleSegment;

#[derive(Debug, Clone, Copy)]
pub enum SubtitleFormat {
    Srt,
    Vtt,
    Ass,
}

impl SubtitleFormat {
    pub fn extension(&self) -> &'static str {
        match self {
            Self::Srt => "srt",
            Self::Vtt => "vtt",
            Self::Ass => "ass",
        }
    }
}

impl TryFrom<&str> for SubtitleFormat {
    type Error = anyhow::Error;

    fn try_from(value: &str) -> Result<Self> {
        match value.to_lowercase().as_str() {
            "srt" => Ok(Self::Srt),
            "vtt" => Ok(Self::Vtt),
            "ass" => Ok(Self::Ass),
            _ => Err(anyhow!("unsupported subtitle format: {value}")),
        }
    }
}

#[derive(Debug, Clone)]
pub struct AssStyle {
    pub name: String,
    pub font_name: String,
    pub font_size: u32,
    pub primary_colour: String,
    pub outline_colour: String,
}

impl Default for AssStyle {
    fn default() -> Self {
        Self {
            name: "Default".to_string(),
            font_name: "Arial".to_string(),
            font_size: 22,
            primary_colour: "&H00FFFFFF".to_string(),
            outline_colour: "&H00000000".to_string(),
        }
    }
}

pub fn render(segments: &[SubtitleSegment], format: SubtitleFormat, bilingual: bool) -> String {
    match format {
        SubtitleFormat::Srt => render_srt(segments, bilingual),
        SubtitleFormat::Vtt => render_vtt(segments, bilingual),
        SubtitleFormat::Ass => render_ass(segments, AssStyle::default(), bilingual),
    }
}

fn render_srt(segments: &[SubtitleSegment], bilingual: bool) -> String {
    segments
        .iter()
        .enumerate()
        .map(|(index, segment)| {
            format!(
                "{}\n{} --> {}\n{}\n",
                index + 1,
                format_timestamp(segment.start, ","),
                format_timestamp(segment.end, ","),
                compose_subtitle_text(segment, bilingual)
            )
        })
        .collect::<Vec<_>>()
        .join("\n")
}

fn render_vtt(segments: &[SubtitleSegment], bilingual: bool) -> String {
    let body = segments
        .iter()
        .map(|segment| {
            format!(
                "{} --> {}\n{}\n",
                format_timestamp(segment.start, "."),
                format_timestamp(segment.end, "."),
                compose_subtitle_text(segment, bilingual)
            )
        })
        .collect::<Vec<_>>()
        .join("\n");

    format!("WEBVTT\n\n{}", body)
}

fn render_ass(segments: &[SubtitleSegment], style: AssStyle, bilingual: bool) -> String {
    let header = format!(
        "[Script Info]\nScriptType: v4.00+\nCollisions: Normal\nPlayResX: 1280\nPlayResY: 720\n\n[V4+ Styles]\nFormat: Name, Fontname, Fontsize, PrimaryColour, SecondaryColour, OutlineColour, BackColour, Bold, Italic, Underline, StrikeOut, ScaleX, ScaleY, Spacing, Angle, BorderStyle, Outline, Shadow, Alignment, MarginL, MarginR, MarginV, Encoding\nStyle: {},{},{},{},&H000000FF,{},&H64000000,0,0,0,0,100,100,0,0,1,2,1,2,32,32,24,1\n\n[Events]\nFormat: Layer, Start, End, Style, Name, MarginL, MarginR, MarginV, Effect, Text\n",
        style.name, style.font_name, style.font_size, style.primary_colour, style.outline_colour
    );

    let body = segments
        .iter()
        .map(|segment| {
            let text = compose_subtitle_text(segment, bilingual).replace('\n', "\\N");
            format!(
                "Dialogue: 0,{}, {},{},{},0,0,0,,{}",
                format_ass_timestamp(segment.start),
                format_ass_timestamp(segment.end),
                style.name,
                "",
                text
            )
        })
        .collect::<Vec<_>>()
        .join("\n");

    format!("{header}{body}\n")
}

fn compose_subtitle_text(segment: &SubtitleSegment, bilingual: bool) -> String {
    match (bilingual, segment.translated_text.as_deref()) {
        (true, Some(translated)) if !translated.trim().is_empty() => {
            format!("{}\n{}", segment.source_text.trim(), translated.trim())
        }
        (_, Some(translated)) if !translated.trim().is_empty() => translated.trim().to_string(),
        _ => segment.source_text.trim().to_string(),
    }
}

fn format_timestamp(seconds: f32, separator: &str) -> String {
    let millis = (seconds * 1000.0).round() as u64;
    let hours = millis / 3_600_000;
    let minutes = (millis % 3_600_000) / 60_000;
    let secs = (millis % 60_000) / 1_000;
    let ms = millis % 1_000;
    format!("{hours:02}:{minutes:02}:{secs:02}{separator}{ms:03}")
}

fn format_ass_timestamp(seconds: f32) -> String {
    let centis = (seconds * 100.0).round() as u64;
    let hours = centis / 360_000;
    let minutes = (centis % 360_000) / 6_000;
    let secs = (centis % 6_000) / 100;
    let cs = centis % 100;
    format!("{hours}:{minutes:02}:{secs:02}.{cs:02}")
}