ImageApi/src/video.rs

use crate::is_video;
use actix::prelude::*;
use futures::TryFutureExt;
use log::{debug, error, info, trace, warn};
use std::io::Result;
use std::path::{Path, PathBuf};
use std::process::{Child, Command, ExitStatus, Stdio};
use std::sync::Arc;
use tokio::sync::Semaphore;
use walkdir::{DirEntry, WalkDir};
// ffmpeg -i test.mp4 -c:v h264 -flags +cgop -g 30 -hls_time 3 out.m3u8
// ffmpeg -i "filename.mp4" -preset veryfast -c:v libx264 -f hls -hls_list_size 100 -hls_time 2 -crf 24 -vf scale=1080:-2,setsar=1:1 attempt/vid_out.m3u8

pub struct StreamActor;

impl Actor for StreamActor {
    type Context = Context<Self>;
}

pub struct ProcessMessage(pub String, pub Child);

impl Message for ProcessMessage {
    type Result = Result<ExitStatus>;
}

impl Handler<ProcessMessage> for StreamActor {
    type Result = Result<ExitStatus>;

    fn handle(&mut self, msg: ProcessMessage, _ctx: &mut Self::Context) -> Self::Result {
        trace!("Message received");
        let mut process = msg.1;
        let result = process.wait();

        debug!(
            "Finished waiting for: {:?}. Code: {:?}",
            msg.0,
            result
                .as_ref()
                .map_or(-1, |status| status.code().unwrap_or(-1))
        );
        result
    }
}

pub async fn create_playlist(video_path: &str, playlist_file: &str) -> Result<Child> {
    if Path::new(playlist_file).exists() {
        debug!("Playlist already exists: {}", playlist_file);
        return Err(std::io::Error::from(std::io::ErrorKind::AlreadyExists));
    }

    let result = Command::new("ffmpeg")
        .arg("-i")
        .arg(video_path)
        .arg("-c:v")
        .arg("h264")
        .arg("-crf")
        .arg("21")
        .arg("-preset")
        .arg("veryfast")
        .arg("-hls_time")
        .arg("3")
        .arg("-hls_list_size")
        .arg("100")
        .arg("-vf")
        .arg("scale=1080:-2,setsar=1:1")
        .arg(playlist_file)
        .stdout(Stdio::null())
        .stderr(Stdio::null())
        .spawn();

    let start_time = std::time::Instant::now();
    loop {
        actix::clock::sleep(std::time::Duration::from_secs(1)).await;

        if Path::new(playlist_file).exists()
            || std::time::Instant::now() - start_time > std::time::Duration::from_secs(5)
        {
            break;
        }
    }

    result
}

pub fn generate_video_thumbnail(path: &Path, destination: &Path) {
    Command::new("ffmpeg")
        .arg("-ss")
        .arg("3")
        .arg("-i")
        .arg(path.to_str().unwrap())
        .arg("-vframes")
        .arg("1")
        .arg("-f")
        .arg("image2")
        .arg(destination)
        .output()
        .expect("Failure to create video frame");
}

pub struct VideoPlaylistManager {
    playlist_dir: PathBuf,
    playlist_generator: Addr<PlaylistGenerator>,
}

impl VideoPlaylistManager {
    pub fn new<P: Into<PathBuf>>(
        playlist_dir: P,
        playlist_generator: Addr<PlaylistGenerator>,
    ) -> Self {
        Self {
            playlist_dir: playlist_dir.into(),
            playlist_generator,
        }
    }
}

impl Actor for VideoPlaylistManager {
    type Context = Context<Self>;
}

impl Handler<ScanDirectoryMessage> for VideoPlaylistManager {
    type Result = ResponseFuture<()>;

