1pub mod transaction;
6
7use crate::error::{Result, UserError};
8#[cfg(test)]
9use crate::providers::Segment;
10use crate::providers::TranscriptionResult;
11use std::io::{self, Write};
12use std::path::Path;
13
14pub use transaction::{commit_text, CommitMode, OutputTransaction, SymlinkPolicy};
15
16pub const DEFAULT_MAX_OUTPUT_BYTES: usize = 32 * 1024 * 1024;
18
19#[derive(Debug, Clone, Copy, PartialEq, Eq)]
21pub enum OutputFormat {
22 Txt,
23 Srt,
24 Json,
25}
26
27impl OutputFormat {
28 pub fn parse(s: &str) -> Result<Self> {
29 match s.to_ascii_lowercase().as_str() {
30 "txt" | "text" => Ok(Self::Txt),
31 "srt" => Ok(Self::Srt),
32 "json" => Ok(Self::Json),
33 other => Err(UserError::InvalidOutputFormat {
34 format: other.to_string(),
35 }
36 .into()),
37 }
38 }
39
40 pub fn as_str(self) -> &'static str {
41 match self {
42 Self::Txt => "txt",
43 Self::Srt => "srt",
44 Self::Json => "json",
45 }
46 }
47
48 pub fn default_extension(self) -> &'static str {
49 self.as_str()
50 }
51}
52
53struct BudgetWriter<'a, W: Write> {
55 inner: &'a mut W,
56 written: usize,
57 max: usize,
58}
59
60impl<'a, W: Write> BudgetWriter<'a, W> {
61 fn new(inner: &'a mut W, max: usize) -> Self {
62 Self {
63 inner,
64 written: 0,
65 max: max.max(1),
66 }
67 }
68}
69
70impl<W: Write> Write for BudgetWriter<'_, W> {
71 fn write(&mut self, buf: &[u8]) -> io::Result<usize> {
72 if self.written.saturating_add(buf.len()) > self.max {
73 return Err(io::Error::other(format!(
74 "output exceeds maximum of {} bytes (would write {})",
75 self.max,
76 self.written.saturating_add(buf.len())
77 )));
78 }
79 let n = self.inner.write(buf)?;
80 self.written = self.written.saturating_add(n);
81 Ok(n)
82 }
83
84 fn flush(&mut self) -> io::Result<()> {
85 self.inner.flush()
86 }
87}
88
89pub fn format_result(result: &TranscriptionResult, format: OutputFormat) -> Result<String> {
93 format_result_with_limit(result, format, DEFAULT_MAX_OUTPUT_BYTES)
94}
95
96pub fn format_result_with_limit(
98 result: &TranscriptionResult,
99 format: OutputFormat,
100 max_bytes: usize,
101) -> Result<String> {
102 let mut buf = Vec::new();
103 write_result_with_limit(result, format, &mut buf, max_bytes)?;
104 match format {
109 OutputFormat::Txt => {
110 if buf.last() == Some(&b'\n') {
111 buf.pop();
112 }
113 String::from_utf8(buf).map_err(|e| {
114 crate::error::TranscriptionError::internal(format!("utf-8 output: {e}"))
115 })
116 }
117 OutputFormat::Srt | OutputFormat::Json => {
118 if buf.last() == Some(&b'\n') && format == OutputFormat::Json {
120 }
123 String::from_utf8(buf).map_err(|e| {
124 crate::error::TranscriptionError::internal(format!("utf-8 output: {e}"))
125 })
126 }
127 }
128}
129
130pub fn write_result<W: Write>(
132 result: &TranscriptionResult,
133 format: OutputFormat,
134 writer: W,
135) -> Result<()> {
136 write_result_with_limit(result, format, writer, DEFAULT_MAX_OUTPUT_BYTES)
137}
138
139pub fn write_result_with_limit<W: Write>(
141 result: &TranscriptionResult,
142 format: OutputFormat,
143 mut writer: W,
144 max_bytes: usize,
145) -> Result<()> {
146 let mut budget = BudgetWriter::new(&mut writer, max_bytes);
147 match format {
148 OutputFormat::Txt => write_txt(result, &mut budget),
149 OutputFormat::Srt => write_srt(result, &mut budget),
150 OutputFormat::Json => write_json(result, &mut budget),
151 }
152 .map_err(map_write_err)?;
153 let _ = budget.written;
155 Ok(())
156}
157
158fn map_write_err(e: io::Error) -> crate::error::TranscriptionError {
159 let msg = e.to_string();
160 if msg.contains("output exceeds maximum") {
161 return UserError::Other {
162 message: format!(
163 "output too large: {msg}\n Hint: raise the output byte budget or split the transcript."
164 ),
165 }
166 .into();
167 }
168 if e.kind() == io::ErrorKind::BrokenPipe {
170 return crate::error::EnvironmentError::Other {
171 message: "broken pipe while writing output".into(),
172 }
173 .into();
174 }
175 crate::error::EnvironmentError::Io(e).into()
176}
177
178pub fn write_result_to_path(
180 result: &TranscriptionResult,
181 format: OutputFormat,
182 path: &Path,
183 mode: CommitMode,
184) -> Result<()> {
185 write_result_to_path_with_limit(result, format, path, mode, DEFAULT_MAX_OUTPUT_BYTES)
186}
187
188pub fn write_result_to_path_with_limit(
190 result: &TranscriptionResult,
191 format: OutputFormat,
192 path: &Path,
193 mode: CommitMode,
194 max_bytes: usize,
195) -> Result<()> {
196 OutputTransaction::new(path, mode).commit_with(|tmp| {
197 let mut file = std::fs::OpenOptions::new()
198 .write(true)
199 .open(tmp)
200 .map_err(crate::error::EnvironmentError::Io)?;
201 write_result_with_limit(result, format, &mut file, max_bytes)?;
202 file.flush().map_err(crate::error::EnvironmentError::Io)?;
203 file.sync_all()
204 .map_err(crate::error::EnvironmentError::Io)?;
205 Ok(())
206 })
207}
208
209fn write_txt<W: Write>(result: &TranscriptionResult, w: &mut W) -> io::Result<()> {
210 let text = result.text.trim();
211 w.write_all(text.as_bytes())?;
212 w.write_all(b"\n")?;
213 Ok(())
214}
215
216fn write_srt<W: Write>(result: &TranscriptionResult, w: &mut W) -> io::Result<()> {
217 if result.segments.is_empty() {
218 if result.text.trim().is_empty() {
219 return Ok(());
220 }
221 let end = if result.duration_secs > 0.0 {
222 result.duration_secs
223 } else {
224 1.0
225 };
226 if !end.is_finite() || result.duration_secs < 0.0 {
227 return Err(io::Error::new(
228 io::ErrorKind::InvalidData,
229 "invalid transcript duration for SRT",
230 )); }
232 write!(
233 w,
234 "1\n{} --> {}\n{}\n",
235 format_ts(0.0),
236 format_ts(end),
237 result.text.trim().replace(['\r', '\n'], " ")
238 )?;
239 return Ok(());
240 }
241
242 let mut cue = 1usize;
243 for seg in &result.segments {
244 if !seg.start().is_finite() || !seg.end().is_finite() || seg.end() < seg.start() {
245 return Err(io::Error::new(
246 io::ErrorKind::InvalidData,
247 format!(
248 "invalid SRT timestamps at cue {cue}: start={} end={}",
249 seg.start(),
250 seg.end()
251 ),
252 ));
253 }
254 let text = seg.text().trim();
255 if text.is_empty() {
256 continue;
257 }
258 let text = text.replace(['\r', '\n'], " ");
259 write!(
260 w,
261 "{}\n{} --> {}\n{}\n\n",
262 cue,
263 format_ts(seg.start()),
264 format_ts(seg.end()),
265 text
266 )?;
267 cue += 1;
268 }
269 Ok(())
270}
271
272fn format_ts(secs: f64) -> String {
274 let total_ms = (secs.max(0.0) * 1000.0).round() as u64;
275 let ms = total_ms % 1000;
276 let total_secs = total_ms / 1000;
277 let s = total_secs % 60;
278 let total_mins = total_secs / 60;
279 let m = total_mins % 60;
280 let h = total_mins / 60;
281 format!("{h:02}:{m:02}:{s:02},{ms:03}")
282}
283
284fn write_json<W: Write>(result: &TranscriptionResult, w: &mut W) -> io::Result<()> {
285 let payload = crate::dto::SttResultDto::from_result(result);
287 serde_json::to_writer_pretty(&mut *w, &payload).map_err(io::Error::other)?;
288 w.write_all(b"\n")?;
289 Ok(())
290}
291
292#[cfg(test)]
293mod tests {
294 use super::*;
295
296 fn sample_result() -> TranscriptionResult {
297 TranscriptionResult::local(
298 "Hello world. This is a test.".into(),
299 vec![
300 Segment::from_parts_unchecked(0.0, 1.2, "Hello world.".to_string()),
301 Segment::from_parts_unchecked(1.4, 3.5, "This is a test.".to_string()),
302 ],
303 Some("en".into()),
304 "base".into(),
305 3.5,
306 )
307 }
308
309 #[test]
310 fn txt_is_plain() {
311 let s = format_result(&sample_result(), OutputFormat::Txt).unwrap();
312 assert_eq!(s, "Hello world. This is a test.");
313 }
314
315 #[test]
316 fn srt_has_cues() {
317 let s = format_result(&sample_result(), OutputFormat::Srt).unwrap();
318 assert!(s.contains("1\n"));
319 assert!(s.contains("00:00:00,000 --> 00:00:01,200"));
320 assert!(s.contains("Hello world."));
321 assert!(s.contains("2\n"));
322 assert!(s.contains("This is a test."));
323 }
324
325 #[test]
326 fn json_roundtrip_fields() {
327 let s = format_result(&sample_result(), OutputFormat::Json).unwrap();
328 let v: serde_json::Value = serde_json::from_str(&s).unwrap();
329 assert_eq!(v["text"], "Hello world. This is a test.");
330 assert_eq!(v["model"], "base");
331 assert_eq!(v["provider"], "local");
332 assert_eq!(v["language"], "en");
333 assert_eq!(v["timestamps_reliable"], true);
334 assert_eq!(v["backend_kind"], "asr");
335 assert_eq!(v["cleanup_style"], "raw");
336 assert_eq!(v["schema_version"], 1);
337 assert!(v.get("cleanup_provider").is_none() || v["cleanup_provider"].is_null());
338 assert!(v.get("original_text").is_none() || v["original_text"].is_null());
339 assert!(v["segments"].as_array().unwrap().len() == 2);
340 }
341
342 #[test]
343 fn json_includes_cleanup_metadata() {
344 let mut r = sample_result();
345 r.cleanup_style = crate::cleanup::CleanupStyle::Clean;
346 r.cleanup_provider = Some(crate::cleanup::CleanupProviderKind::Rules);
347 r.original_text = Some("um hello".into());
348 r.text = "Hello.".into();
349 let s = format_result(&r, OutputFormat::Json).unwrap();
350 let v: serde_json::Value = serde_json::from_str(&s).unwrap();
351 assert_eq!(v["cleanup_style"], "clean");
352 assert_eq!(v["cleanup_provider"], "rules");
353 assert_eq!(v["original_text"], "um hello");
354 assert_eq!(v["text"], "Hello.");
355 }
356
357 #[test]
358 fn parse_formats() {
359 assert_eq!(OutputFormat::parse("txt").unwrap(), OutputFormat::Txt);
360 assert_eq!(OutputFormat::parse("SRT").unwrap(), OutputFormat::Srt);
361 assert_eq!(OutputFormat::parse("json").unwrap(), OutputFormat::Json);
362 assert!(OutputFormat::parse("docx").is_err());
363 }
364
365 #[test]
366 fn format_ts_values() {
367 assert_eq!(format_ts(0.0), "00:00:00,000");
368 assert_eq!(format_ts(61.5), "00:01:01,500");
369 assert_eq!(format_ts(3661.001), "01:01:01,001");
370 }
371
372 #[test]
373 fn stream_matches_format_for_srt() {
374 let r = sample_result();
375 let mut stream = Vec::new();
376 write_result(&r, OutputFormat::Srt, &mut stream).unwrap();
377 let via_format = format_result(&r, OutputFormat::Srt).unwrap();
378 let stream_s = String::from_utf8(stream).unwrap();
380 assert!(stream_s.contains("Hello world."));
381 assert!(via_format.contains("Hello world."));
382 assert_eq!(
383 stream_s.replace("\r\n", "\n").trim(),
384 via_format.replace("\r\n", "\n").trim()
385 );
386 }
387
388 #[test]
389 fn budget_rejects_large_output() {
390 let r = sample_result();
391 let err = format_result_with_limit(&r, OutputFormat::Json, 32).unwrap_err();
392 assert!(err.to_string().contains("too large") || err.to_string().contains("maximum"));
393 }
394
395 #[test]
396 fn invalid_srt_timestamps_rejected() {
397 let mut r = sample_result();
398 r.segments[0].set_end(f64::NAN);
399 let err = write_result(&r, OutputFormat::Srt, Vec::new()).unwrap_err();
400 assert!(err.to_string().contains("timestamp") || err.to_string().contains("invalid"));
401 }
402
403 #[test]
404 fn large_segment_set_streams() {
405 let segs: Vec<Segment> = (0..5000)
406 .map(|i| Segment::from_parts_unchecked(i as f64, i as f64 + 0.5, format!("cue {i}")))
407 .collect();
408 let r = TranscriptionResult::local(
409 "many".into(),
410 segs,
411 Some("en".into()),
412 "base".into(),
413 5000.0,
414 );
415 let mut out = Vec::new();
416 write_result_with_limit(&r, OutputFormat::Srt, &mut out, 8 * 1024 * 1024).unwrap();
417 assert!(out.len() > 10_000);
418 assert!(out
419 .windows(5)
420 .any(|w| w == b"cue 0" || w.starts_with(b"cue ")));
421 }
422}