pikadick/commands/
sauce_nao.rs

1use crate::{
2    checks::ENABLED_CHECK,
3    client_data::{
4        CacheStatsBuilder,
5        CacheStatsProvider,
6    },
7    util::{
8        LoadingReaction,
9        TimedCache,
10        TimedCacheEntry,
11    },
12    ClientDataKey,
13};
14use anyhow::Context as _;
15use serenity::{
16    builder::{
17        CreateEmbed,
18        CreateMessage,
19    },
20    framework::standard::{
21        macros::command,
22        Args,
23        CommandResult,
24    },
25    model::prelude::*,
26    prelude::*,
27};
28use std::sync::Arc;
29use tracing::error;
30
31#[derive(Clone, Debug)]
32pub struct SauceNaoClient {
33    client: sauce_nao::Client,
34    search_cache: TimedCache<String, sauce_nao::OkResponse>,
35}
36
37impl SauceNaoClient {
38    pub fn new(api_key: &str) -> Self {
39        Self {
40            client: sauce_nao::Client::new(api_key),
41            search_cache: TimedCache::new(),
42        }
43    }
44
45    /// Search for an image, with caching
46    pub async fn search(
47        &self,
48        query: &str,
49    ) -> anyhow::Result<Arc<TimedCacheEntry<sauce_nao::OkResponse>>> {
50        if let Some(entry) = self.search_cache.get_if_fresh(query) {
51            return Ok(entry);
52        }
53
54        let search_results = self
55            .client
56            .search(query)
57            .await
58            .context("failed to search for image")?;
59
60        self.search_cache
61            .insert(String::from(query), search_results);
62
63        self.search_cache
64            .get_if_fresh(query)
65            .context("cache data expired")
66    }
67}
68
69impl CacheStatsProvider for SauceNaoClient {
70    fn publish_cache_stats(&self, cache_stats_builder: &mut CacheStatsBuilder) {
71        cache_stats_builder.publish_stat(
72            "sauce-nao",
73            "search_cache",
74            self.search_cache.len() as f32,
75        );
76    }
77}
78
79#[command("sauce-nao")]
80#[description("Search SauceNao for an image at a url")]
81#[usage("<img_url>")]
82#[example("https://konachan.com/image/5982d8946ae503351e960f097f84cd90/Konachan.com%20-%20330136%20animal%20nobody%20original%20signed%20yutaka_kana.jpg")]
83#[checks(Enabled)]
84#[min_args(1)]
85#[max_args(1)]
86async fn sauce_nao(ctx: &Context, msg: &Message, mut args: Args) -> CommandResult {
87    let data_lock = ctx.data.read().await;
88    let client_data = data_lock
89        .get::<ClientDataKey>()
90        .expect("missing client data");
91    let client = client_data.sauce_nao_client.clone();
92    drop(data_lock);
93
94    let query = args.trimmed().current().expect("missing query");
95
96    let mut loading = LoadingReaction::new(ctx.http.clone(), msg);
97
98    let result = client
99        .search(query)
100        .await
101        .context("failed to search for image");
102
103    match result {
104        Ok(data) => {
105            let data = data.data();
106
107            match data.results.first() {
108                Some(data) => {
109                    let mut embed_builder = CreateEmbed::new()
110                        .title("SauceNao Best Match")
111                        .image(data.header.thumbnail.as_str());
112                    if let Some(ext_url) = data.data.ext_urls.first() {
113                        embed_builder = embed_builder
114                            .description(ext_url.as_str())
115                            .url(ext_url.as_str());
116                    }
117
118                    if let Some(source) = data.data.source.as_deref() {
119                        embed_builder = embed_builder.field("Source", source, true);
120                    }
121
122                    if let Some(eng_name) = data.data.eng_name.as_deref() {
123                        embed_builder = embed_builder.field("English Name", eng_name, true);
124                    }
125
126                    if let Some(jp_name) = data.data.jp_name.as_deref() {
127                        embed_builder = embed_builder.field("Jap Name", jp_name, true);
128                    }
129
130                    let message_builder = CreateMessage::new().embed(embed_builder);
131
132                    msg.channel_id
133                        .send_message(&ctx.http, message_builder)
134                        .await?;
135
136                    loading.send_ok();
137                }
138                None => {
139                    msg.channel_id
140                        .say(&ctx.http, format!("No results on SauceNao for \"{query}\""))
141                        .await?;
142                }
143            }
144        }
145        Err(error) => {
146            error!("{error:?}");
147            msg.channel_id.say(&ctx.http, format!("{error:?}")).await?;
148        }
149    }
150
151    Ok(())
152}
pikadick/commands/sauce_nao.rs

pikadick/commands/
sauce_nao.rs
