1use crate::errors::{Result, TypecastError};
9use base64::{engine::general_purpose::STANDARD as B64, Engine};
10use serde::{Deserialize, Serialize};
11use std::fs;
12use std::path::Path;
13
14#[derive(Debug, Clone, Deserialize, Serialize, PartialEq)]
20pub struct AlignmentSegmentWord {
21 pub text: String,
23 pub start: f64,
25 pub end: f64,
27}
28
29#[derive(Debug, Clone, Deserialize, Serialize, PartialEq)]
31pub struct AlignmentSegmentCharacter {
32 pub text: String,
34 pub start: f64,
36 pub end: f64,
38}
39
40#[derive(Debug, Clone, Serialize)]
42pub struct TTSRequestWithTimestamps {
43 pub voice_id: String,
45 pub text: String,
47 pub model: crate::models::TTSModel,
49 #[serde(skip_serializing_if = "Option::is_none")]
51 pub language: Option<String>,
52 #[serde(skip_serializing_if = "Option::is_none")]
55 pub prompt: Option<serde_json::Value>,
56 #[serde(skip_serializing_if = "Option::is_none")]
59 pub output: Option<serde_json::Value>,
60 #[serde(skip_serializing_if = "Option::is_none")]
62 pub seed: Option<u32>,
63}
64
65impl TTSRequestWithTimestamps {
66 pub fn new(
68 voice_id: impl Into<String>,
69 text: impl Into<String>,
70 model: crate::models::TTSModel,
71 ) -> Self {
72 Self {
73 voice_id: voice_id.into(),
74 text: text.into(),
75 model,
76 language: None,
77 prompt: None,
78 output: None,
79 seed: None,
80 }
81 }
82
83 pub fn language(mut self, language: impl Into<String>) -> Self {
85 self.language = Some(language.into());
86 self
87 }
88
89 pub fn prompt(mut self, prompt: serde_json::Value) -> Self {
91 self.prompt = Some(prompt);
92 self
93 }
94
95 pub fn output(mut self, output: serde_json::Value) -> Self {
97 self.output = Some(output);
98 self
99 }
100
101 pub fn seed(mut self, seed: u32) -> Self {
103 self.seed = Some(seed);
104 self
105 }
106}
107
108#[derive(Debug, Clone, Deserialize, Serialize)]
118pub struct TTSWithTimestampsResponse {
119 pub audio: String,
121 pub audio_format: String,
123 pub audio_duration: f64,
125 pub words: Option<Vec<AlignmentSegmentWord>>,
128 pub characters: Option<Vec<AlignmentSegmentCharacter>>,
131}
132
133impl TTSWithTimestampsResponse {
134 pub fn audio_bytes(&self) -> Result<Vec<u8>> {
136 B64.decode(&self.audio)
137 .map_err(|e| TypecastError::DecodeError(e.to_string()))
138 }
139
140 pub fn save_audio<P: AsRef<Path>>(&self, path: P) -> Result<()> {
142 let bytes = self.audio_bytes()?;
143 fs::write(path, bytes).map_err(|e| TypecastError::IoError(e.to_string()))
144 }
145
146 pub fn to_srt(&self) -> Result<String> {
151 format_captions(self, true)
152 }
153
154 pub fn to_vtt(&self) -> Result<String> {
159 format_captions(self, false)
160 }
161}
162
163const MAX_CAPTION_SECONDS: f64 = 7.0;
168const MAX_CAPTION_CHARS: usize = 42;
169const SENTENCE_TERMINATORS: &[&str] = &[".", "?", "!", "\u{3002}", "\u{ff1f}", "\u{ff01}"];
170
171struct Segment {
172 text: String,
173 start: f64,
174 end: f64,
175}
176
177struct Cue {
178 text: String,
179 start: f64,
180 end: f64,
181}
182
183fn pick_segments(resp: &TTSWithTimestampsResponse) -> Result<(Vec<Segment>, bool)> {
188 let word_segs = |words: &[crate::timestamps::AlignmentSegmentWord]| -> Vec<Segment> {
189 words
190 .iter()
191 .map(|w| Segment {
192 text: w.text.clone(),
193 start: w.start,
194 end: w.end,
195 })
196 .collect()
197 };
198 let char_segs = |chars: &[crate::timestamps::AlignmentSegmentCharacter]| -> Vec<Segment> {
199 chars
200 .iter()
201 .map(|c| Segment {
202 text: c.text.clone(),
203 start: c.start,
204 end: c.end,
205 })
206 .collect()
207 };
208
209 let multi_words = resp.words.as_deref().filter(|w| w.len() >= 2);
211 let chars = resp.characters.as_deref().filter(|c| !c.is_empty());
213 let single_word = resp.words.as_deref().filter(|w| w.len() == 1);
215
216 if let Some(words) = multi_words {
217 Ok((word_segs(words), true))
218 } else if let Some(c) = chars {
219 Ok((char_segs(c), false))
220 } else if let Some(words) = single_word {
221 Ok((word_segs(words), true))
222 } else {
223 Err(TypecastError::CaptioningError(
224 "no alignment segments to caption from".into(),
225 ))
226 }
227}
228
229fn join_parts(parts: &[String], word_mode: bool) -> String {
232 let sep = if word_mode { " " } else { "" };
233 parts.join(sep).trim().to_string()
234}
235
236fn ends_in_sentence(text: &str) -> bool {
238 let trimmed = text.trim_end();
239 SENTENCE_TERMINATORS.iter().any(|t| trimmed.ends_with(t))
240}
241
242fn group_into_cues(segs: &[Segment], word_mode: bool) -> Vec<Cue> {
248 let mut cues: Vec<Cue> = Vec::new();
249 let mut parts: Vec<String> = Vec::new();
250 let mut cur_start: f64 = 0.0;
253 let mut last_end: f64 = 0.0;
254
255 fn emit(cues: &mut Vec<Cue>, text: String, start: f64, end: f64) {
257 if !text.is_empty() {
258 cues.push(Cue { text, start, end });
259 }
260 }
261
262 for seg in segs {
263 if !parts.is_empty() {
266 let mut tentative = parts.clone();
267 tentative.push(seg.text.clone());
268 let would_be = join_parts(&tentative, word_mode);
269 let too_long_secs = (seg.end - cur_start) > MAX_CAPTION_SECONDS;
270 let too_long_chars = would_be.chars().count() > MAX_CAPTION_CHARS;
271 if too_long_secs || too_long_chars {
272 emit(
274 &mut cues,
275 join_parts(&parts, word_mode),
276 cur_start,
277 last_end,
278 );
279 parts.clear();
280 }
281 }
282
283 if parts.is_empty() {
285 cur_start = seg.start;
286 }
287 parts.push(seg.text.clone());
288 last_end = seg.end;
289
290 if ends_in_sentence(&seg.text) {
292 emit(&mut cues, join_parts(&parts, word_mode), cur_start, seg.end);
293 parts.clear();
294 }
295 }
296
297 if !parts.is_empty() {
299 emit(
300 &mut cues,
301 join_parts(&parts, word_mode),
302 cur_start,
303 last_end,
304 );
305 }
306
307 cues
308}
309
310fn format_srt_time(seconds: f64) -> String {
312 let total_ms = (seconds * 1000.0).round() as i64;
313 let ms = total_ms % 1000;
314 let total_sec = total_ms / 1000;
315 let ss = total_sec % 60;
316 let total_min = total_sec / 60;
317 let mm = total_min % 60;
318 let hh = total_min / 60;
319 format!("{:02}:{:02}:{:02},{:03}", hh, mm, ss, ms)
320}
321
322fn format_vtt_time(seconds: f64) -> String {
324 format_srt_time(seconds).replace(',', ".")
325}
326
327fn format_captions(resp: &TTSWithTimestampsResponse, srt: bool) -> Result<String> {
329 let (segs, word_mode) = pick_segments(resp)?;
330 let cues = group_into_cues(&segs, word_mode);
331 if cues.is_empty() {
332 return Err(TypecastError::CaptioningError(
333 "no alignment segments to caption from".into(),
334 ));
335 }
336
337 let mut out = String::new();
338 if !srt {
339 out.push_str("WEBVTT\n\n");
340 }
341 for (i, cue) in cues.iter().enumerate() {
342 if srt {
343 out.push_str(&format!("{}\n", i + 1));
344 }
345 let (s, e) = if srt {
346 (format_srt_time(cue.start), format_srt_time(cue.end))
347 } else {
348 (format_vtt_time(cue.start), format_vtt_time(cue.end))
349 };
350 out.push_str(&format!("{} --> {}\n", s, e));
351 out.push_str(&cue.text);
352 out.push_str("\n\n");
353 }
354 Ok(out)
355}