Compare commits

...

4 Commits

Author SHA1 Message Date
270878e4fa Added "TODO" to README 2025-05-16 10:51:18 -05:00
88d4b19655 Improved color processing, improved logging 2025-05-16 10:50:00 -05:00
f66f89dc04 Added avatar display 2025-05-16 06:25:18 -05:00
792d9615d6 Added http caching of lastfm images 2025-05-16 02:28:02 -05:00
7 changed files with 973 additions and 57 deletions

1
.gitignore vendored
View File

@ -1,3 +1,4 @@
/target /target
/.env /.env
/watcat.db* /watcat.db*
/http-cacache

761
Cargo.lock generated

File diff suppressed because it is too large Load Diff

View File

@ -6,13 +6,15 @@ edition = "2024"
[dependencies] [dependencies]
serenity = "0.12" serenity = "0.12"
lastfm = "0.10.0" lastfm = "0.10.0"
dotenvy = "0.15.7"
reqwest = "0.12"
reqwest-middleware = "0.4.2"
http-cache-reqwest = "0.15.1"
palette = "0.7.6"
[dependencies.magick_rust] [dependencies.magick_rust]
path = "./magick-rust" path = "./magick-rust"
[dependencies.dotenvy]
version = "0.15.7"
[dependencies.tokio] [dependencies.tokio]
version = "1.21.2" version = "1.21.2"
features = ["macros", "rt-multi-thread"] features = ["macros", "rt-multi-thread"]

View File

@ -24,3 +24,8 @@ create the database ``sqlx database create``
(note: make sure you have sqlx installed first) (note: make sure you have sqlx installed first)
init the database ``sqlx migrate run`` init the database ``sqlx migrate run``
TODO
----
literally just add the text

View File

Before

Width:  |  Height:  |  Size: 1008 B

After

Width:  |  Height:  |  Size: 1008 B

BIN
src/avatar_mask.png Normal file

Binary file not shown.

After

Width:  |  Height:  |  Size: 1.9 KiB

View File

