pikadick/
client_data.rs

1use crate::{
2    checks::EnabledCheckData,
3    commands::{
4        deviantart::DeviantartClient,
5        fml::FmlClient,
6        iqdb::IqdbClient,
7        nekos::NekosClient,
8        quizizz::QuizizzClient,
9        r6stats::R6StatsClient,
10        r6tracker::R6TrackerClient,
11        reddit_embed::RedditEmbedData,
12        rule34::Rule34Client,
13        sauce_nao::SauceNaoClient,
14        shift::ShiftClient,
15        tic_tac_toe::TicTacToeData,
16        tiktok_embed::TikTokData,
17        urban::UrbanClient,
18    },
19    config::Config,
20    database::Database,
21    util::EncoderTask,
22};
23use anyhow::Context;
24use serenity::gateway::ShardManager;
25use std::{
26    collections::BTreeMap,
27    fmt::Debug,
28    sync::Arc,
29};
30use tracing::error;
31
32/// A tool to build cache stats
33#[derive(Debug)]
34pub struct CacheStatsBuilder {
35    stats: BTreeMap<&'static str, BTreeMap<&'static str, f32>>,
36}
37
38impl CacheStatsBuilder {
39    /// Make a new [`CacheStatsBuilder`].
40    pub fn new() -> Self {
41        Self {
42            stats: BTreeMap::new(),
43        }
44    }
45
46    /// Publish a stat to a section
47    pub fn publish_stat(&mut self, section: &'static str, name: &'static str, value: f32) {
48        self.stats.entry(section).or_default().insert(name, value);
49    }
50
51    /// Get the inner stats
52    pub fn into_inner(self) -> BTreeMap<&'static str, BTreeMap<&'static str, f32>> {
53        self.stats
54    }
55}
56
57impl Default for CacheStatsBuilder {
58    fn default() -> Self {
59        Self::new()
60    }
61}
62
63/// A type that can provide cache stats
64pub trait CacheStatsProvider {
65    /// Publish stats to the provided [`CacheStatsBuilder`].
66    fn publish_cache_stats(&self, cache_stats_builder: &mut CacheStatsBuilder);
67}
68
69/// The [`ClientData`].
70#[derive(Debug)]
71pub struct ClientData {
72    /// The discord shard_manager
73    pub shard_manager: Arc<ShardManager>,
74
75    /// The client for nekos
76    pub nekos_client: NekosClient,
77    /// The R6Stats client
78    pub r6stats_client: R6StatsClient,
79    /// The r6tracker client
80    pub r6tracker_client: R6TrackerClient,
81    /// The rule34 client
82    pub rule34_client: Rule34Client,
83    /// The quizizz client
84    pub quizizz_client: QuizizzClient,
85    /// The fml client
86    pub fml_client: FmlClient,
87    /// The shift client
88    pub shift_client: ShiftClient,
89    /// The reddit embed data
90    pub reddit_embed_data: RedditEmbedData,
91    /// The enabled check data
92    pub enabled_check_data: EnabledCheckData,
93    /// The insta client data
94    pub insta_client: insta::Client,
95    /// The deviantart client
96    pub deviantart_client: DeviantartClient,
97    /// The urban dictionary client
98    pub urban_client: UrbanClient,
99    /// The xkcd client
100    pub xkcd_client: xkcd::Client,
101    /// The tic tac toe data
102    pub tic_tac_toe_data: TicTacToeData,
103    /// The iqdb client
104    pub iqdb_client: IqdbClient,
105    /// The sauce nao client
106    pub sauce_nao_client: SauceNaoClient,
107    /// The open ai client
108    pub open_ai_client: open_ai::Client,
109    /// The yodaspeak client
110    pub yodaspeak: yodaspeak::Client,
111    /// TikTokData
112    pub tiktok_data: TikTokData,
113    /// Encoder Task
114    pub encoder_task: EncoderTask,
115
116    /// The database
117    pub db: Database,
118
119    /// The config
120    pub config: Arc<Config>,
121}
122
123impl ClientData {
124    /// Init this client data
125    pub async fn init(
126        shard_manager: Arc<ShardManager>,
127        config: Arc<Config>,
128        db: Database,
129    ) -> anyhow::Result<Self> {
130        // TODO: Standardize an async init system with allocated data per command somehow. Maybe boxes?
131
132        let cache_dir = config.cache_dir();
133        let encoder_task = EncoderTask::new();
134
135        let deviantart_client = DeviantartClient::new(&db)
136            .await
137            .context("failed to init deviantart client")?;
138        let tiktok_data = TikTokData::new(&cache_dir, encoder_task.clone())
139            .await
140            .context("failed to init tiktok data")?;
141
142        Ok(ClientData {
143            shard_manager,
144
145            nekos_client: Default::default(),
146            r6stats_client: Default::default(),
147            r6tracker_client: Default::default(),
148            rule34_client: Default::default(),
149            quizizz_client: Default::default(),
150            fml_client: FmlClient::new(config.fml.key.to_string()),
151            shift_client: ShiftClient::new(),
152            reddit_embed_data: Default::default(),
153            enabled_check_data: Default::default(),
154            insta_client: insta::Client::new(),
155            deviantart_client,
156            urban_client: Default::default(),
157            xkcd_client: Default::default(),
158            tic_tac_toe_data: Default::default(),
159            iqdb_client: Default::default(),
160            sauce_nao_client: SauceNaoClient::new(config.sauce_nao.api_key.as_str()),
161            open_ai_client: open_ai::Client::new(config.open_ai.api_key.as_str()),
162            yodaspeak: yodaspeak::Client::new(),
163            tiktok_data,
164            encoder_task,
165
166            db,
167
168            config,
169        })
170    }
171
172    /// Generate cache stats
173    /// Currently, In order for something to show up in cache-stats it must be added here.
174    /// More automation is desirable in the future.
175    pub fn generate_cache_stats(&self) -> BTreeMap<&'static str, BTreeMap<&'static str, f32>> {
176        let mut stat_builder = CacheStatsBuilder::new();
177
178        let cache_stat_providers: &[&dyn CacheStatsProvider] = &[
179            &self.fml_client,
180            &self.nekos_client,
181            &self.r6stats_client,
182            &self.r6tracker_client,
183            &self.reddit_embed_data,
184            &self.rule34_client,
185            &self.shift_client,
186            &self.deviantart_client,
187            &self.urban_client,
188            &self.iqdb_client,
189        ];
190
191        for cache_stat_provider in cache_stat_providers {
192            cache_stat_provider.publish_cache_stats(&mut stat_builder);
193        }
194
195        stat_builder.into_inner()
196    }
197
198    /// Shutdown anything that needs to be shut down.
199    ///
200    /// Errors are logged to the console,
201    /// but not returned to the user as it is assumed that they don't matter in the middle of a shutdown.
202    pub async fn shutdown(&self) {
203        if let Err(e) = self.encoder_task.shutdown().await {
204            error!("{:?}", e);
205        }
206    }
207}
pikadick/client_data.rs

pikadick/
client_data.rs
