FrenBot/src/discord/album.rs

use crate::album_manager::{Album, Image};
use crate::discord::admin::is_admin;
use crate::error::Error;
use crate::{command, group, GlobalData};
use j_db::model::JdbModel;
use serenity::client::Context;
use serenity::framework::standard::{Args, CommandResult};
use serenity::model::channel::Message;

#[group]
#[commands(add_image, remove_album, list_albums, import_from_file)]
pub struct AlbumCmd;

#[command]
#[only_in(guilds)]
#[min_args(1)]
#[description("Add an image to an album")]
#[usage("<album_name>")]
async fn add_image(ctx: &Context, msg: &Message, mut args: Args) -> CommandResult {
    let album_name = args.parse::<String>()?;
    args.advance();
    let tags: Vec<String> = args.rest().split(',').map(|s| s.to_string()).collect();

    let mut data = ctx.data.write().await;

    let global_data = data.get_mut::<GlobalData>().unwrap();

    let album = Album::find_album_by_name_or_alias(&global_data.db, &album_name).unwrap();

    let mut album = if let Some(album) = album {
        album
    } else {
        Album::add_album(&global_data.db, &album_name).unwrap()
    };

    for attachment in &msg.attachments {
        let data = attachment.download().await.unwrap();
        let img = Image::new(
            &data,
            &attachment.filename,
            global_data.cfg.img_path.clone(),
            tags.clone(),
        )
        .await;
        album.images.push(img);
    }

    global_data.db.insert(album).unwrap();

    let plural = if msg.attachments.len() > 1 { "s" } else { "" };

    msg.reply(
        &ctx.http,
        format!("Image{} added to {}!", plural, album_name),
    )
    .await?;

    Ok(())
}

#[command]
#[only_in(guilds)]
#[min_args(1)]
#[description("Bulk add images")]
#[usage("<path>")]
async fn import_from_file(ctx: &Context, msg: &Message, mut args: Args) -> CommandResult {
    let album_name = args.parse::<String>()?;
    args.advance();
    let album_path = args.parse::<String>()?;
    let album_path = album_path.replace('"', "");

    let mut data = ctx.data.write().await;

    let global_data = data.get_mut::<GlobalData>().unwrap();

    if !is_admin(&msg.author.id, &global_data.cfg) {
        return Ok(());
    }

    let album = Album::find_album_by_name_or_alias(&global_data.db, &album_name).unwrap();

    let mut album = if let Some(album) = album {
        album
    } else {
        Album::add_album(&global_data.db, &album_name).unwrap()
    };

    let mut dir = tokio::fs::read_dir(album_path).await.unwrap();

    while let Some(entry) = dir.next_entry().await? {
        let data = tokio::fs::read(entry.path()).await?;
        let img = Image::new(
            &data,
            entry.path().file_name().unwrap().to_str().unwrap(),
            global_data.cfg.img_path.clone(),
            vec![],
        )
        .await;
        album.images.push(img);
    }

    global_data.db.insert(album).unwrap();

    let plural = if msg.attachments.len() > 1 { "s" } else { "" };

    msg.reply(
        &ctx.http,
        format!("Image{} added to {}!", plural, album_name),
    )
    .await?;

    Ok(())
}

#[command]
#[only_in(guilds)]
#[max_args(1)]
#[description("Remove an imgur album command.")]
#[usage("<name>")]
async fn remove_album(ctx: &Context, msg: &Message, args: Args) -> CommandResult {
    let album_name = args.parse::<String>()?;

    let mut data = ctx.data.write().await;

    let global_data = data.get_mut::<GlobalData>().unwrap();

    let album = Album::find_album_by_name_or_alias(&global_data.db, &album_name).unwrap();

    match album {
        None => {
            msg.reply(
                &ctx.http,
                "I already got rid of that one, or I never had it. Who knows!",
            )
            .await?;
        }
        Some(album) => {
            for img in &album.images {
                tokio::fs::remove_file(img.full_path(global_data.cfg.img_path.clone()))
                    .await
                    .unwrap();
            }

            global_data.db.remove::<Album>(album.id().unwrap())?;

            msg.reply(&ctx.http, format!("{} album removed!", album_name))
                .await?;
        }
    }

    Ok(())
}

#[command]
#[aliases("albums")]
#[description("List all album commands.")]
#[only_in(guilds)]
async fn list_albums(ctx: &Context, msg: &Message, _args: Args) -> CommandResult {
    let data = ctx.data.read().await;

    let global_data = data.get::<GlobalData>().unwrap();

    let album_names: Vec<String> = global_data
        .db
        .filter(|_, _: &Album| true)?
        .map(|album| album.album_name)
        .collect();

    if album_names.is_empty() {
        msg.reply(&ctx.http, "There are no albums configured!")
            .await?;
    } else {
        msg.reply(
            &ctx.http,
            format!("**Albums**:\n{}", album_names.join("\n")),
        )
        .await?;
    }

    Ok(())
}

pub async fn parse_album(
    ctx: &Context,
    msg: &Message,
    album_name: &str,
    tags: Vec<&str>,
) -> Result<bool, Error> {
    let data = ctx.data.read().await;
    let global_data = data.get::<GlobalData>().unwrap();

    let img = match Album::get_random_image(&global_data.db, album_name, tags) {
        Ok(img) => img,
        Err(err) => {
            return match err {
                Error::NoAlbumFound => Ok(false),
                _ => Err(err),
            }
        }
    };

    if let Some(img) = img {
        msg.reply(&ctx.http, img.link(&global_data.cfg.base_url))
            .await?;
    } else {
        msg.reply(&ctx.http, "No image found :(").await?;
    }

    Ok(true)
}