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, EncoderBackend, 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 if let Some(ref tx) = tx {
142 let _ = tx.try_send(p.clone());
143 }
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 let args = vec![
161 "-y".to_string(),
162 "-ss".into(),
163 format!("{start:.6}"),
164 "-i".into(),
165 input.into(),
166 "-t".into(),
167 format!("{duration:.6}"),
168 "-c".into(),
169 "copy".into(),
170 "-avoid_negative_ts".into(),
171 "make_zero".into(),
172 output.into(),
173 ];
174
175 let output = Command::new(ffmpeg_path())
176 .args(&args)
177 .stderr(std::process::Stdio::piped())
178 .output()
179 .await?;
180
181 if !output.status.success() {
182 let stderr = String::from_utf8_lossy(&output.stderr);
183 anyhow::bail!("ffmpeg extract failed: {stderr}");
184 }
185 Ok(())
186}
187
188pub async fn concat(inputs: &[String], output: &str) -> anyhow::Result<()> {
190 if inputs.is_empty() {
191 anyhow::bail!("cannot concat an empty input list");
192 }
193
194 let list_path = make_concat_list_path(output);
195 let list_body = inputs
196 .iter()
197 .map(|path| format!("file '{}'", path.replace('\'', "'\\''")))
198 .collect::<Vec<_>>()
199 .join("\n");
200 std::fs::write(&list_path, format!("{list_body}\n"))?;
201
202 let args = vec![
203 "-y".to_string(),
204 "-f".into(),
205 "concat".into(),
206 "-safe".into(),
207 "0".into(),
208 "-i".into(),
209 list_path.to_string_lossy().into_owned(),
210 "-c".into(),
211 "copy".into(),
212 output.into(),
213 ];
214
215 let result = run_ffmpeg(args, None).await;
216 let _ = std::fs::remove_file(&list_path);
217 result
218}
219
220enum EncodePass<'a> {
221 Single,
222 First(&'a Path),
223 Second(&'a Path),
224}
225
226fn build_encode_args(job: &EncodeJob, pass: EncodePass<'_>) -> anyhow::Result<Vec<String>> {
227 let mut args = vec!["-y".into(), "-i".into(), job.input.clone(), "-an".into()];
228
229 if !matches!(pass, EncodePass::First(_)) {
230 args.extend(["-progress".into(), "pipe:1".into(), "-nostats".into()]);
231 }
232
233 args.extend(["-c:v".into(), job.codec.as_str().into()]);
234
235 if job.codec.is_hardware() {
236 build_hw_args(&mut args, job, &pass)?;
237 } else {
238 build_sw_args(&mut args, job, &pass)?;
239 }
240
241 if !job.preset.is_empty() {
242 if job.codec.is_hardware() {
243 add_hw_preset(&mut args, job.codec, &job.preset);
244 } else {
245 args.extend(["-preset".into(), job.preset.clone()]);
246 }
247 }
248
249 if let Some(ref res) = job.resolution {
250 if res.width > 0 && res.height > 0 {
251 args.extend([
252 "-vf".into(),
253 format!("scale={}:{}:flags=lanczos", res.width, res.height),
254 ]);
255 }
256 }
257
258 args.extend(job.extra_args.iter().cloned());
259
260 match pass {
261 EncodePass::First(_) => {
262 args.extend(["-f".into(), "null".into()]);
263 args.push(null_output_path().into());
264 }
265 EncodePass::Single | EncodePass::Second(_) => args.push(job.output.clone()),
266 }
267
268 Ok(args)
269}
270
271fn build_sw_args(
272 args: &mut Vec<String>,
273 job: &EncodeJob,
274 pass: &EncodePass<'_>,
275) -> anyhow::Result<()> {
276 match job.rate_control {
277 RateControlMode::Qp => {
278 if job.codec == Codec::SvtAv1 {
279 args.extend(["-qp".into(), job.crf.to_string()]);
280 args.extend(["-svtav1-params".into(), "enable-adaptive-quantization=0".into()]);
281 } else {
282 args.extend(["-qp".into(), job.crf.to_string()]);
283 }
284 }
285 RateControlMode::CappedCrf => {
286 if job.max_bitrate <= 0.0 {
287 anyhow::bail!("max bitrate must be greater than zero for capped CRF mode");
288 }
289 let bufsize = if job.bufsize > 0.0 { job.bufsize } else { job.max_bitrate * 2.0 };
290 args.extend(["-crf".into(), job.crf.to_string()]);
291 args.extend(["-maxrate".into(), format!("{:.0}k", job.max_bitrate)]);
292 args.extend(["-bufsize".into(), format!("{bufsize:.0}k")]);
293 }
294 RateControlMode::Vbr => {
295 if job.target_bitrate <= 0.0 {
296 anyhow::bail!("target bitrate must be greater than zero for VBR mode");
297 }
298 args.extend(["-b:v".into(), format!("{:.0}k", job.target_bitrate)]);
299 args.extend(["-maxrate".into(), format!("{:.0}k", job.target_bitrate * 2.0)]);
300 args.extend(["-bufsize".into(), format!("{:.0}k", job.target_bitrate * 4.0)]);
301
302 let passlog = match pass {
303 EncodePass::First(path) => {
304 args.extend(["-pass".into(), "1".into()]);
305 path
306 }
307 EncodePass::Second(path) => {
308 args.extend(["-pass".into(), "2".into()]);
309 path
310 }
311 EncodePass::Single => {
312 anyhow::bail!("VBR mode requires a two-pass encode flow");
313 }
314 };
315 args.extend(["-passlogfile".into(), passlog.to_string_lossy().into_owned()]);
316 }
317 RateControlMode::Crf => {
318 args.extend(["-crf".into(), job.crf.to_string()]);
319 }
320 }
321 Ok(())
322}
323
324fn build_hw_args(
325 args: &mut Vec<String>,
326 job: &EncodeJob,
327 _pass: &EncodePass<'_>,
328) -> anyhow::Result<()> {
329 let backend = job.codec.backend();
330 match job.rate_control {
331 RateControlMode::Crf | RateControlMode::CappedCrf => {
332 match backend {
333 EncoderBackend::Nvenc => {
334 let cq = crf_to_nvenc_cq(job.crf);
335 args.extend(["-cq".into(), cq.to_string()]);
336 args.extend(["-rc".into(), "constqp".into()]);
337 }
338 EncoderBackend::Qsv => {
339 let gq = crf_to_qsv_quality(job.crf);
340 args.extend(["-global_quality".into(), gq.to_string()]);
341 }
342 EncoderBackend::VideoToolbox => {
343 let qual = crf_to_vt_quality(job.crf);
344 args.extend(["-quality".into(), qual.to_string()]);
345 }
346 EncoderBackend::Vaapi => {
347 let gq = crf_to_qsv_quality(job.crf);
348 args.extend(["-global_quality".into(), gq.to_string()]);
349 }
350 EncoderBackend::Amf => {
351 args.extend(["-qp_i".into(), job.crf.to_string()]);
352 args.extend(["-qp_p".into(), (job.crf + 2).to_string()]);
353 args.extend(["-usage".into(), "transcoding".into()]);
354 }
355 EncoderBackend::Software => unreachable!(),
356 }
357 if let RateControlMode::CappedCrf = job.rate_control {
359 if job.max_bitrate <= 0.0 {
360 anyhow::bail!("max bitrate must be greater than zero for capped CRF mode");
361 }
362 let bufsize = if job.bufsize > 0.0 { job.bufsize } else { job.max_bitrate * 2.0 };
363 args.extend(["-maxrate".into(), format!("{:.0}k", job.max_bitrate)]);
364 args.extend(["-bufsize".into(), format!("{bufsize:.0}k")]);
365 }
366 }
367 RateControlMode::Qp => match backend {
368 EncoderBackend::VideoToolbox => {
369 anyhow::bail!("VideoToolbox does not support QP rate control mode");
370 }
371 _ => {
372 args.extend(["-qp".into(), job.crf.to_string()]);
373 }
374 },
375 RateControlMode::Vbr => {
376 if job.target_bitrate <= 0.0 {
377 anyhow::bail!("target bitrate must be greater than zero for VBR mode");
378 }
379 args.extend(["-b:v".into(), format!("{:.0}k", job.target_bitrate)]);
380 args.extend(["-maxrate".into(), format!("{:.0}k", job.target_bitrate * 2.0)]);
381 args.extend(["-bufsize".into(), format!("{:.0}k", job.target_bitrate * 4.0)]);
382
383 if backend == EncoderBackend::Nvenc {
384 args.extend(["-rc".into(), "vbr_hq".into()]);
385 }
386 }
387 }
388 Ok(())
389}
390
391fn crf_to_nvenc_cq(crf: i32) -> i32 {
392 let cq = (crf * 51) / 63;
393 cq.clamp(1, 51)
394}
395
396fn crf_to_qsv_quality(crf: i32) -> i32 {
397 let gq = 100 - ((crf * 100) / 51);
398 gq.clamp(1, 100)
399}
400
401fn crf_to_vt_quality(crf: i32) -> f64 {
402 let q = 1.0 - (crf as f64 / 51.0);
403 q.clamp(0.0, 1.0)
404}
405
406fn add_hw_preset(args: &mut Vec<String>, codec: Codec, preset: &str) {
407 match codec.backend() {
408 EncoderBackend::Nvenc => {
409 let p = map_nvenc_preset(preset);
410 args.extend(["-preset".into(), p.into()]);
411 }
412 EncoderBackend::Qsv => {
413 args.extend(["-preset".into(), preset.to_string()]);
414 }
415 EncoderBackend::Vaapi => {
416 args.extend(["-compression_level".into(), map_vaapi_preset(preset).into()]);
417 }
418 EncoderBackend::Amf => {
419 args.extend(["-quality".into(), map_amf_quality(preset).into()]);
420 }
421 EncoderBackend::VideoToolbox => {
422 if preset == "ultrafast" || preset == "superfast" || preset == "veryfast" {
423 args.extend(["-realtime".into(), "1".into()]);
424 }
425 }
426 EncoderBackend::Software => unreachable!(),
427 }
428}
429
430fn map_nvenc_preset(preset: &str) -> &str {
431 match preset {
432 "ultrafast" | "superfast" => "p1",
433 "veryfast" => "p2",
434 "faster" => "p3",
435 "fast" => "p4",
436 "medium" => "p5",
437 "slow" => "p6",
438 "slower" | "veryslow" => "p7",
439 other => other,
440 }
441}
442
443fn map_vaapi_preset(preset: &str) -> &str {
444 match preset {
445 "ultrafast" | "superfast" => "1",
446 "veryfast" | "faster" => "2",
447 "fast" | "medium" => "3",
448 "slow" => "4",
449 "slower" | "veryslow" => "5",
450 other => other,
451 }
452}
453
454fn map_amf_quality(preset: &str) -> &str {
455 match preset {
456 "ultrafast" | "superfast" => "speed",
457 "veryfast" | "faster" | "fast" => "balanced",
458 "medium" | "slow" | "slower" | "veryslow" => "quality",
459 other => other,
460 }
461}
462
463fn make_passlog_prefix(output: &str) -> PathBuf {
464 let output_path = Path::new(output);
465 let parent =
466 output_path.parent().filter(|p| !p.as_os_str().is_empty()).unwrap_or(Path::new("."));
467 let stem = output_path.file_stem().and_then(|s| s.to_str()).unwrap_or("viser");
468 let unique = SystemTime::now().duration_since(UNIX_EPOCH).map(|d| d.as_millis()).unwrap_or(0);
469 parent.join(format!(".{stem}.viser-passlog-{unique}-{}", std::process::id()))
470}
471
472fn make_concat_list_path(output: &str) -> PathBuf {
473 let output_path = Path::new(output);
474 let parent =
475 output_path.parent().filter(|p| !p.as_os_str().is_empty()).unwrap_or(Path::new("."));
476 let stem = output_path.file_stem().and_then(|s| s.to_str()).unwrap_or("viser");
477 let unique = SystemTime::now().duration_since(UNIX_EPOCH).map(|d| d.as_millis()).unwrap_or(0);
478 parent.join(format!(".{stem}.viser-concat-{unique}-{}.txt", std::process::id()))
479}
480
481fn null_output_path() -> &'static str {
482 if cfg!(windows) { "NUL" } else { "/dev/null" }
483}
484
485struct PasslogCleanup {
486 parent: PathBuf,
487 prefix: String,
488}
489
490impl PasslogCleanup {
491 fn new(path: PathBuf) -> Self {
492 let parent = path.parent().unwrap_or(Path::new(".")).to_path_buf();
493 let prefix = path.file_name().and_then(|name| name.to_str()).unwrap_or_default().to_owned();
494 Self { parent, prefix }
495 }
496
497 fn run(&self) {
498 let Ok(entries) = std::fs::read_dir(&self.parent) else {
499 return;
500 };
501
502 for entry in entries.flatten() {
503 let path = entry.path();
504 let Some(name) = path.file_name().and_then(|name| name.to_str()) else {
505 continue;
506 };
507 if !name.starts_with(&self.prefix) {
508 continue;
509 }
510 if let Err(err) = std::fs::remove_file(&path) {
511 tracing::debug!(?path, ?err, "failed to remove ffmpeg two-pass log file");
512 }
513 }
514 }
515}
516
517fn parse_progress_line(line: &str, p: &mut Progress) -> bool {
519 let Some((key, value)) = line.split_once('=') else {
520 return false;
521 };
522
523 match key {
524 "frame" => {
525 p.frame = value.parse().unwrap_or(0);
526 }
527 "fps" => {
528 p.fps = value.parse().unwrap_or(0.0);
529 }
530 "bitrate" => {
531 let v = value.trim_end_matches("kbits/s");
532 p.bitrate = v.parse().unwrap_or(0.0);
533 }
534 "speed" => {
535 let v = value.trim_end_matches('x');
536 p.speed = v.parse().unwrap_or(0.0);
537 }
538 "out_time_us" => {
539 let us: u64 = value.parse().unwrap_or(0);
540 p.time = Duration::from_micros(us);
541 }
542 "progress" => return true,
543 _ => {}
544 }
545 false
546}
547
548#[cfg(test)]
549mod tests {
550 use super::*;
551 use crate::Codec;
552
553 fn sample_job(mode: RateControlMode) -> EncodeJob {
554 EncodeJob {
555 input: "input.mp4".into(),
556 output: "output.mp4".into(),
557 resolution: Some(crate::Resolution::new(1280, 720)),
558 codec: Codec::X264,
559 crf: 23,
560 rate_control: mode,
561 target_bitrate: 2500.0,
562 max_bitrate: 3000.0,
563 bufsize: 6000.0,
564 preset: "medium".into(),
565 extra_args: vec![],
566 }
567 }
568
569 #[test]
570 fn test_build_encode_args_crf_single_pass() {
571 let args =
572 build_encode_args(&sample_job(RateControlMode::Crf), EncodePass::Single).unwrap();
573 assert!(args.windows(2).any(|w| w == ["-crf", "23"]));
574 assert_eq!(args.last().unwrap(), "output.mp4");
575 }
576
577 #[test]
578 fn test_build_encode_args_vbr_first_pass_uses_null_output() {
579 let job = sample_job(RateControlMode::Vbr);
580 let passlog = Path::new("/tmp/viser-passlog");
581 let args = build_encode_args(&job, EncodePass::First(passlog)).unwrap();
582 assert!(args.windows(2).any(|w| w == ["-pass", "1"]));
583 assert!(args.windows(2).any(|w| w == ["-f", "null"]));
584 assert_eq!(args.last().unwrap(), null_output_path());
585 }
586
587 #[test]
588 fn test_build_encode_args_vbr_second_pass_writes_output() {
589 let job = sample_job(RateControlMode::Vbr);
590 let passlog = Path::new("/tmp/viser-passlog");
591 let args = build_encode_args(&job, EncodePass::Second(passlog)).unwrap();
592 assert!(args.windows(2).any(|w| w == ["-pass", "2"]));
593 assert_eq!(args.last().unwrap(), "output.mp4");
594 }
595
596 #[test]
597 fn test_build_encode_args_capped_crf_sets_vbv() {
598 let args =
599 build_encode_args(&sample_job(RateControlMode::CappedCrf), EncodePass::Single).unwrap();
600 assert!(args.windows(2).any(|w| w == ["-crf", "23"]));
601 assert!(args.windows(2).any(|w| w == ["-maxrate", "3000k"]));
602 assert!(args.windows(2).any(|w| w == ["-bufsize", "6000k"]));
603 }
604}