582 lines
18 KiB
Rust
582 lines
18 KiB
Rust
extern crate dotenvy;
|
|
extern crate lastfm;
|
|
extern crate reqwest;
|
|
extern crate reqwest_middleware;
|
|
use http_cache_reqwest::{CACacheManager, Cache, CacheMode, HttpCache, HttpCacheOptions};
|
|
|
|
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, FilterType};
|
|
|
|
use std::env;
|
|
use std::sync::Once;
|
|
|
|
use palette::color_difference::ImprovedCiede2000;
|
|
use palette::lab::Lab;
|
|
use palette::white_point::D65;
|
|
|
|
extern crate sqlx;
|
|
use sqlx::sqlite::SqlitePool;
|
|
|
|
use imagetext::fontdb::FontDB;
|
|
use imagetext::superfont::SuperFont;
|
|
|
|
mod text;
|
|
|
|
static START: Once = Once::new();
|
|
|
|
struct Handler;
|
|
|
|
enum Reply {
|
|
Image(Vec<u8>),
|
|
Text(String),
|
|
}
|
|
|
|
macro_rules! now {
|
|
() => {
|
|
std::time::SystemTime::now()
|
|
.duration_since(std::time::UNIX_EPOCH)
|
|
.unwrap()
|
|
.as_secs_f32() as usize
|
|
};
|
|
}
|
|
|
|
macro_rules! log {
|
|
($a:expr) => {
|
|
if env::var("DEBUG").unwrap_or("0".to_string()) == "1" {
|
|
println!(concat!("debug: {} ", $a), now!())
|
|
}
|
|
};
|
|
|
|
($a:expr, $($b:expr),+) => {
|
|
if env::var("DEBUG").unwrap_or("0".to_string()) == "1" {
|
|
println!(concat!("debug: {} ", $a), now!(), $($b),+)
|
|
}
|
|
};
|
|
}
|
|
|
|
macro_rules! handle_magick_result {
|
|
($a: expr, $b: literal) => {
|
|
match $a {
|
|
Ok(_) => {
|
|
log!("{} run successfully", stringify!($a));
|
|
}
|
|
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()),
|
|
}
|
|
};
|
|
}
|
|
|
|
// this function is almost assuredly the least idiomatic rust code I will ever write.
|
|
// but until I know what the "correct" way is, it'll have to do
|
|
fn parse_cielab(s: String) -> Result<(f32, f32, f32), &'static str> {
|
|
if s.len() < 13 || &s[0..7] != "cielab(" || s.chars().nth(s.len() - 1).unwrap() != ')' {
|
|
return Err("Error parsing cielab string");
|
|
}
|
|
let vals = s[7..(s.len() - 1)]
|
|
.split(",")
|
|
.map(|x| match x.parse() {
|
|
Ok(f) => f,
|
|
Err(_) => f32::NAN,
|
|
})
|
|
.collect::<Vec<f32>>();
|
|
if vals.iter().any(|x| x.is_nan()) {
|
|
return Err("Error parsing cielab string");
|
|
}
|
|
match &vals[..] {
|
|
&[a, b, c] => Ok((a, b, c)),
|
|
_ => Err("Error parsing cielab string"),
|
|
}
|
|
}
|
|
|
|
fn validate_color(col: &PixelWand) -> Option<Lab> {
|
|
let color_str = match col.get_color_as_string() {
|
|
Ok(s) => s,
|
|
Err(_) => return None,
|
|
};
|
|
let color_raw = match parse_cielab(color_str) {
|
|
Ok(f) => f,
|
|
Err(_) => return None,
|
|
};
|
|
Some(Lab::<D65>::from_components(color_raw))
|
|
}
|
|
|
|
async fn art(ctx: &Context, arg: &str, id: UserId) -> Reply {
|
|
let lastfm_user = match arg {
|
|
"" => get_lastfm_username(ctx, id).await,
|
|
_ => Some(arg.to_string()),
|
|
};
|
|
let lastfm_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 lastfm_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 track_art = match track.image.extralarge {
|
|
Some(track_art) => track_art,
|
|
None => return Reply::Text("Error: getting image uri failed".to_string()),
|
|
};
|
|
Reply::Text(track_art.to_string())
|
|
}
|
|
|
|
async fn url(ctx: &Context, arg: &str, id: UserId) -> Reply {
|
|
let lastfm_user = match arg {
|
|
"" => get_lastfm_username(ctx, id).await,
|
|
_ => Some(arg.to_string()),
|
|
};
|
|
let lastfm_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 lastfm_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()),
|
|
};
|
|
Reply::Text(track.url)
|
|
}
|
|
|
|
async fn fmi(ctx: &Context, arg: &str, id: UserId, avatar: Option<String>) -> Reply {
|
|
let lastfm_user = match arg {
|
|
"" => get_lastfm_username(ctx, id).await,
|
|
_ => Some(arg.to_string()),
|
|
};
|
|
let lastfm_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 lastfm_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 track_info = text::TrackInfo {
|
|
title: track.name,
|
|
artist: track.artist.name,
|
|
album: track.album,
|
|
};
|
|
let image_uri = match track.image.extralarge {
|
|
Some(iu) => iu,
|
|
None => return Reply::Text("Error: getting image uri failed".to_string()),
|
|
};
|
|
|
|
let image = match get_image(ctx, image_uri.as_str()).await {
|
|
Ok(i) => i,
|
|
Err(e) => return Reply::Text(format!("{e}")),
|
|
};
|
|
let mut base_color = PixelWand::new();
|
|
let mut white = PixelWand::new();
|
|
let main_wand = MagickWand::new();
|
|
let art_wand = MagickWand::new();
|
|
let mut art_wand_cluster = MagickWand::new();
|
|
let mask_wand = MagickWand::new();
|
|
let text_image_wand = MagickWand::new();
|
|
|
|
handle_magick_result!(
|
|
base_color.set_color("#7f7f7f"),
|
|
"Failed to set init base color"
|
|
);
|
|
handle_magick_result!(
|
|
art_wand.read_image_blob(image),
|
|
"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(10, 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 base_lab = validate_color(&colors[0]).expect("Failed to validate base color");
|
|
let other_colors = colors
|
|
.iter()
|
|
.filter_map(|col| {
|
|
let testing_color = validate_color(col)?;
|
|
match testing_color.improved_difference(base_lab) {
|
|
10.0.. => Some(testing_color),
|
|
_ => None,
|
|
}
|
|
})
|
|
.collect::<Vec<_>>();
|
|
let mut accent_color = PixelWand::new();
|
|
handle_magick_result!(
|
|
accent_color.set_color("cielab(0.0,0.0,0.0)"),
|
|
"Failed to init accent color"
|
|
);
|
|
if let Some(color) = other_colors.first() {
|
|
let col = format!("cielab{:?}", color.into_components());
|
|
handle_magick_result!(
|
|
accent_color.set_color(col.as_str()),
|
|
"Failed to set accent color"
|
|
);
|
|
}
|
|
|
|
|
|
let use_dark_color = match validate_color(&colors[0]) {
|
|
Some(color) => {
|
|
let black_delta_e = Lab::<D65>::new(0.0, 0.0, 0.0).improved_difference(color);
|
|
let white_delta_e = Lab::<D65>::new(100.0, 0.0, 0.0).improved_difference(color);
|
|
black_delta_e > white_delta_e
|
|
}
|
|
_ => false
|
|
};
|
|
|
|
handle_magick_result!(white.set_color("#ffffff"), "Failed to init white");
|
|
handle_magick_result!(
|
|
mask_wand.read_image_blob(include_bytes!("accent_mask.png")),
|
|
"Failed to load mask"
|
|
);
|
|
handle_magick_result!(
|
|
mask_wand.opaque_paint_image(&white, &accent_color, 0.0, false),
|
|
"Failed to paint accent mask"
|
|
);
|
|
|
|
let avatar_wand = MagickWand::new();
|
|
match avatar {
|
|
Some(avatar_uri) => {
|
|
if let Ok(avatar) = get_image(ctx, avatar_uri.as_str()).await {
|
|
handle_magick_result!(
|
|
avatar_wand.read_image_blob(avatar),
|
|
"Failed to load user avatar"
|
|
);
|
|
} else {
|
|
handle_magick_result!(
|
|
avatar_wand.new_image(100, 100, &white),
|
|
"Failed to load dummy avatar"
|
|
);
|
|
}
|
|
}
|
|
None => handle_magick_result!(
|
|
avatar_wand.new_image(100, 100, &white),
|
|
"Failed to load dummy avatar"
|
|
),
|
|
}
|
|
let avatar_mask_wand = MagickWand::new();
|
|
handle_magick_result!(
|
|
avatar_mask_wand.read_image_blob(include_bytes!("avatar_mask.png")),
|
|
"Failed to load avatar mask"
|
|
);
|
|
handle_magick_result!(
|
|
avatar_wand.compose_images(&avatar_mask_wand, CompositeOperator::CopyAlpha, false, 0, 0),
|
|
"Failed to mask avatar"
|
|
);
|
|
handle_magick_result!(
|
|
avatar_wand.resize_image(64, 64, FilterType::RobidouxSharp),
|
|
"Failed to resize avatar"
|
|
);
|
|
handle_magick_result!(
|
|
avatar_wand.transform_image_colorspace(ColorspaceType::Lab),
|
|
"Failed to set avatar colorspace"
|
|
);
|
|
|
|
let (image_width, image_height) = (548, 147);
|
|
let data = ctx.data.write().await;
|
|
let fonts = data
|
|
.get::<FontsHaver>().expect("Failed to load fonts");
|
|
let text_image_blob = match text::fmi_text(image_width, image_height, track_info, fonts, 19.0, use_dark_color) {
|
|
Ok(bytes) => bytes,
|
|
Err(e) => return Reply::Text(e),
|
|
};
|
|
handle_magick_result!(text_image_wand.read_image_blob(&text_image_blob), "Failed to load text image");
|
|
handle_magick_result!(
|
|
text_image_wand.transform_image_colorspace(ColorspaceType::Lab),
|
|
"Failed to set avatar colorspace"
|
|
);
|
|
|
|
handle_magick_result!(
|
|
main_wand.new_image(image_width as usize, image_height as usize, &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 art image"
|
|
);
|
|
handle_magick_result!(
|
|
main_wand.compose_images(&mask_wand, CompositeOperator::SrcOver, false, 401, 0),
|
|
"Failed to combine mask"
|
|
);
|
|
#[cfg(debug_assertions)]
|
|
if env::var("DEBUG_IMAGE").unwrap_or("0".to_string()) == "1" {
|
|
handle_magick_result!(
|
|
main_wand.compose_images(
|
|
&art_wand_cluster,
|
|
CompositeOperator::SrcOver,
|
|
false,
|
|
160,
|
|
12
|
|
),
|
|
"Failed to combine debug image"
|
|
);
|
|
}
|
|
handle_magick_result!(
|
|
main_wand.compose_images(&text_image_wand, CompositeOperator::SrcOver, false, 0, 0),
|
|
"Failed to combine mask"
|
|
);
|
|
handle_magick_result!(
|
|
main_wand.compose_images(&avatar_wand, CompositeOperator::SrcOver, false, 473, 73),
|
|
"Failed to combine avatar image"
|
|
);
|
|
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 arg1 = cmd_iter.next().unwrap_or("");
|
|
let arg2 = cmd_iter.next().unwrap_or("");
|
|
let resp = match (cmd, arg1, arg2) {
|
|
(".set", arg, "") | (".k", "set", arg) | (".kset", arg, "") => {
|
|
log!("{} received", msg.content);
|
|
Some(set(&ctx, arg, msg.author.id).await)
|
|
},
|
|
(".k", "art", arg) | (".k", arg, "art") | (".ka" | ".kart", arg, "") => {
|
|
log!("{} received", msg.content);
|
|
Some(art(&ctx, arg, msg.author.id).await)
|
|
},
|
|
(".k", "url" | "uri" | "link", arg) | (".k", arg, "url" | "uri" | "link") | (".ku" | ".kl" | ".kurl", arg, "") => {
|
|
log!("{} received", msg.content);
|
|
Some(url(&ctx, arg, msg.author.id).await)
|
|
},
|
|
(".fmi" | ".k" | ".kf" | ".ki" | ".kfmi", arg, "") | (".k", "fm" | "fmi" | "get" | "img", arg) => {
|
|
log!("{} received", msg.content);
|
|
Some(fmi(&ctx, arg, msg.author.id, msg.author.avatar_url()).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 HttpHaver;
|
|
impl TypeMapKey for HttpHaver {
|
|
type Value = reqwest_middleware::ClientWithMiddleware;
|
|
}
|
|
|
|
struct PoolHaver;
|
|
impl TypeMapKey for PoolHaver {
|
|
type Value = SqlitePool;
|
|
}
|
|
|
|
struct FontsHaver;
|
|
impl TypeMapKey for FontsHaver {
|
|
type Value = (SuperFont<'static>, SuperFont<'static>);
|
|
}
|
|
|
|
struct DBResponse {
|
|
lastfm_username: String,
|
|
}
|
|
|
|
async fn get_lastfm_username(ctx: &Context, id: UserId) -> Option<String> {
|
|
log!("get db user {}", id);
|
|
let data = ctx.data.write().await;
|
|
let pool = data
|
|
.get::<PoolHaver>()
|
|
.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) => {
|
|
log!("got user {}", r.lastfm_username);
|
|
Some(r.lastfm_username)
|
|
}
|
|
Err(_) => None,
|
|
}
|
|
}
|
|
|
|
async fn set_lastfm_username(ctx: &Context, id: UserId, user: String) {
|
|
log!("set db user {} {}", id, user);
|
|
let data = ctx.data.write().await;
|
|
let pool = data
|
|
.get::<PoolHaver>()
|
|
.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");
|
|
}
|
|
|
|
async fn get_image(ctx: &Context, url: &str) -> Result<Vec<u8>, reqwest_middleware::Error> {
|
|
log!("get {}", url);
|
|
let data = ctx.data.write().await;
|
|
let http = data
|
|
.get::<HttpHaver>()
|
|
.expect("Failed to get http client");
|
|
match http.get(url).send().await {
|
|
Ok(resp) => {
|
|
log!("response received");
|
|
Ok(resp
|
|
.bytes()
|
|
.await
|
|
.expect("Unable to resolve bytes")
|
|
.to_vec())
|
|
}
|
|
Err(e) => Err(e),
|
|
}
|
|
}
|
|
|
|
#[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 http = reqwest_middleware::ClientBuilder::new(reqwest::Client::new())
|
|
.with(Cache(HttpCache {
|
|
mode: CacheMode::Default,
|
|
manager: CACacheManager::default(),
|
|
options: HttpCacheOptions::default(),
|
|
}))
|
|
.build();
|
|
|
|
FontDB::load_from_dir("fonts"); //singlet instance??? wha?????
|
|
let regular_fonts = FontDB::query("
|
|
NotoSans-Regular
|
|
NotoSansHK-Regular
|
|
NotoSansJP-Regular
|
|
NotoSansKR-Regular
|
|
NotoSansSC-Regular
|
|
NotoSansTC-Regular
|
|
NotoSansArabic-Regular
|
|
Heebo-Regular
|
|
NotoEmoji-Regular
|
|
Symbola
|
|
Unifont
|
|
").expect("Failed to load regular fonts");
|
|
let bold_fonts = FontDB::query("
|
|
NotoSans-SemiBold
|
|
NotoSansJP-Medium
|
|
NotoSansKR-Medium
|
|
NotoSansSC-Medium
|
|
NotoSansTC-Medium
|
|
NotoSansArabic-SemiBold
|
|
Heebo-SemiBold
|
|
NotoEmoji-Medium
|
|
Symbola
|
|
Unifont
|
|
").expect("Failed to load bold fonts");
|
|
|
|
let mut discord_client = Client::builder(&token, intents)
|
|
.event_handler(Handler)
|
|
.await
|
|
.expect("Err creating Discord client");
|
|
|
|
{
|
|
let mut data = discord_client.data.write().await;
|
|
data.insert::<PoolHaver>(pool);
|
|
data.insert::<HttpHaver>(http);
|
|
data.insert::<FontsHaver>((regular_fonts, bold_fonts));
|
|
}
|
|
|
|
if let Err(why) = discord_client.start().await {
|
|
println!("Client error: {why:?}");
|
|
}
|
|
}
|