pikadick/commands/
rule34.rs

1use crate::{
2    client_data::{
3        CacheStatsBuilder,
4        CacheStatsProvider,
5    },
6    util::{
7        TimedCache,
8        TimedCacheEntry,
9    },
10    ClientDataKey,
11};
12use anyhow::Context as _;
13use rand::seq::SliceRandom;
14use serenity::builder::{
15    CreateInteractionResponse,
16    CreateInteractionResponseMessage,
17};
18use std::sync::Arc;
19use tracing::{
20    error,
21    info,
22};
23
24/// A caching rule34 client
25#[derive(Clone, Default, Debug)]
26pub struct Rule34Client {
27    client: rule34::Client,
28    // Ideally, this would be an LRU.
29    // However, we would also need to add time tracking to
30    // get new data when it goes stale.
31    // We would end up duplicating 90% of the logic from [`TimedCache`],
32    // so directly using an LRU isn't worth it.
33    // However, we could add an LRU based on [`TimedCache`]
34    // in the future, or add a setting to it to cap the maximum
35    // number of entries.
36    list_cache: TimedCache<String, rule34::PostList>,
37}
38
39impl Rule34Client {
40    /// Make a new [`Rule34Client`].
41    pub fn new() -> Rule34Client {
42        Rule34Client {
43            client: rule34::Client::new(),
44            list_cache: TimedCache::new(),
45        }
46    }
47
48    /// Search for a query.
49    #[tracing::instrument(skip(self))]
50    pub async fn list(&self, tags: &str) -> anyhow::Result<Arc<TimedCacheEntry<rule34::PostList>>> {
51        if let Some(entry) = self.list_cache.get_if_fresh(tags) {
52            return Ok(entry);
53        }
54
55        let results = self
56            .client
57            .list_posts()
58            .tags(Some(tags))
59            .limit(Some(1_000))
60            .execute()
61            .await
62            .context("failed to search rule34")?;
63        Ok(self.list_cache.insert_and_get(String::from(tags), results))
64    }
65}
66
67impl CacheStatsProvider for Rule34Client {
68    fn publish_cache_stats(&self, cache_stats_builder: &mut CacheStatsBuilder) {
69        cache_stats_builder.publish_stat("rule34", "list_cache", self.list_cache.len() as f32);
70    }
71}
72
73/// Options for the rule34 command
74#[derive(Debug, pikadick_slash_framework::FromOptions)]
75pub struct Rule34Options {
76    // The search query
77    query: String,
78}
79
80/// Create a slash command
81pub fn create_slash_command() -> anyhow::Result<pikadick_slash_framework::Command> {
82    pikadick_slash_framework::CommandBuilder::new()
83        .name("rule34")
84        .description("Look up rule34 for almost anything")
85        .argument(
86            pikadick_slash_framework::ArgumentParamBuilder::new()
87                .name("query")
88                .description("The search query")
89                .kind(pikadick_slash_framework::ArgumentKind::String)
90                .required(true)
91                .build()?,
92        )
93        .on_process(|ctx, interaction, args: Rule34Options| async move {
94            let data_lock = ctx.data.read().await;
95            let client_data = data_lock
96                .get::<ClientDataKey>()
97                .expect("missing client data");
98            let client = client_data.rule34_client.clone();
99            drop(data_lock);
100
101            let query_str = rule34::SearchQueryBuilder::new()
102                .add_tag_iter(args.query.split(' '))
103                .take_query_string();
104
105            info!("searching rule34 for \"{query_str}\"");
106
107            let result = client
108                .list(&query_str)
109                .await
110                .context("failed to get search results");
111
112            let mut message_builder = CreateInteractionResponseMessage::new();
113            match result {
114                Ok(list_results) => {
115                    let maybe_list_result: Option<String> = list_results
116                        .data()
117                        .posts
118                        .choose(&mut rand::thread_rng())
119                        .map(|list_result| list_result.file_url.to_string());
120
121                    if let Some(file_url) = maybe_list_result {
122                        info!("sending \"{file_url}\"");
123                        message_builder = message_builder.content(file_url);
124                    } else {
125                        info!("no results");
126                        message_builder =
127                            message_builder.content(format!("No results for \"{query_str}\""));
128                    }
129                }
130                Err(error) => {
131                    error!("{error:?}");
132                    message_builder = message_builder.content(format!("{error:?}"));
133                }
134            }
135            let response = CreateInteractionResponse::Message(message_builder);
136            interaction.create_response(&ctx.http, response).await?;
137
138            client.list_cache.trim();
139
140            Ok(())
141        })
142        .build()
143        .context("failed to build rule34 command")
144}
pikadick/commands/rule34.rs

pikadick/commands/
rule34.rs
