277 lines
8.1 KiB
Rust
277 lines
8.1 KiB
Rust
extern crate dotenvy;
|
|
extern crate lastfm;
|
|
|
|
use serenity::async_trait;
|
|
use serenity::builder::{CreateAttachment, CreateMessage};
|
|
use serenity::model::channel::Message;
|
|
use serenity::model::id::UserId;
|
|
use serenity::prelude::*;
|
|
|
|
extern crate magick_rust;
|
|
use magick_rust::{ColorspaceType, CompositeOperator, MagickWand, PixelWand, magick_wand_genesis};
|
|
|
|
use std::env;
|
|
use std::sync::Once;
|
|
|
|
extern crate sqlx;
|
|
use sqlx::sqlite::SqlitePool;
|
|
|
|
static START: Once = Once::new();
|
|
|
|
struct Handler;
|
|
|
|
enum Reply {
|
|
Image(Vec<u8>),
|
|
Text(String),
|
|
}
|
|
|
|
macro_rules! handle_magick_result {
|
|
($a: expr, $b: literal) => {
|
|
match $a {
|
|
Ok(_) => {}
|
|
Err(e) => return Reply::Text(format!("Error: {} {}", $b, e)),
|
|
}
|
|
};
|
|
}
|
|
|
|
macro_rules! handle_magick_option {
|
|
($a: expr, $b: literal) => {
|
|
match $a {
|
|
Some(res) => res,
|
|
None => return Reply::Text($b.to_string()),
|
|
}
|
|
};
|
|
}
|
|
|
|
#[allow(unused)]
|
|
async fn fmi(ctx: &Context, arg: &str, id: UserId) -> Reply {
|
|
let lastfm_user = match arg {
|
|
"" => get_lastfm_username(ctx, id).await,
|
|
_ => Some(arg.to_string()),
|
|
};
|
|
let client = match lastfm_user {
|
|
Some(s) => lastfm::Client::<String, String>::from_env(s),
|
|
None => return Reply::Text("No last.fm username set.".to_string()),
|
|
};
|
|
let now_playing = match client.now_playing().await {
|
|
Ok(np) => np,
|
|
Err(e) => return Reply::Text(format!("Error: grabbing last.fm user data failed {e}")),
|
|
};
|
|
let track = match now_playing {
|
|
Some(track) => track,
|
|
None => return Reply::Text("Nothing playing.".to_string()),
|
|
};
|
|
let image_uri = match track.image.extralarge {
|
|
Some(iu) => iu,
|
|
None => return Reply::Text("Error: getting image uri failed".to_string()),
|
|
};
|
|
let mut base_color = PixelWand::new();
|
|
let mut white = PixelWand::new();
|
|
let mut main_wand = MagickWand::new();
|
|
let mut art_wand = MagickWand::new();
|
|
let mut art_wand_cluster = MagickWand::new();
|
|
let mut mask_wand = MagickWand::new();
|
|
|
|
handle_magick_result!(
|
|
base_color.set_color("#7f7f7f"),
|
|
"Failed to set init base color"
|
|
);
|
|
handle_magick_result!(
|
|
art_wand.read_image(image_uri.as_str()),
|
|
"Failed to read image from uri"
|
|
);
|
|
handle_magick_result!(
|
|
art_wand.transform_image_colorspace(ColorspaceType::Lab),
|
|
"Failed to set art colorspace"
|
|
);
|
|
|
|
let (art_width, art_height) = (art_wand.get_image_width(), art_wand.get_image_height());
|
|
handle_magick_result!(
|
|
art_wand_cluster.new_image(art_width, art_height, &base_color),
|
|
"Failed to init cluster image"
|
|
);
|
|
handle_magick_result!(
|
|
art_wand_cluster.compose_images(&art_wand, CompositeOperator::SrcOver, true, 0, 0),
|
|
"Failed to copy onto cluster image"
|
|
);
|
|
// least stupid way I could figure out to do this, since it seems that MagickWand::new_from_wand() doesn't work
|
|
handle_magick_result!(
|
|
art_wand_cluster.set_image_colorspace(ColorspaceType::Lab),
|
|
"Failed to set cluster colorspace"
|
|
);
|
|
handle_magick_result!(
|
|
art_wand_cluster.kmeans(5, 200, 0.001),
|
|
"Failed to run kmeans"
|
|
);
|
|
let mut colors = handle_magick_option!(
|
|
art_wand_cluster.get_image_histogram(),
|
|
"Failed to get color histogram"
|
|
);
|
|
colors.sort_by_cached_key(|color| -(color.get_color_count() as isize));
|
|
let mut accent_color = match colors.len() {
|
|
1 => &colors[0],
|
|
_ => &colors[1],
|
|
};
|
|
|
|
handle_magick_result!(white.set_color("#ffffff"), "Failed to set init white");
|
|
handle_magick_result!(mask_wand.read_image("mask.png"), "Failed to load mask");
|
|
handle_magick_result!(
|
|
mask_wand.opaque_paint_image(&white, accent_color, 0.0, false),
|
|
"Failed to paint mask"
|
|
);
|
|
|
|
handle_magick_result!(
|
|
main_wand.new_image(548, 147, &colors[0]),
|
|
"Failed to create wand"
|
|
);
|
|
handle_magick_result!(
|
|
main_wand.transform_image_colorspace(ColorspaceType::Lab),
|
|
"Failed to set main colorspace"
|
|
);
|
|
handle_magick_result!(
|
|
art_wand.adaptive_resize_image(124, 124),
|
|
"Failed to resize art"
|
|
);
|
|
handle_magick_result!(
|
|
art_wand_cluster.adaptive_resize_image(124, 124),
|
|
"Failed to resize art_cluster"
|
|
);
|
|
handle_magick_result!(
|
|
main_wand.compose_images(&art_wand, CompositeOperator::SrcOver, false, 12, 12),
|
|
"Failed to combine images"
|
|
);
|
|
handle_magick_result!(
|
|
main_wand.compose_images(
|
|
&art_wand_cluster,
|
|
CompositeOperator::SrcOver,
|
|
false,
|
|
160,
|
|
12
|
|
),
|
|
"Failed to combine debug images"
|
|
);
|
|
handle_magick_result!(
|
|
main_wand.compose_images(&mask_wand, CompositeOperator::SrcOver, false, 401, 0),
|
|
"Failed to combine mask"
|
|
);
|
|
Reply::Image(main_wand.write_image_blob("png").unwrap())
|
|
}
|
|
|
|
async fn set(ctx: &Context, arg: &str, id: UserId) -> Reply {
|
|
set_lastfm_username(ctx, id, arg.to_string()).await;
|
|
Reply::Text(format!("set user {}", arg))
|
|
}
|
|
|
|
#[async_trait]
|
|
impl EventHandler for Handler {
|
|
async fn message(&self, ctx: Context, msg: Message) {
|
|
let mut cmd_iter = msg.content.split(" ");
|
|
let cmd = cmd_iter.next().unwrap_or("");
|
|
let arg = cmd_iter.next().unwrap_or("");
|
|
let resp = match cmd {
|
|
".fmi" => Some(fmi(&ctx, arg, msg.author.id).await),
|
|
".set" => Some(set(&ctx, arg, msg.author.id).await),
|
|
_ => None,
|
|
};
|
|
if let Some(reply) = resp {
|
|
let m = match reply {
|
|
Reply::Image(img) => {
|
|
let file = CreateAttachment::bytes(img, "fmi.png");
|
|
CreateMessage::new().add_file(file)
|
|
}
|
|
Reply::Text(txt) => CreateMessage::new().content(txt),
|
|
};
|
|
|
|
let message = msg.channel_id.send_message(&ctx.http, m).await;
|
|
if let Err(why) = message {
|
|
println!("Error sending message: {why:?}");
|
|
}
|
|
}
|
|
}
|
|
}
|
|
|
|
struct PoolContainer;
|
|
|
|
impl TypeMapKey for PoolContainer {
|
|
type Value = SqlitePool;
|
|
}
|
|
|
|
struct DBResponse {
|
|
lastfm_username: String,
|
|
}
|
|
|
|
async fn get_lastfm_username(ctx: &Context, id: UserId) -> Option<String> {
|
|
let data = ctx.data.write().await;
|
|
let pool = data
|
|
.get::<PoolContainer>()
|
|
.expect("Failed to get pool container");
|
|
let id = i64::from(id);
|
|
let resp = sqlx::query_as!(
|
|
DBResponse,
|
|
r#"
|
|
SELECT lastfm_username
|
|
FROM users
|
|
WHERE discord_userid = ?1
|
|
"#,
|
|
id
|
|
)
|
|
.fetch_one(pool)
|
|
.await;
|
|
match resp {
|
|
Ok(r) => Some(r.lastfm_username),
|
|
Err(_) => None,
|
|
}
|
|
}
|
|
|
|
async fn set_lastfm_username(ctx: &Context, id: UserId, user: String) {
|
|
let data = ctx.data.write().await;
|
|
let pool = data
|
|
.get::<PoolContainer>()
|
|
.expect("Failed to get pool container");
|
|
let id = i64::from(id);
|
|
sqlx::query!(
|
|
r#"
|
|
INSERT INTO users
|
|
VALUES (?1, ?2)
|
|
"#,
|
|
id, user
|
|
)
|
|
.execute(pool)
|
|
.await.expect("Failed to insert");
|
|
}
|
|
|
|
#[tokio::main]
|
|
async fn main() {
|
|
START.call_once(|| {
|
|
dotenvy::dotenv().expect("Failed to load .env");
|
|
magick_wand_genesis();
|
|
});
|
|
|
|
let db_url = env::var("DATABASE_URL").expect("Failed to load DATABASE_URL");
|
|
let pool = SqlitePool::connect(&db_url)
|
|
.await
|
|
.expect("Failed to load db");
|
|
|
|
let token = env::var("DISCORD_TOKEN").expect("Failed to load DISCORD_TOKEN");
|
|
// Set gateway intents, which decides what events the bot will be notified about
|
|
let intents = GatewayIntents::GUILD_MESSAGES
|
|
| GatewayIntents::DIRECT_MESSAGES
|
|
| GatewayIntents::MESSAGE_CONTENT;
|
|
|
|
let mut client = Client::builder(&token, intents)
|
|
.event_handler(Handler)
|
|
.await
|
|
.expect("Err creating client");
|
|
|
|
{
|
|
let mut data = client.data.write().await;
|
|
data.insert::<PoolContainer>(pool);
|
|
}
|
|
|
|
if let Err(why) = client.start().await {
|
|
println!("Client error: {why:?}");
|
|
} else {
|
|
println!("watcat started");
|
|
}
|
|
}
|