ffmpeg.rs

use super::{cmd, filter::Filter};
use crate::{
	render::filter::channel,
	time::{format_time, Time},
	Resolution
};
use anyhow::bail;
use camino::Utf8PathBuf as PathBuf;
use rational::Rational;
use std::{borrow::Cow, fmt::Write as _, process::Command};

pub(crate) struct FfmpegInput {
	pub(crate) concat: bool,
	pub(crate) loop_input: bool,
	pub(crate) fps: Option<Rational>,
	pub(crate) start: Option<Time>,
	pub(crate) duration: Option<Time>,
	pub(crate) path: PathBuf
}

impl FfmpegInput {
	pub(crate) fn new(path: PathBuf) -> Self {
		Self {
			concat: false,
			loop_input: false,
			fps: None,
			start: None,
			duration: None,
			path
		}
	}

	fn append_to_cmd(self, cmd: &mut Command) {
		if self.concat {
			cmd.arg("-f").arg("concat").arg("-safe").arg("0");
		}
		if self.loop_input {
			cmd.arg("-loop").arg("1");
		}
		if let Some(fps) = self.fps {
			cmd.arg("-r").arg(fps.to_string());
		}
		if let Some(start) = self.start {
			if self.path.ends_with(".mp4") {
				cmd.arg("-seek_streams_individualy").arg("false");
			}
			cmd.arg("-ss").arg(format_time(start));
		}
		if let Some(duration) = self.duration {
			cmd.arg("-t").arg(format_time(duration));
		}
		cmd.arg("-i").arg(self.path);
	}
}

#[derive(Clone, Copy)]
pub(crate) enum FfmpegOutputFormat {
	/// AV1 / FLAC
	Av1Flac,
	/// AV1 / OPUS
	Av1Opus,
	/// AVC (H.264) / AAC
	AvcAac
}

pub(crate) struct FfmpegOutput {
	pub(crate) format: FfmpegOutputFormat,
	pub(crate) audio_bitrate: Option<u64>,
	pub(crate) video_bitrate: Option<u64>,

	pub(crate) fps: Option<Rational>,
	pub(crate) duration: Option<Time>,
	pub(crate) time_base: Option<Rational>,
	pub(crate) fps_mode_vfr: bool,
	pub(crate) faststart: bool,

	pub(crate) path: PathBuf
}

impl FfmpegOutput {
	pub(crate) fn new(format: FfmpegOutputFormat, path: PathBuf) -> Self {
		Self {
			format,
			audio_bitrate: None,
			video_bitrate: None,
			fps: None,
			duration: None,
			time_base: None,
			fps_mode_vfr: false,
			faststart: false,
			path
		}
	}

	pub(crate) fn enable_faststart(mut self) -> Self {
		// only enable faststart for MP4 containers
		if matches!(self.format, FfmpegOutputFormat::AvcAac) {
			self.faststart = true;
		}
		self
	}

	fn append_to_cmd(self, cmd: &mut Command, venc: bool, _aenc: bool, vaapi: bool) {
		// select codec and bitrate
		const QUALITY: &str = "22";
		if venc {
			let vcodec = match (self.format, vaapi) {
				(FfmpegOutputFormat::Av1Flac, false) | (FfmpegOutputFormat::Av1Opus, false) => "libsvtav1",
				(FfmpegOutputFormat::Av1Flac, true) | (FfmpegOutputFormat::Av1Opus, true) => "av1_vaapi",
				(FfmpegOutputFormat::AvcAac, false) => "h264",
				(FfmpegOutputFormat::AvcAac, true) => "h264_vaapi"
			};
			cmd.arg("-c:v").arg(vcodec);

			if let Some(bv) = self.video_bitrate {
				cmd.arg("-b:v").arg(bv.to_string());
			} else if vaapi {
				cmd.arg("-rc_mode").arg("CQP");
				cmd.arg("-global_quality").arg(QUALITY);
			} else {
				cmd.arg("-crf").arg(QUALITY);
			}
		} else {
			cmd.arg("-c:v").arg("copy");
		}
		cmd.arg("-c:a").arg(match self.format {
			FfmpegOutputFormat::Av1Flac => "flac",
			FfmpegOutputFormat::Av1Opus => "libopus",
			FfmpegOutputFormat::AvcAac => "aac"
		});
		if let Some(ba) = self.audio_bitrate {
			cmd.arg("-b:a").arg(ba.to_string());
		} else if !matches!(self.format, FfmpegOutputFormat::Av1Flac) {
			cmd.arg("-b:a").arg("128k");
		}

		// other output options
		if let Some(fps) = self.fps {
			cmd.arg("-r").arg(fps.to_string());
		}
		if let Some(duration) = self.duration {
			cmd.arg("-t").arg(format_time(duration));
		}
		if let Some(time_base) = self.time_base {
			cmd.arg("-enc_time_base").arg(time_base.to_string());
		}
		if self.fps_mode_vfr {
			cmd.arg("-fps_mode").arg("vfr");
		}
		if self.faststart {
			cmd.arg("-movflags").arg("+faststart");
		}
		cmd.arg(self.path);
	}
}

enum FfmpegFilter {
	None,
	Filters {
		filters: Vec<Filter>,
		output: Cow<'static, str>
	},
	Loudnorm {
		stereo: bool
	},
	Rescale(Resolution)
}