    fn handle(&mut self, msg: ScanDirectoryMessage, _ctx: &mut Self::Context) -> Self::Result {
        let start = std::time::Instant::now();
        info!(
            "Starting scan directory for video playlist generation: {}",
            msg.directory
        );

        let video_files = WalkDir::new(&msg.directory)
            .into_iter()
            .filter_map(|e| e.ok())
            .filter(|e| e.file_type().is_file())
            .filter(is_video)
            .collect::<Vec<DirEntry>>();

        let scan_dir_name = msg.directory.clone();
        let playlist_output_dir = self.playlist_dir.clone();
        let playlist_generator = self.playlist_generator.clone();

        Box::pin(async move {
            for e in video_files {
                let path = e.path();
                let path_as_str = path.to_str().unwrap();
                debug!(
                    "Sending generate playlist message for path: {}",
                    path_as_str
                );

                match playlist_generator
                    .send(GeneratePlaylistMessage {
                        playlist_path: playlist_output_dir.to_str().unwrap().to_string(),
                        video_path: PathBuf::from(path),
                    })
                    .await
                    .expect("Failed to send generate playlist message")
                {
                    Ok(_) => {
                        debug!(
                            "Successfully generated playlist for file: '{}'",
                            path_as_str
                        );
                    }
                    Err(e) => {
                        warn!("Failed to generate playlist for path '{:?}'. {:?}", path, e);
                    }
                }
            }

            info!(
                "Finished directory scan of '{}' in {:?}",
                scan_dir_name,
                start.elapsed()
            );
        })
    }
}

#[derive(Message)]
#[rtype(result = "()")]
pub struct ScanDirectoryMessage {
    pub(crate) directory: String,
}

#[derive(Message)]
#[rtype(result = "Result<()>")]
struct GeneratePlaylistMessage {
    video_path: PathBuf,
    playlist_path: String,
}

pub struct PlaylistGenerator {
    semaphore: Arc<Semaphore>,
}

impl PlaylistGenerator {
    pub(crate) fn new() -> Self {
        PlaylistGenerator {
            semaphore: Arc::new(Semaphore::new(2)),
        }
    }
}

impl Actor for PlaylistGenerator {
    type Context = Context<Self>;
}

impl Handler<GeneratePlaylistMessage> for PlaylistGenerator {
    type Result = ResponseFuture<Result<()>>;

    fn handle(&mut self, msg: GeneratePlaylistMessage, _ctx: &mut Self::Context) -> Self::Result {
        let video_file = msg.video_path.to_str().unwrap().to_owned();
        let playlist_path = msg.playlist_path.as_str().to_owned();
        let semaphore = self.semaphore.clone();

        let playlist_file = format!(
            "{}/{}.m3u8",
            playlist_path,
            msg.video_path.file_name().unwrap().to_str().unwrap()
        );
        Box::pin(async move {
            let wait_start = std::time::Instant::now();
            let permit = semaphore
                .acquire_owned()
                .await
                .expect("Unable to acquire semaphore");

            debug!(
                "Waited for {:?} before starting ffmpeg",
                wait_start.elapsed()
            );

            if Path::new(&playlist_file).exists() {
                debug!("Playlist already exists: {}", playlist_file);
                return Err(std::io::Error::from(std::io::ErrorKind::AlreadyExists));
            }

            tokio::spawn(async move {
                let ffmpeg_result = tokio::process::Command::new("ffmpeg")
                    .arg("-i")
                    .arg(&video_file)
                    .arg("-c:v")
                    .arg("h264")
                    .arg("-crf")
                    .arg("21")
                    .arg("-preset")
                    .arg("veryfast")
                    .arg("-hls_time")
                    .arg("3")
                    .arg("-hls_list_size")
                    .arg("100")
                    .arg("-vf")
                    .arg("scale=1080:-2,setsar=1:1")
                    .arg(playlist_file)
                    .stdout(Stdio::null())
                    .stderr(Stdio::piped())
                    .output()
                    .inspect_err(|e| error!("Failed to run ffmpeg on child process: {}", e))
                    .map_err(|e| std::io::Error::new(std::io::ErrorKind::Other, e.to_string()))
                    .await;

                // Hang on to the permit until we're done decoding and then explicitly drop
                drop(permit);

                if let Ok(ref res) = ffmpeg_result {
                    debug!("ffmpeg output: {:?}", res);
                }

                ffmpeg_result
            });

            Ok(())
        })
    }
}