1use std::path::{Path, PathBuf};
2use std::time::{Duration, Instant, SystemTime, UNIX_EPOCH};
3use tokio::io::AsyncBufReadExt;
4use tokio::process::Command;
5
6use crate::{Codec, RateControlMode, Resolution, ffmpeg_path, probe};
7
8#[derive(Debug, Clone)]
10pub struct EncodeJob {
11 pub input: String,
13 pub output: String,
15 pub resolution: Option<Resolution>,
17 pub codec: Codec,
19 pub crf: i32,
21 pub rate_control: RateControlMode,
23 pub target_bitrate: f64, pub max_bitrate: f64, pub bufsize: f64, pub preset: String,
31 pub extra_args: Vec<String>,
33}
34
35#[derive(Debug, Clone)]
37pub struct EncodeResult {
38 pub job: EncodeJob,
40 pub bitrate: f64, pub file_size: u64, pub duration: Duration, }
47
48#[derive(Debug, Clone, Default)]
50pub struct Progress {
51 pub frame: i64,
53 pub fps: f64,
55 pub bitrate: f64, pub speed: f64, pub time: Duration,
61}
62
63pub async fn encode(
65 job: EncodeJob,
66 progress_tx: Option<tokio::sync::mpsc::Sender<Progress>>,
67) -> anyhow::Result<EncodeResult> {
68 match job.rate_control {
69 RateControlMode::Vbr => encode_two_pass(job, progress_tx).await,
70 _ => encode_single_pass(job, progress_tx).await,
71 }
72}
73
74async fn encode_single_pass(
75 job: EncodeJob,
76 progress_tx: Option<tokio::sync::mpsc::Sender<Progress>>,
77) -> anyhow::Result<EncodeResult> {
78 let args = build_encode_args(&job, EncodePass::Single)?;
79 run_encode(job, args, progress_tx).await
80}
81
82async fn encode_two_pass(
83 job: EncodeJob,
84 progress_tx: Option<tokio::sync::mpsc::Sender<Progress>>,
85) -> anyhow::Result<EncodeResult> {
86 if job.target_bitrate <= 0.0 {
87 anyhow::bail!("target bitrate must be greater than zero for VBR mode");
88 }
89
90 let passlog_prefix = make_passlog_prefix(&job.output);
91 let cleanup = PasslogCleanup::new(passlog_prefix.clone());
92
93 let first_pass_args = build_encode_args(&job, EncodePass::First(&passlog_prefix))?;
94 run_ffmpeg(first_pass_args, None).await?;
95
96 let second_pass_args = build_encode_args(&job, EncodePass::Second(&passlog_prefix))?;
97 let result = run_encode(job, second_pass_args, progress_tx).await;
98
99 cleanup.run();
100 result
101}
102
103async fn run_encode(
104 job: EncodeJob,
105 args: Vec<String>,
106 progress_tx: Option<tokio::sync::mpsc::Sender<Progress>>,
107) -> anyhow::Result<EncodeResult> {
108 let start = Instant::now();
109 run_ffmpeg(args, progress_tx).await?;
110
111 let elapsed = start.elapsed();
112
113 let meta = std::fs::metadata(&job.output)
115 .map_err(|e| anyhow::anyhow!("failed to stat output: {e}"))?;
116
117 let probe_result = probe(&job.output).await?;
118 let bitrate = probe_result.format.bit_rate as f64 / 1000.0;
119
120 Ok(EncodeResult { job, bitrate, file_size: meta.len(), duration: elapsed })
121}
122
123async fn run_ffmpeg(
124 args: Vec<String>,
125 progress_tx: Option<tokio::sync::mpsc::Sender<Progress>>,
126) -> anyhow::Result<()> {
127 let mut cmd = Command::new(ffmpeg_path());
128 cmd.args(&args).stdout(std::process::Stdio::piped()).stderr(std::process::Stdio::piped());
129
130 let mut child = cmd.spawn().map_err(|e| anyhow::anyhow!("failed to start ffmpeg: {e}"))?;
131
132 if let Some(stdout) = child.stdout.take() {
134 let tx = progress_tx.clone();
135 tokio::spawn(async move {
136 let reader = tokio::io::BufReader::new(stdout);
137 let mut lines = reader.lines();
138 let mut p = Progress::default();
139 while let Ok(Some(line)) = lines.next_line().await {
140 if parse_progress_line(&line, &mut p)
141 && let Some(ref tx) = tx
142 {
143 let _ = tx.try_send(p.clone());
144 }
145 }
146 });
147 }
148
149 let output = child.wait_with_output().await?;
150 if !output.status.success() {
151 let stderr = String::from_utf8_lossy(&output.stderr);
152 anyhow::bail!("ffmpeg encode failed: {stderr}");
153 }
154
155 Ok(())
156}
157
158pub async fn extract(input: &str, output: &str, start: f64, duration: f64) -> anyhow::Result<()> {
160 if start.is_finite() && start < 0.0 {
161 anyhow::bail!("extract start must be non-negative, got {start}");
162 }
163 if !duration.is_finite() || duration <= 0.0 {
164 anyhow::bail!("extract duration must be positive, got {duration}");
165 }
166
167 let args = vec![
168 "-y".to_string(),
169 "-ss".into(),
170 format!("{start:.6}"),
171 "-i".into(),
172 input.into(),
173 "-t".into(),
174 format!("{duration:.6}"),
175 "-c".into(),
176 "copy".into(),
177 "-avoid_negative_ts".into(),
178 "make_zero".into(),
179 output.into(),
180 ];
181
182 let output = Command::new(ffmpeg_path())
183 .args(&args)
184 .stderr(std::process::Stdio::piped())
185 .output()
186 .await?;
187
188 if !output.status.success() {
189 let stderr = String::from_utf8_lossy(&output.stderr);
190 anyhow::bail!("ffmpeg extract failed: {stderr}");
191 }
192 Ok(())
193}
194
195pub async fn concat(inputs: &[String], output: &str) -> anyhow::Result<()> {
197 if inputs.is_empty() {
198 anyhow::bail!("cannot concat an empty input list");
199 }
200
201 let list_path = make_concat_list_path(output);
202 let list_body = inputs
203 .iter()
204 .map(|path| format!("file '{}'", escape_concat_path(path)))
205 .collect::<Vec<_>>()
206 .join("\n");
207 std::fs::write(&list_path, format!("{list_body}\n"))?;
208
209 let args = vec![
210 "-y".to_string(),
211 "-f".into(),
212 "concat".into(),
213 "-safe".into(),
214 "0".into(),
215 "-i".into(),
216 list_path.to_string_lossy().into_owned(),
217 "-c".into(),
218 "copy".into(),
219 output.into(),
220 ];
221
222 let result = run_ffmpeg(args, None).await;
223 let _ = std::fs::remove_file(&list_path);
224 result
225}
226
227enum EncodePass<'a> {
228 Single,
229 First(&'a Path),
230 Second(&'a Path),
231}
232
233fn build_encode_args(job: &EncodeJob, pass: EncodePass<'_>) -> anyhow::Result<Vec<String>> {
234 let mut args = vec!["-y".into(), "-i".into(), job.input.clone(), "-an".into()];
235
236 if !matches!(pass, EncodePass::First(_)) {
237 args.extend(["-progress".into(), "pipe:1".into(), "-nostats".into()]);
238 }
239
240 args.extend(["-c:v".into(), job.codec.as_str().into()]);
241
242 match job.rate_control {
244 RateControlMode::Qp => match job.codec {
245 Codec::SvtAv1 => {
246 args.extend(["-qp".into(), job.crf.to_string()]);
247 args.extend(["-svtav1-params".into(), "enable-adaptive-quantization=0".into()]);
248 }
249 _ => {
250 args.extend(["-qp".into(), job.crf.to_string()]);
251 }
252 },
253 RateControlMode::CappedCrf => {
254 if job.max_bitrate <= 0.0 {
255 anyhow::bail!("max bitrate must be greater than zero for capped CRF mode");
256 }
257 let bufsize = if job.bufsize > 0.0 { job.bufsize } else { job.max_bitrate * 2.0 };
258 args.extend(["-crf".into(), job.crf.to_string()]);
259 args.extend(["-maxrate".into(), format!("{:.0}k", job.max_bitrate)]);
260 args.extend(["-bufsize".into(), format!("{bufsize:.0}k")]);
261 }
262 RateControlMode::Vbr => {
263 if job.target_bitrate <= 0.0 {
264 anyhow::bail!("target bitrate must be greater than zero for VBR mode");
265 }
266 args.extend(["-b:v".into(), format!("{:.0}k", job.target_bitrate)]);
267 args.extend(["-maxrate".into(), format!("{:.0}k", job.target_bitrate * 2.0)]);
268 args.extend(["-bufsize".into(), format!("{:.0}k", job.target_bitrate * 4.0)]);
269
270 let passlog = match pass {
271 EncodePass::First(path) => {
272 args.extend(["-pass".into(), "1".into()]);
273 path
274 }
275 EncodePass::Second(path) => {
276 args.extend(["-pass".into(), "2".into()]);
277 path
278 }
279 EncodePass::Single => {
280 anyhow::bail!("VBR mode requires a two-pass encode flow");
281 }
282 };
283
284 args.extend(["-passlogfile".into(), passlog.to_string_lossy().into_owned()]);
285 }
286 RateControlMode::Crf => {
287 args.extend(["-crf".into(), job.crf.to_string()]);
288 }
289 }
290
291 if !job.preset.is_empty() {
292 args.extend(["-preset".into(), job.preset.clone()]);
293 }
294
295 if let Some(ref res) = job.resolution
296 && res.width > 0
297 && res.height > 0
298 {
299 args.extend(["-vf".into(), format!("scale={}:{}:flags=lanczos", res.width, res.height)]);
300 }
301
302 args.extend(job.extra_args.iter().cloned());
303
304 match pass {
305 EncodePass::First(_) => {
306 args.extend(["-f".into(), "null".into()]);
307 args.push(null_output_path().into());
308 }
309 EncodePass::Single | EncodePass::Second(_) => args.push(job.output.clone()),
310 }
311
312 Ok(args)
313}
314
315fn escape_concat_path(path: &str) -> String {
319 path.replace('\\', "\\\\").replace('\'', "\\'")
320}
321
322fn make_passlog_prefix(output: &str) -> PathBuf {
323 let output_path = Path::new(output);
324 let parent =
325 output_path.parent().filter(|p| !p.as_os_str().is_empty()).unwrap_or(Path::new("."));
326 let stem = output_path.file_stem().and_then(|s| s.to_str()).unwrap_or("viser");
327 let unique = SystemTime::now().duration_since(UNIX_EPOCH).map(|d| d.as_millis()).unwrap_or(0);
328 parent.join(format!(".{stem}.viser-passlog-{unique}-{}", std::process::id()))
329}
330
331fn make_concat_list_path(output: &str) -> PathBuf {
332 let output_path = Path::new(output);
333 let parent =
334 output_path.parent().filter(|p| !p.as_os_str().is_empty()).unwrap_or(Path::new("."));
335 let stem = output_path.file_stem().and_then(|s| s.to_str()).unwrap_or("viser");
336 let unique = SystemTime::now().duration_since(UNIX_EPOCH).map(|d| d.as_millis()).unwrap_or(0);
337 parent.join(format!(".{stem}.viser-concat-{unique}-{}.txt", std::process::id()))
338}
339
340fn null_output_path() -> &'static str {
341 if cfg!(windows) { "NUL" } else { "/dev/null" }
342}
343
344struct PasslogCleanup {
345 parent: PathBuf,
346 prefix: String,
347}
348
349impl PasslogCleanup {
350 fn new(path: PathBuf) -> Self {
351 let parent = path.parent().unwrap_or(Path::new(".")).to_path_buf();
352 let prefix = path.file_name().and_then(|name| name.to_str()).unwrap_or_default().to_owned();
353 Self { parent, prefix }
354 }
355
356 fn run(&self) {
357 let Ok(entries) = std::fs::read_dir(&self.parent) else {
358 return;
359 };
360
361 for entry in entries.flatten() {
362 let path = entry.path();
363 let Some(name) = path.file_name().and_then(|name| name.to_str()) else {
364 continue;
365 };
366 if !name.starts_with(&self.prefix) {
367 continue;
368 }
369 if let Err(err) = std::fs::remove_file(&path) {
370 tracing::debug!(?path, ?err, "failed to remove ffmpeg two-pass log file");
371 }
372 }
373 }
374}
375
376fn parse_progress_line(line: &str, p: &mut Progress) -> bool {
378 let Some((key, value)) = line.split_once('=') else {
379 return false;
380 };
381
382 match key {
383 "frame" => {
384 p.frame = value.parse().unwrap_or(0);
385 }
386 "fps" => {
387 p.fps = value.parse().unwrap_or(0.0);
388 }
389 "bitrate" => {
390 let v = value.trim_end_matches("kbits/s");
391 p.bitrate = v.parse().unwrap_or(0.0);
392 }
393 "speed" => {
394 let v = value.trim_end_matches('x');
395 p.speed = v.parse().unwrap_or(0.0);
396 }
397 "out_time_us" => {
398 let us: u64 = value.parse().unwrap_or(0);
399 p.time = Duration::from_micros(us);
400 }
401 "progress" => return true,
402 _ => {}
403 }
404 false
405}
406
407#[cfg(test)]
408mod tests {
409 use super::*;
410 use crate::Codec;
411
412 fn sample_job(mode: RateControlMode) -> EncodeJob {
413 EncodeJob {
414 input: "input.mp4".into(),
415 output: "output.mp4".into(),
416 resolution: Some(crate::Resolution::new(1280, 720)),
417 codec: Codec::X264,
418 crf: 23,
419 rate_control: mode,
420 target_bitrate: 2500.0,
421 max_bitrate: 3000.0,
422 bufsize: 6000.0,
423 preset: "medium".into(),
424 extra_args: vec![],
425 }
426 }
427
428 #[test]
429 fn test_build_encode_args_crf_single_pass() {
430 let args =
431 build_encode_args(&sample_job(RateControlMode::Crf), EncodePass::Single).unwrap();
432 assert!(args.windows(2).any(|w| w == ["-crf", "23"]));
433 assert_eq!(args.last().unwrap(), "output.mp4");
434 }
435
436 #[test]
437 fn test_build_encode_args_vbr_first_pass_uses_null_output() {
438 let job = sample_job(RateControlMode::Vbr);
439 let passlog = Path::new("/tmp/viser-passlog");
440 let args = build_encode_args(&job, EncodePass::First(passlog)).unwrap();
441 assert!(args.windows(2).any(|w| w == ["-pass", "1"]));
442 assert!(args.windows(2).any(|w| w == ["-f", "null"]));
443 assert_eq!(args.last().unwrap(), null_output_path());
444 }
445
446 #[test]
447 fn test_build_encode_args_vbr_second_pass_writes_output() {
448 let job = sample_job(RateControlMode::Vbr);
449 let passlog = Path::new("/tmp/viser-passlog");
450 let args = build_encode_args(&job, EncodePass::Second(passlog)).unwrap();
451 assert!(args.windows(2).any(|w| w == ["-pass", "2"]));
452 assert_eq!(args.last().unwrap(), "output.mp4");
453 }
454
455 #[test]
456 fn test_build_encode_args_capped_crf_sets_vbv() {
457 let args =
458 build_encode_args(&sample_job(RateControlMode::CappedCrf), EncodePass::Single).unwrap();
459 assert!(args.windows(2).any(|w| w == ["-crf", "23"]));
460 assert!(args.windows(2).any(|w| w == ["-maxrate", "3000k"]));
461 assert!(args.windows(2).any(|w| w == ["-bufsize", "6000k"]));
462 }
463}