Skip to main content

aurum_core/output/
mod.rs

1//! Output formatters: txt, srt, json.
2
3use crate::error::{Result, UserError};
4use crate::providers::{Segment, TranscriptionResult};
5use serde::Serialize;
6use std::io::Write;
7
8/// Supported output formats.
9#[derive(Debug, Clone, Copy, PartialEq, Eq)]
10pub enum OutputFormat {
11    Txt,
12    Srt,
13    Json,
14}
15
16impl OutputFormat {
17    pub fn parse(s: &str) -> Result<Self> {
18        match s.to_ascii_lowercase().as_str() {
19            "txt" | "text" => Ok(Self::Txt),
20            "srt" => Ok(Self::Srt),
21            "json" => Ok(Self::Json),
22            other => Err(UserError::InvalidOutputFormat {
23                format: other.to_string(),
24            }
25            .into()),
26        }
27    }
28
29    pub fn as_str(self) -> &'static str {
30        match self {
31            Self::Txt => "txt",
32            Self::Srt => "srt",
33            Self::Json => "json",
34        }
35    }
36
37    pub fn default_extension(self) -> &'static str {
38        self.as_str()
39    }
40}
41
42/// Format a transcription result into the requested representation.
43pub fn format_result(result: &TranscriptionResult, format: OutputFormat) -> Result<String> {
44    match format {
45        OutputFormat::Txt => Ok(format_txt(result)),
46        OutputFormat::Srt => Ok(format_srt(result)),
47        OutputFormat::Json => format_json(result),
48    }
49}
50
51/// Write formatted output to a writer.
52pub fn write_result<W: Write>(
53    result: &TranscriptionResult,
54    format: OutputFormat,
55    mut writer: W,
56) -> Result<()> {
57    let body = format_result(result, format)?;
58    writer.write_all(body.as_bytes())?;
59    if !body.ends_with('\n') {
60        writer.write_all(b"\n")?;
61    }
62    Ok(())
63}
64
65fn format_txt(result: &TranscriptionResult) -> String {
66    result.text.trim().to_string()
67}
68
69fn format_srt(result: &TranscriptionResult) -> String {
70    if result.segments.is_empty() {
71        // Single cue covering the whole transcript when no segments exist.
72        if result.text.trim().is_empty() {
73            return String::new();
74        }
75        let end = if result.duration_secs > 0.0 {
76            result.duration_secs
77        } else {
78            1.0
79        };
80        return format!(
81            "1\n{} --> {}\n{}\n",
82            format_ts(0.0),
83            format_ts(end),
84            result.text.trim()
85        );
86    }
87
88    let mut out = String::new();
89    let mut cue = 1usize;
90    for seg in &result.segments {
91        let text = seg.text.trim();
92        if text.is_empty() {
93            continue;
94        }
95        // SRT forbids unescaped newlines inside cue text — flatten.
96        let text = text.replace(['\r', '\n'], " ");
97        out.push_str(&format!(
98            "{}\n{} --> {}\n{}\n\n",
99            cue,
100            format_ts(seg.start),
101            format_ts(seg.end),
102            text
103        ));
104        cue += 1;
105    }
106    out
107}
108
109/// SRT timestamp: HH:MM:SS,mmm
110fn format_ts(secs: f64) -> String {
111    let total_ms = (secs.max(0.0) * 1000.0).round() as u64;
112    let ms = total_ms % 1000;
113    let total_secs = total_ms / 1000;
114    let s = total_secs % 60;
115    let total_mins = total_secs / 60;
116    let m = total_mins % 60;
117    let h = total_mins / 60;
118    format!("{h:02}:{m:02}:{s:02},{ms:03}")
119}
120
121#[derive(Serialize)]
122struct JsonOutput<'a> {
123    text: &'a str,
124    language: Option<&'a str>,
125    model: &'a str,
126    provider: &'a str,
127    duration_secs: f64,
128    backend_kind: crate::providers::BackendKind,
129    /// False for LLM-assisted backends — segment times are best-effort only.
130    timestamps_reliable: bool,
131    cleanup_style: crate::cleanup::CleanupStyle,
132    #[serde(skip_serializing_if = "Option::is_none")]
133    cleanup_provider: Option<crate::cleanup::CleanupProviderKind>,
134    #[serde(skip_serializing_if = "Option::is_none")]
135    original_text: Option<&'a str>,
136    segments: &'a [Segment],
137}
138
139fn format_json(result: &TranscriptionResult) -> Result<String> {
140    let payload = JsonOutput {
141        text: &result.text,
142        language: result.language.as_deref(),
143        model: &result.model,
144        provider: &result.provider,
145        duration_secs: result.duration_secs,
146        backend_kind: result.backend_kind,
147        timestamps_reliable: result.timestamps_reliable,
148        cleanup_style: result.cleanup_style,
149        cleanup_provider: result.cleanup_provider,
150        original_text: result.original_text.as_deref(),
151        segments: &result.segments,
152    };
153    serde_json::to_string_pretty(&payload).map_err(|e| {
154        crate::error::TranscriptionError::internal(format!("json serialize failed: {e}"))
155    })
156}
157
158#[cfg(test)]
159mod tests {
160    use super::*;
161
162    fn sample_result() -> TranscriptionResult {
163        TranscriptionResult::local(
164            "Hello world. This is a test.".into(),
165            vec![
166                Segment {
167                    start: 0.0,
168                    end: 1.2,
169                    text: "Hello world.".into(),
170                },
171                Segment {
172                    start: 1.4,
173                    end: 3.5,
174                    text: "This is a test.".into(),
175                },
176            ],
177            Some("en".into()),
178            "base".into(),
179            3.5,
180        )
181    }
182
183    #[test]
184    fn txt_is_plain() {
185        let s = format_result(&sample_result(), OutputFormat::Txt).unwrap();
186        assert_eq!(s, "Hello world. This is a test.");
187    }
188
189    #[test]
190    fn srt_has_cues() {
191        let s = format_result(&sample_result(), OutputFormat::Srt).unwrap();
192        assert!(s.contains("1\n"));
193        assert!(s.contains("00:00:00,000 --> 00:00:01,200"));
194        assert!(s.contains("Hello world."));
195        assert!(s.contains("2\n"));
196        assert!(s.contains("This is a test."));
197    }
198
199    #[test]
200    fn json_roundtrip_fields() {
201        let s = format_result(&sample_result(), OutputFormat::Json).unwrap();
202        let v: serde_json::Value = serde_json::from_str(&s).unwrap();
203        assert_eq!(v["text"], "Hello world. This is a test.");
204        assert_eq!(v["model"], "base");
205        assert_eq!(v["provider"], "local");
206        assert_eq!(v["language"], "en");
207        assert_eq!(v["timestamps_reliable"], true);
208        assert_eq!(v["backend_kind"], "asr");
209        assert_eq!(v["cleanup_style"], "raw");
210        assert!(v.get("cleanup_provider").is_none() || v["cleanup_provider"].is_null());
211        assert!(v.get("original_text").is_none() || v["original_text"].is_null());
212        assert!(v["segments"].as_array().unwrap().len() == 2);
213    }
214
215    #[test]
216    fn json_includes_cleanup_metadata() {
217        let mut r = sample_result();
218        r.cleanup_style = crate::cleanup::CleanupStyle::Clean;
219        r.cleanup_provider = Some(crate::cleanup::CleanupProviderKind::Rules);
220        r.original_text = Some("um hello".into());
221        r.text = "Hello.".into();
222        let s = format_result(&r, OutputFormat::Json).unwrap();
223        let v: serde_json::Value = serde_json::from_str(&s).unwrap();
224        assert_eq!(v["cleanup_style"], "clean");
225        assert_eq!(v["cleanup_provider"], "rules");
226        assert_eq!(v["original_text"], "um hello");
227        assert_eq!(v["text"], "Hello.");
228    }
229
230    #[test]
231    fn parse_formats() {
232        assert_eq!(OutputFormat::parse("txt").unwrap(), OutputFormat::Txt);
233        assert_eq!(OutputFormat::parse("SRT").unwrap(), OutputFormat::Srt);
234        assert_eq!(OutputFormat::parse("json").unwrap(), OutputFormat::Json);
235        assert!(OutputFormat::parse("docx").is_err());
236    }
237
238    #[test]
239    fn format_ts_values() {
240        assert_eq!(format_ts(0.0), "00:00:00,000");
241        assert_eq!(format_ts(61.5), "00:01:01,500");
242        assert_eq!(format_ts(3661.001), "01:01:01,001");
243    }
244}