novel-cli 0.17.0

use std::env;
use std::path::{Path, PathBuf};
use std::sync::Arc;

use bytes::BytesMut;
use clap::Args;
use color_eyre::eyre::{self, Result};
use fluent_templates::Loader;
use image::ImageReader;
use tokio::fs::File;
use tokio::io::{AsyncReadExt, BufReader};
use tokio::process::Command;
use tokio::sync::Semaphore;
use walkdir::WalkDir;

use crate::utils::{self, ProgressBar};
use crate::{LANG_ID, LOCALES};

#[must_use]
#[derive(Args)]
#[command(about = LOCALES.lookup(&LANG_ID, "real_cugan_command"))]
pub struct RealCugan {
    #[arg(help = LOCALES.lookup(&LANG_ID, "image_path"))]
    pub image_path: Option<PathBuf>,

    #[arg(short, long, default_value_t = utils::maximum_concurrency(),
        help = LOCALES.lookup(&LANG_ID, "maximum_concurrency"))]
    pub maximum_concurrency: usize,
}

pub async fn execute(config: RealCugan) -> Result<()> {
    utils::ensure_executable_exists("realcugan-ncnn-vulkan")?;

    let mut handles = Vec::new();
    let mut to_delete = Vec::new();

    let semaphore = Arc::new(Semaphore::new(config.maximum_concurrency));
    let image_paths = image_paths(config).await?;
    let pb = ProgressBar::new(image_paths.len() as u64)?;

    let curr_path = env::current_dir()?;

    for input_path in image_paths {
        let image = ImageReader::open(&input_path)?
            .with_guessed_format()?
            .decode()?;
        let scale = calc_scale(image.width(), image.height());

        let ext = utils::new_image_ext(&image);
        if let Err(error) = ext {
            tracing::error!("{error}: {}", input_path.display());
            continue;
        }

        let output_path = input_path.with_extension(ext.unwrap());

        if input_path != output_path {
            to_delete.push(input_path.clone());
        }

        tracing::info!(
            "Run realcugan-ncnn-vulkan with {}, {}x{}, scale: {}",
            input_path.display(),
            image.width(),
            image.height(),
            scale
        );

        let permit = semaphore.clone().acquire_owned().await.unwrap();
        let mut pb = pb.clone();

        let absolute_path = dunce::canonicalize(&input_path)?;
        let diff_path = pathdiff::diff_paths(absolute_path, &curr_path).unwrap();

        handles.push(tokio::spawn(async move {
            pb.inc(diff_path.display().to_string(), 1)?;
            create_child(input_path, output_path, scale).await?;

            drop(permit);

            eyre::Ok(())
        }));
    }

    for handle in handles {
        handle.await??;
    }

    for path in to_delete {
        utils::remove_file_or_dir(path)?;
    }

    pb.finish()?;

    Ok(())
}

async fn image_paths(config: RealCugan) -> Result<Vec<PathBuf>> {
    let image_path = if let Some(image_path) = config.image_path {
        image_path
    } else {
        env::current_dir()?
    };

    let mut result = Vec::new();

    for entry in WalkDir::new(image_path).max_depth(1) {
        let input_path = entry?.path().to_path_buf();

        if is_image(&input_path).await? {
            result.push(input_path);
        }
    }

    if result.is_empty() {
        eyre::bail!("There is no image in this directory");
    }

    Ok(result)
}

async fn is_image<T>(path: T) -> Result<bool>
where
    T: AsRef<Path>,
{
    if !path.as_ref().is_file() {
        return Ok(false);
    }

    let file = File::open(path).await?;
    let mut reader = BufReader::new(file);
    let mut buffer = BytesMut::with_capacity(128);

    reader.read_buf(&mut buffer).await?;

    Ok(infer::is_image(&buffer))
}

// 5k: 5120*2880=  14745600
// 4k: 3840*2160=   8294400
// 2k: 2560*1440=   3686400
// 1080p: 1920×1080=2073600
#[must_use]
#[inline]
const fn calc_scale(width: u32, height: u32) -> u8 {
    let pixel = width * height;
    let n = (5120 * 2880 / pixel) as u8;

    if n >= 16 {
        4
    } else if n >= 9 {
        3
    } else {
        2
    }
}

async fn create_child<T, E>(input_path: T, output_path: E, scale: u8) -> Result<()>
where
    T: AsRef<Path>,
    E: AsRef<Path>,
{
    let output = Command::new("realcugan-ncnn-vulkan")
        .arg("-i")
        .arg(input_path.as_ref())
        .arg("-o")
        .arg(output_path.as_ref())
        .arg("-s")
        .arg(scale.to_string())
        .output()
        .await?;

    tracing::info!("{}", simdutf8::basic::from_utf8(&output.stdout)?);

    if !output.status.success() {
        tracing::error!("{}", simdutf8::basic::from_utf8(&output.stderr)?);
        eyre::bail!("`realcugan-ncnn-vulkan` failed to execute");
    }

    Ok(())
}