glimbus/src/main.rs

use axum::body::Bytes;
use axum::extract::{Multipart, Query};
use axum::http::StatusCode;
use axum::response::{IntoResponse, Response};
use axum::routing::post;
use axum::{Json, Router};
use axum_macros::debug_handler;
use filepath::FilePath;
use reqwest::header::WWW_AUTHENTICATE;
use serde::__private::de::FlatInternallyTaggedAccess;
use serde::{Deserialize, Serialize};
use std::collections::HashSet;
use std::io::Read;
use std::io::Write;

#[tokio::main]
async fn main() {
    pretty_env_logger::init();
    // build our application with a single route
    let app = Router::new().route("/", post(upload_file));

    // run it with hyper on localhost:3000
    axum::Server::bind(&"0.0.0.0:6679".parse().unwrap())
        .serve(app.into_make_service())
        .await
        .unwrap();
}

#[derive(Deserialize)]
struct Options {
    threshold: String,
}

#[derive(Serialize, Deserialize)]
#[serde(untagged)]
enum WrappedResponse {
    Tags(Vec<String>),
    Error(String),
}

// Make our own error that wraps `anyhow::Error`.
struct AppError(anyhow::Error);

// Tell axum how to convert `AppError` into a response.
impl IntoResponse for AppError {
    fn into_response(self) -> Response {
        println!("amogus");
        (
            StatusCode::INTERNAL_SERVER_ERROR,
            format!("Something went wrong: {}", self.0),
        )
            .into_response()
    }
}

// This enables using `?` on functions that return `Result<_, anyhow::Error>` to turn them into
// `Result<_, AppError>`. That way you don't need to do that manually.
impl<E> From<E> for AppError
where
    E: Into<anyhow::Error>,
{
    fn from(err: E) -> Self {
        Self(err.into())
    }
}

fn try_thing() -> Result<(), anyhow::Error> {
    anyhow::bail!("it failed!")
}

async fn test_handler() -> Result<(), AppError> {
    try_thing()?;
    Ok(())
}

async fn send_image_to_dd(
    file_contents: Vec<u8>,
    file_name: String,
    file_mime_type: &str,
    options: &Options,
) -> anyhow::Result<WrappedResponse> {
    let part = reqwest::multipart::Part::bytes(file_contents)
        .file_name(file_name)
        .mime_str(file_mime_type)
        .unwrap();
    let form = reqwest::multipart::Form::new().part("file", part);

    log::debug!("calling dd");

    let resp = reqwest::Client::new()
        .post("http://localhost:4443")
        .multipart(form)
        .header("authorization", "Bearer 123")
        .query(&[("threshold", options.threshold.clone())])
        .send()
        .await?;

    let body = resp.text().await?;
    log::info!("body: {}", &body);
    let json_response: WrappedResponse = serde_json::from_str(&body)?;

    log::debug!("called!");
    Ok(json_response)
}

use std::process::Stdio;

use ffmpeg_cli::Parameter;

use futures_util::{future::ready, StreamExt};

async fn fetch_frame_as_image(
    input_path: &str,
    output_path: &str,
    frame_index: usize,
) -> anyhow::Result<()> {
    let frame_index_param = format!("select=eq(n\\,{})", frame_index);
    let builder = ffmpeg_cli::FfmpegBuilder::new()
        .stderr(Stdio::piped())
        .option(Parameter::Single("nostdin"))
        .option(Parameter::Single("y"))
        .input(ffmpeg_cli::File::new(input_path))
        .output(
            ffmpeg_cli::File::new(output_path)
                .option(Parameter::KeyValue("vf", &frame_index_param))
                .option(Parameter::KeyValue("vframes", "1")),
        );

    log::debug!("running {:?}", builder);
    let ffmpeg = builder.run().await.unwrap();

    log::debug!("run");
    ffmpeg
        .progress
        .for_each(|x| {
            dbg!(&x);
            // lmao
            //         x.unwrap();
            ready(())
        })
        .await;

    log::debug!("run");
    let output = ffmpeg.process.wait_with_output().unwrap();

    log::debug!("out");
    println!(
        "{}\nstderr:\n{}",
        output.status,
        std::str::from_utf8(&output.stderr).unwrap()
    );

    Ok(())
}