pub(crate) struct Ffmpeg {
	inputs: Vec<FfmpegInput>,
	filter: FfmpegFilter,
	output: FfmpegOutput,

	filter_idx: usize
}

impl Ffmpeg {
	pub fn new(output: FfmpegOutput) -> Self {
		Self {
			inputs: Vec::new(),
			filter: FfmpegFilter::None,
			output,

			filter_idx: 0
		}
	}

	pub fn add_input(&mut self, input: FfmpegInput) -> String {
		self.inputs.push(input);
		(self.inputs.len() - 1).to_string()
	}

	pub fn add_filter(&mut self, filter: Filter) -> &mut Self {
		match &mut self.filter {
			FfmpegFilter::None => {
				self.filter = FfmpegFilter::Filters {
					filters: vec![filter],
					output: "0".into()
				}
			},
			FfmpegFilter::Filters { filters, .. } => filters.push(filter),
			_ => panic!("An incompatible type of filter has been set before")
		}
		self
	}

	pub fn set_filter_output<T: Into<Cow<'static, str>>>(
		&mut self,
		filter_output: T
	) -> &mut Self {
		match &mut self.filter {
			FfmpegFilter::None => {
				self.filter = FfmpegFilter::Filters {
					filters: vec![],
					output: filter_output.into()
				}
			},
			FfmpegFilter::Filters { output, .. } => *output = filter_output.into(),
			_ => panic!("An incompatible type of filter has been set before")
		}
		self
	}

	pub fn enable_loudnorm(&mut self, loudnorm_stereo: bool) -> &mut Self {
		match &mut self.filter {
			FfmpegFilter::None => {
				self.filter = FfmpegFilter::Loudnorm {
					stereo: loudnorm_stereo
				}
			},
			FfmpegFilter::Loudnorm { stereo } if *stereo == loudnorm_stereo => {},
			_ => panic!("An incompatible type of filter has been set before")
		}
		self
	}

	pub fn rescale_video(&mut self, res: Resolution) -> &mut Self {
		match &mut self.filter {
			FfmpegFilter::None => self.filter = FfmpegFilter::Rescale(res),
			_ => panic!("An incompatible type of filter has been set before")
		}
		self
	}

	pub fn run(mut self) -> anyhow::Result<()> {
		let mut cmd = cmd();
		cmd.arg("ffmpeg").arg("-hide_banner").arg("-y");

		// determine whether the video need to be re-encoded
		// vdec is only true if the video should be decoded on hardware
		let (vdec, venc, aenc) = match &self.filter {
			FfmpegFilter::None => (false, false, false),
			FfmpegFilter::Filters { .. } => (false, true, true),
			FfmpegFilter::Loudnorm { .. } => (false, false, true),
			FfmpegFilter::Rescale(_) => (true, true, false)
		};

		// initialise a vaapi device if one exists
		let vaapi_device: PathBuf = "/dev/dri/renderD128".into();
		let vaapi = vaapi_device.exists();
		if vaapi && venc {
			if vdec {
				cmd.arg("-hwaccel").arg("vaapi");
				cmd.arg("-hwaccel_device").arg(vaapi_device);
				cmd.arg("-hwaccel_output_format").arg("vaapi");
			} else {
				cmd.arg("-vaapi_device").arg(&vaapi_device);
			}
		}

		// append all the inputs
		for i in self.inputs {
			i.append_to_cmd(&mut cmd);
		}

		// always try to synchronise audio
		cmd.arg("-async").arg("1");

		// apply filters
		match self.filter {
			FfmpegFilter::None => {},
			FfmpegFilter::Filters { filters, output } => {
				let mut complex = String::new();
				for filter in filters {
					filter
						.append_to_complex_filter(&mut complex, &mut self.filter_idx)?;
				}
				if vaapi {
					write!(complex, "{}format=nv12,hwupload[v]", channel('v', &output))?;
				} else {
					write!(complex, "{}null[v]", channel('v', &output))?;
				}
				cmd.arg("-filter_complex").arg(complex);
				cmd.arg("-map").arg("[v]");
				cmd.arg("-map").arg(channel('a', &output));
			},
			FfmpegFilter::Loudnorm { stereo: false } => {
				cmd.arg("-af").arg(concat!(
					"pan=mono|c0=FR,",
					"loudnorm=dual_mono=true:print_format=summary,",
					"pan=stereo|c0=c0|c1=c0,",
					"aformat=sample_rates=48000"
				));
			},
			FfmpegFilter::Loudnorm { stereo: true } => {
				cmd.arg("-af")
					.arg("loudnorm=print_format=summary,aformat=sample_rates=48000");
			},
			FfmpegFilter::Rescale(res) => {
				cmd.arg("-vf").arg(if vaapi {
					format!("scale_vaapi=w={}:h={}", res.width(), res.height())
				} else {
					format!("scale=w={}:h={}", res.width(), res.height())
				});
			}
		}

		self.output.append_to_cmd(&mut cmd, venc, aenc, vaapi);

		let status = cmd.status()?;
		if status.success() {
			Ok(())
		} else {
			bail!("ffmpeg failed with exit code {:?}", status.code())
		}
	}
}