@ -1,5 +1,8 @@
extern crate dotenvy; extern crate dotenvy;
extern crate lastfm; 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::async_trait;
use serenity::builder::{CreateAttachment, CreateMessage}; use serenity::builder::{CreateAttachment, CreateMessage};
@ -13,6 +16,10 @@ use magick_rust::{ColorspaceType, CompositeOperator, MagickWand, PixelWand, magi
use std::env; use std::env;
use std::sync::Once; use std::sync::Once;
use palette::color_difference::ImprovedCiede2000;
use palette::lab::Lab;
use palette::white_point::D65;
extern crate sqlx; extern crate sqlx;
use sqlx::sqlite::SqlitePool; use sqlx::sqlite::SqlitePool;
@ -25,10 +32,35 @@ enum Reply {
Text(String), 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 { macro_rules! handle_magick_result {
($a: expr, $b: literal) => { ($a: expr, $b: literal) => {
match $a { match $a {
Ok(_) => {} Ok(_) => {
log!("{} run successfully", stringify!($a));
}
Err(e) => return Reply::Text(format!("Error: {} {}", $b, e)), Err(e) => return Reply::Text(format!("Error: {} {}", $b, e)),
} }
}; };
@ -43,17 +75,50 @@ macro_rules! handle_magick_option {
}; };
} }
#[allow(unused)] // this function is almost assuredly the least idiomatic rust code I will ever write.
async fn fmi(ctx: &Context, arg: &str, id: UserId) -> Reply { // 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 fmi(ctx: &Context, arg: &str, id: UserId, avatar: Option<String>) -> Reply {
let lastfm_user = match arg { let lastfm_user = match arg {
"" => get_lastfm_username(ctx, id).await, "" => get_lastfm_username(ctx, id).await,
_ => Some(arg.to_string()), _ => Some(arg.to_string()),
}; };
let client = match lastfm_user { let lastfm_client = match lastfm_user {
Some(s) => lastfm::Client::<String, String>::from_env(s), Some(s) => lastfm::Client::<String, String>::from_env(s),
None => return Reply::Text("No last.fm username set.".to_string()), None => return Reply::Text("No last.fm username set.".to_string()),
}; };
let now_playing = match client.now_playing().await { let now_playing = match lastfm_client.now_playing().await {
Ok(np) => np, Ok(np) => np,
Err(e) => return Reply::Text(format!("Error: grabbing last.fm user data failed {e}")), Err(e) => return Reply::Text(format!("Error: grabbing last.fm user data failed {e}")),
}; };
@ -65,19 +130,24 @@ async fn fmi(ctx: &Context, arg: &str, id: UserId) -> Reply {
Some(iu) => iu, Some(iu) => iu,
None => return Reply::Text("Error: getting image uri failed".to_string()), 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 base_color = PixelWand::new();
let mut white = PixelWand::new(); let mut white = PixelWand::new();
let mut main_wand = MagickWand::new(); let main_wand = MagickWand::new();
let mut art_wand = MagickWand::new(); let art_wand = MagickWand::new();
let mut art_wand_cluster = MagickWand::new(); let mut art_wand_cluster = MagickWand::new();
let mut mask_wand = MagickWand::new(); let mask_wand = MagickWand::new();
handle_magick_result!( handle_magick_result!(
base_color.set_color("#7f7f7f"), base_color.set_color("#7f7f7f"),
"Failed to set init base color" "Failed to set init base color"
); );
handle_magick_result!( handle_magick_result!(
art_wand.read_image(image_uri.as_str()), art_wand.read_image_blob(image),
"Failed to read image from uri" "Failed to read image from uri"
); );
handle_magick_result!( handle_magick_result!(
@ -100,7 +170,7 @@ async fn fmi(ctx: &Context, arg: &str, id: UserId) -> Reply {
"Failed to set cluster colorspace" "Failed to set cluster colorspace"
); );
handle_magick_result!( handle_magick_result!(
art_wand_cluster.kmeans(5, 200, 0.001), art_wand_cluster.kmeans(10, 200, 0.001),
"Failed to run kmeans" "Failed to run kmeans"
); );
let mut colors = handle_magick_option!( let mut colors = handle_magick_option!(
@ -108,16 +178,76 @@ async fn fmi(ctx: &Context, arg: &str, id: UserId) -> Reply {
"Failed to get color histogram" "Failed to get color histogram"
); );
colors.sort_by_cached_key(|color| -(color.get_color_count() as isize)); colors.sort_by_cached_key(|color| -(color.get_color_count() as isize));
let mut accent_color = match colors.len() { let base_lab = validate_color(&colors[0]).expect("Failed to validate base color");
1 => &colors[0], let other_colors = colors
_ => &colors[1], .iter()
}; .filter_map(|col| {
let testing_color = validate_color(col)?;
handle_magick_result!(white.set_color("#ffffff"), "Failed to set init white"); match testing_color.improved_difference(base_lab) {
handle_magick_result!(mask_wand.read_image("mask.png"), "Failed to load mask"); 10.0.. => Some(testing_color),
_ => None,
}
})
.collect::<Vec<_>>();
let mut accent_color = PixelWand::new();
handle_magick_result!( handle_magick_result!(
mask_wand.opaque_paint_image(&white, accent_color, 0.0, false), accent_color.set_color("#ffffff"),
"Failed to paint mask" "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"
);
}
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.adaptive_resize_image(64, 64),
"Failed to resize avatar"
);
handle_magick_result!(
avatar_wand.transform_image_colorspace(ColorspaceType::Lab),
"Failed to set avatar colorspace"
); );
handle_magick_result!( handle_magick_result!(
@ -138,7 +268,11 @@ async fn fmi(ctx: &Context, arg: &str, id: UserId) -> Reply {
); );
handle_magick_result!( handle_magick_result!(
main_wand.compose_images(&art_wand, CompositeOperator::SrcOver, false, 12, 12), main_wand.compose_images(&art_wand, CompositeOperator::SrcOver, false, 12, 12),
"Failed to combine images" "Failed to combine art image"
);
handle_magick_result!(
main_wand.compose_images(&mask_wand, CompositeOperator::SrcOver, false, 401, 0),
"Failed to combine mask"
); );
handle_magick_result!( handle_magick_result!(
main_wand.compose_images( main_wand.compose_images(
@ -148,18 +282,18 @@ async fn fmi(ctx: &Context, arg: &str, id: UserId) -> Reply {
160, 160,
12 12
), ),
"Failed to combine debug images" "Failed to combine debug image"
); );
handle_magick_result!( handle_magick_result!(
main_wand.compose_images(&mask_wand, CompositeOperator::SrcOver, false, 401, 0), main_wand.compose_images(&avatar_wand, CompositeOperator::SrcOver, false, 473, 73),
"Failed to combine mask" "Failed to combine avatar image"
); );
Reply::Image(main_wand.write_image_blob("png").unwrap()) Reply::Image(main_wand.write_image_blob("png").unwrap())
} }
async fn set(ctx: &Context, arg: &str, id: UserId) -> Reply { async fn set(ctx: &Context, arg: &str, id: UserId) -> Reply {
set_lastfm_username(ctx, id, arg.to_string()).await; set_lastfm_username(ctx, id, arg.to_string()).await;
Reply::Text(format!("set user {}", arg)) Reply::Text(format!("set user {}", arg))
} }
#[async_trait] #[async_trait]
@ -169,8 +303,14 @@ impl EventHandler for Handler {
let cmd = cmd_iter.next().unwrap_or(""); let cmd = cmd_iter.next().unwrap_or("");
let arg = cmd_iter.next().unwrap_or(""); let arg = cmd_iter.next().unwrap_or("");
let resp = match cmd { let resp = match cmd {
".fmi" => Some(fmi(&ctx, arg, msg.author.id).await), ".fmi" => {
".set" => Some(set(&ctx, arg, msg.author.id).await), log!("{} received", msg.content);
Some(fmi(&ctx, arg, msg.author.id, msg.author.avatar_url()).await)
}
".set" => {
log!("{} received", msg.content);
Some(set(&ctx, arg, msg.author.id).await)
}
_ => None, _ => None,
}; };
if let Some(reply) = resp { if let Some(reply) = resp {
@ -190,8 +330,12 @@ impl EventHandler for Handler {
} }
} }
struct PoolContainer; struct HttpContainer;
impl TypeMapKey for HttpContainer {
type Value = reqwest_middleware::ClientWithMiddleware;
}
struct PoolContainer;
impl TypeMapKey for PoolContainer { impl TypeMapKey for PoolContainer {
type Value = SqlitePool; type Value = SqlitePool;
} }
@ -201,6 +345,7 @@ struct DBResponse {
} }
async fn get_lastfm_username(ctx: &Context, id: UserId) -> Option<String> { async fn get_lastfm_username(ctx: &Context, id: UserId) -> Option<String> {
log!("get db user {}", id);
let data = ctx.data.write().await; let data = ctx.data.write().await;
let pool = data let pool = data
.get::<PoolContainer>() .get::<PoolContainer>()
@ -218,12 +363,16 @@ async fn get_lastfm_username(ctx: &Context, id: UserId) -> Option<String> {
.fetch_one(pool) .fetch_one(pool)
.await; .await;
match resp { match resp {
Ok(r) => Some(r.lastfm_username), Ok(r) => {
log!("got user {}", r.lastfm_username);
Some(r.lastfm_username)
}
Err(_) => None, Err(_) => None,
} }
} }
async fn set_lastfm_username(ctx: &Context, id: UserId, user: String) { async fn set_lastfm_username(ctx: &Context, id: UserId, user: String) {
log!("set db user {} {}", id, user);
let data = ctx.data.write().await; let data = ctx.data.write().await;
let pool = data let pool = data
.get::<PoolContainer>() .get::<PoolContainer>()
@ -234,10 +383,31 @@ async fn set_lastfm_username(ctx: &Context, id: UserId, user: String) {
INSERT INTO users INSERT INTO users
VALUES (?1, ?2) VALUES (?1, ?2)
"#, "#,
id, user id,
user
) )
.execute(pool) .execute(pool)
.await.expect("Failed to insert"); .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::<HttpContainer>()
.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] #[tokio::main]
@ -258,19 +428,24 @@ async fn main() {
| GatewayIntents::DIRECT_MESSAGES | GatewayIntents::DIRECT_MESSAGES
| GatewayIntents::MESSAGE_CONTENT; | GatewayIntents::MESSAGE_CONTENT;
let mut client = Client::builder(&token, intents) let mut discord_client = Client::builder(&token, intents)
.event_handler(Handler) .event_handler(Handler)
.await .await
.expect("Err creating client"); .expect("Err creating Discord client");
let http = reqwest_middleware::ClientBuilder::new(reqwest::Client::new())
.with(Cache(HttpCache {
mode: CacheMode::Default,
manager: CACacheManager::default(),
options: HttpCacheOptions::default(),
}))
.build();
{ {
let mut data = client.data.write().await; let mut data = discord_client.data.write().await;
data.insert::<PoolContainer>(pool); data.insert::<PoolContainer>(pool);
data.insert::<HttpContainer>(http);
} }
if let Err(why) = client.start().await { if let Err(why) = discord_client.start().await {
println!("Client error: {why:?}"); println!("Client error: {why:?}");
} else {
println!("watcat started");
} }
} }