#[debug_handler]
async fn upload_file(
    options: Query<Options>,
    mut multipart: Multipart,
) -> Result<(StatusCode, Json<WrappedResponse>), AppError> {
    let mut maybe_file_contents: Option<axum::body::Bytes> = None;
    let mut maybe_file_type: Option<String> = None;
    let mut maybe_file_name: Option<String> = None;

    while let Some(field) = multipart.next_field().await.unwrap() {
        let name = field.name().unwrap().to_string();
        let content_type = field.content_type().unwrap().to_string();
        let filename = field.file_name().unwrap().to_string();
        let data = field.bytes().await.unwrap();

        log::info!("Length of `{}` is {} bytes", name, data.len());
        if name == "file" {
            maybe_file_contents = Some(data);
            maybe_file_type = Some(content_type);
            maybe_file_name = Some(filename);
        }
    }

    if let Some(file_contents) = maybe_file_contents {
        let file_type = maybe_file_type.unwrap();
        let file_name = maybe_file_name.unwrap();
        let is_video = file_type.starts_with("video/") || file_name.ends_with(".mp4");
        if is_video {
            let mut final_tag_set = HashSet::new();

            let mut temp_file = tempfile::NamedTempFile::new()?;
            temp_file.write_all(&file_contents.to_vec())?;

            log::debug!("tmp path: {:?}", temp_file.path());

            let info = ffprobe::ffprobe(temp_file.path())?;
            let stream = info.streams.get(0).unwrap();

            let all_frames = stream.nb_frames.clone().unwrap().parse::<u32>()?;
            let frame_rate_str = stream.r_frame_rate.clone();

            let parts = frame_rate_str.split("/").into_iter().collect::<Vec<_>>();
            let frame_rate =
                parts.get(0).unwrap().parse::<u32>()? / parts.get(1).unwrap().parse::<u32>()?;

            let wanted_frame_skip = frame_rate.try_into().unwrap(); // every second

            let temporary_frame_dir = tempfile::tempdir()?;

            let temporary_frame_path =
                format!("{}/frame.png", temporary_frame_dir.path().to_string_lossy());
            log::info!("path:{}", &temporary_frame_path);

            for frame_number in (0..all_frames).step_by(wanted_frame_skip) {
                log::info!("extracting frame {:?}", frame_number);
                fetch_frame_as_image(
                    temp_file.path().to_str().unwrap(),
                    &temporary_frame_path,
                    frame_number.try_into().unwrap(),
                )
                .await?;

                let mut actual_frame_file = std::fs::File::open(&temporary_frame_path)?;
                let mut frame_data = vec![];
                actual_frame_file.read_to_end(&mut frame_data)?;
                let tags_from_frame = if let WrappedResponse::Tags(tags_from_frame) =
                    send_image_to_dd(frame_data, "amongus.png".to_string(), "image/png", &options)
                        .await?
                {
                    tags_from_frame
                } else {
                    todo!()
                };

                for tag in tags_from_frame {
                    final_tag_set.insert(tag);
                }
            }

            let response = WrappedResponse::Tags(final_tag_set.into_iter().collect::<Vec<_>>());
            Ok((StatusCode::OK, Json(response)))
        } else {
            let json_response =
                send_image_to_dd(file_contents.to_vec(), file_name, &file_type, &options).await?;
            Ok((StatusCode::OK, Json(json_response)))
        }
    } else {
        Ok((
            StatusCode::INTERNAL_SERVER_ERROR,
            Json(WrappedResponse::Error(
                "no file found in request".to_string(),
            )),
        ))
    }
}