1use crate::error::{Result, UserError};
4use crate::providers::{Segment, TranscriptionResult};
5use serde::Serialize;
6use std::io::Write;
7
8#[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
42pub 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
51pub 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 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 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
109fn 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 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}