use std::{
path::Path,
sync::mpsc,
thread::{self, JoinHandle},
};
use anyhow::anyhow;
use lumen_engine::{
audio::{AUDIO_CHANNELS, AUDIO_SAMPLE_RATE, AudioBuffer, AudioMixer, duration_samples},
composition::Composition,
};
use lumen_ffmpeg::{
AudioEncoderConfig, AudioFrame, CpuVideoFrame, MuxedEncoder, PixelFormat, SampleFormat,
VideoCodec, VideoEncoderConfig,
};
use super::media::LocalMediaStore;
pub(super) const ENCODER_FRAME_QUEUE_CAPACITY: usize = 2;
pub(super) struct EncoderFrame {
pub frame: u32,
pub pixels: Vec<u8>,
pub recycle_tx: mpsc::SyncSender<Vec<u8>>,
}
enum EncoderMessage {
Video(EncoderFrame),
Audio(AudioFrame),
}
pub(super) struct LumenFfmpegEncoder {
message_tx: Option<mpsc::SyncSender<EncoderMessage>>,
writer_handle: Option<JoinHandle<anyhow::Result<()>>>,
}
impl LumenFfmpegEncoder {
pub(super) fn create(
output: &Path,
width: u32,
height: u32,
fps: f32,
encoder: &str,
codec: VideoCodec,
include_audio: bool,
) -> anyhow::Result<Self> {
let mut config =
VideoEncoderConfig::cpu_rgba(width, height, fps.round().max(1.0) as u32, codec);
config.encoder_name = Some(encoder.to_string());
config.bit_rate = 14_000_000;
let audio = include_audio
.then(|| AudioEncoderConfig::aac(AUDIO_SAMPLE_RATE, AUDIO_CHANNELS as u16));
let (message_tx, message_rx) =
mpsc::sync_channel::<EncoderMessage>(ENCODER_FRAME_QUEUE_CAPACITY);
let (startup_tx, startup_rx) = mpsc::sync_channel::<anyhow::Result<()>>(1);
let output = output.to_string_lossy().to_string();
let writer_handle = thread::spawn(move || {
let mut encoder = match MuxedEncoder::create_with_audio(output, config, audio) {
Ok(encoder) => {
let _ = startup_tx.send(Ok(()));
encoder
}
Err(error) => {
let _ = startup_tx.send(Err(anyhow!(
"lumen-ffmpeg encoder failed to start: {error}"
)));
return Ok(());
}
};
while let Ok(message) = message_rx.recv() {
match message {
EncoderMessage::Video(frame) => {
let frame_index = frame.frame;
let cpu_frame = CpuVideoFrame {
width,
height,
stride: (width as usize) * 4,
pixel_format: PixelFormat::Rgba8,
pts: Some(i64::from(frame_index)),
data: frame.pixels,
};
encoder.write_video_frame(&cpu_frame).map_err(|error| {
anyhow!("lumen-ffmpeg encode failed at frame {frame_index}: {error}")
})?;
let _ = frame.recycle_tx.try_send(cpu_frame.data);
}
EncoderMessage::Audio(frame) => {
encoder.write_audio_frame(&frame).map_err(|error| {
anyhow!("lumen-ffmpeg audio encode failed: {error}")
})?;
}
}
}
encoder
.finish()
.map_err(|error| anyhow!("lumen-ffmpeg encoder finish failed: {error}"))
});
startup_rx
.recv()
.map_err(|_| anyhow!("lumen-ffmpeg encoder startup thread stopped"))??;
Ok(Self {
message_tx: Some(message_tx),
writer_handle: Some(writer_handle),
})
}
pub(super) fn send(&self, frame: EncoderFrame) -> anyhow::Result<()> {
self.message_tx
.as_ref()
.ok_or_else(|| anyhow!("lumen-ffmpeg encoder writer unavailable"))?
.send(EncoderMessage::Video(frame))
.map_err(|_| anyhow!("lumen-ffmpeg encoder writer stopped"))
}
fn send_audio(&self, frame: AudioFrame) -> anyhow::Result<()> {
self.message_tx
.as_ref()
.ok_or_else(|| anyhow!("lumen-ffmpeg encoder writer unavailable"))?
.send(EncoderMessage::Audio(frame))
.map_err(|_| anyhow!("lumen-ffmpeg encoder writer stopped"))
}
pub(super) fn finish(mut self) -> anyhow::Result<()> {
self.message_tx.take();
self.writer_handle
.take()
.ok_or_else(|| anyhow!("lumen-ffmpeg encoder writer thread missing"))?
.join()
.map_err(|_| anyhow!("lumen-ffmpeg encoder writer thread panicked"))?
}
}
pub(super) fn write_composited_audio(
composition: &Composition,
media_store: &LocalMediaStore,
encoder: &LumenFfmpegEncoder,
) -> anyhow::Result<()> {
write_composited_audio_with(composition, media_store, |frame| encoder.send_audio(frame))
}
pub(super) fn write_composited_audio_with(
composition: &Composition,
media_store: &LocalMediaStore,
mut write_frame: impl FnMut(AudioFrame) -> anyhow::Result<()>,
) -> anyhow::Result<()> {
let Some(audio) = composition.audio.as_ref() else {
return Ok(());
};
let total_samples = duration_samples(
composition.timeline.duration_frames,
composition.timeline.fps,
);
let mixer = AudioMixer::new(audio, media_store);
let mut start_sample = 0_u64;
let chunk_frames = (AUDIO_SAMPLE_RATE as usize).saturating_mul(30);
while start_sample < total_samples {
let frames = usize::try_from((total_samples - start_sample).min(chunk_frames as u64))
.unwrap_or(chunk_frames);
let mixed = mixer
.mix_range(start_sample, frames)
.map_err(|err| anyhow!("audio mix failed at sample {start_sample}: {err}"))?;
write_frame(audio_frame_from_buffer(&mixed, start_sample))?;
start_sample = start_sample.saturating_add(frames as u64);
}
Ok(())
}
pub(super) fn has_audio(composition: &Composition) -> bool {
composition
.audio
.as_ref()
.is_some_and(|audio| !audio.clips.is_empty())
}
fn audio_frame_from_buffer(buffer: &AudioBuffer, start_sample: u64) -> AudioFrame {
AudioFrame {
sample_rate: buffer.sample_rate(),
channels: buffer.channel_count() as u16,
sample_format: SampleFormat::F32,
pts: Some(start_sample as i64),
samples: buffer.frames(),
interleaved_f32: buffer.interleaved_f32(),
}
}