Files
kirara/src/main.rs
T
2025-07-30 17:56:01 -05:00

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:?}");
}
}