diff options
Diffstat (limited to '')
| -rw-r--r-- | src/main.rs | 190 |
1 files changed, 190 insertions, 0 deletions
diff --git a/src/main.rs b/src/main.rs new file mode 100644 index 0000000..4a89315 --- /dev/null +++ b/src/main.rs | |||
| @@ -0,0 +1,190 @@ | |||
| 1 | use std::collections::HashMap; | ||
| 2 | |||
| 3 | use poise::serenity_prelude as serenity; | ||
| 4 | use songbird::SerenityInit; | ||
| 5 | use tracing_subscriber::filter::EnvFilter; | ||
| 6 | |||
| 7 | mod commands; | ||
| 8 | mod config; | ||
| 9 | mod copyparty; | ||
| 10 | mod database; | ||
| 11 | mod error; | ||
| 12 | mod events; | ||
| 13 | mod lang; | ||
| 14 | mod openrouter; | ||
| 15 | mod tasks; | ||
| 16 | mod voice; | ||
| 17 | |||
| 18 | use database::DbPool; | ||
| 19 | |||
| 20 | pub struct Data { | ||
| 21 | pub db_pool: DbPool, | ||
| 22 | pub config: config::Config, | ||
| 23 | pub voice_manager: voice::VoiceManager, | ||
| 24 | pub copyparty: copyparty::CopypartyClient, | ||
| 25 | pub http_client: reqwest::Client, | ||
| 26 | pub volume: std::sync::atomic::AtomicU64, | ||
| 27 | pub triggers: tokio::sync::Mutex<Vec<(database::Trigger, Vec<String>, Vec<String>)>>, | ||
| 28 | /// Autocomplete lookup, keyed per-user to avoid cross-command collisions. | ||
| 29 | /// | ||
| 30 | /// Mirrors Python's per-cog `_cached_urls[user_id]` map: when two users | ||
| 31 | /// search the same filename simultaneously, each invocation must resolve its | ||
| 32 | /// own stashed `SearchHit`, so the cache is keyed by `(UserId, filename)`. | ||
| 33 | pub ac_cache: tokio::sync::Mutex<HashMap<(serenity::UserId, String), copyparty::SearchHit>>, | ||
| 34 | /// Active music track per guild (only one music source at a time). | ||
| 35 | pub music_tracks: tokio::sync::Mutex<HashMap<u64, songbird::tracks::TrackHandle>>, | ||
| 36 | } | ||
| 37 | |||
| 38 | impl Data { | ||
| 39 | pub async fn reload_triggers(&self) { | ||
| 40 | let triggers = database::load_all_triggers_with_keywords_and_responses(&self.db_pool) | ||
| 41 | .await | ||
| 42 | .unwrap_or_default(); | ||
| 43 | *self.triggers.lock().await = triggers; | ||
| 44 | } | ||
| 45 | |||
| 46 | pub async fn ac_stash( | ||
| 47 | &self, | ||
| 48 | user_id: serenity::UserId, | ||
| 49 | key: String, | ||
| 50 | hit: copyparty::SearchHit, | ||
| 51 | ) { | ||
| 52 | self.ac_cache.lock().await.insert((user_id, key), hit); | ||
| 53 | } | ||
| 54 | |||
| 55 | pub async fn ac_lookup( | ||
| 56 | &self, | ||
| 57 | user_id: serenity::UserId, | ||
| 58 | key: &str, | ||
| 59 | ) -> Option<copyparty::SearchHit> { | ||
| 60 | self.ac_cache | ||
| 61 | .lock() | ||
| 62 | .await | ||
| 63 | .get(&(user_id, key.to_string())) | ||
| 64 | .cloned() | ||
| 65 | } | ||
| 66 | } | ||
| 67 | |||
| 68 | #[allow(dead_code)] | ||
| 69 | type Error = error::Error; | ||
| 70 | |||
| 71 | #[tokio::main] | ||
| 72 | async fn main() { | ||
| 73 | let _ = rustls::crypto::ring::default_provider().install_default(); | ||
| 74 | |||
| 75 | tracing_subscriber::fmt() | ||
| 76 | .with_env_filter( | ||
| 77 | EnvFilter::builder() | ||
| 78 | .parse(std::env::var("RUST_LOG").unwrap_or_else(|_| "info".into())) | ||
| 79 | .unwrap(), | ||
| 80 | ) | ||
| 81 | .init(); | ||
| 82 | |||
| 83 | dotenvy::dotenv().ok(); | ||
| 84 | config::check_env_vars(); | ||
| 85 | let config = config::Config::from_env(); | ||
| 86 | |||
| 87 | let db_pool = database::init_pool(&config.database_path) | ||
| 88 | .await | ||
| 89 | .expect("Failed to initialize database"); | ||
| 90 | |||
| 91 | let http_client = reqwest::Client::new(); | ||
| 92 | |||
| 93 | let copyparty = copyparty::CopypartyClient::new( | ||
| 94 | config.copyparty_url.clone(), | ||
| 95 | config.copyparty_url_upload.clone(), | ||
| 96 | config.copyparty_password.clone(), | ||
| 97 | ); | ||
| 98 | |||
| 99 | let triggers = database::load_all_triggers_with_keywords_and_responses(&db_pool) | ||
| 100 | .await | ||
| 101 | .unwrap_or_default(); | ||
| 102 | |||
| 103 | let commands_list = commands::all_commands(); | ||
| 104 | let config_for_client = config.clone(); | ||
| 105 | |||
| 106 | let framework = poise::Framework::builder() | ||
| 107 | .options(poise::FrameworkOptions { | ||
| 108 | commands: commands_list, | ||
| 109 | event_handler: |ctx, event, _framework, user_data| { | ||
| 110 | Box::pin(async move { | ||
| 111 | handle_event(ctx, event, user_data).await; | ||
| 112 | Ok(()) | ||
| 113 | }) | ||
| 114 | }, | ||
| 115 | ..Default::default() | ||
| 116 | }) | ||
| 117 | .setup(move |ctx, _ready, framework| { | ||
| 118 | let config = config.clone(); | ||
| 119 | let db_pool = db_pool.clone(); | ||
| 120 | let http_client = http_client.clone(); | ||
| 121 | let copyparty = copyparty.clone(); | ||
| 122 | let triggers = triggers.clone(); | ||
| 123 | Box::pin(async move { | ||
| 124 | let vm = voice::VoiceManager::new(ctx).await; | ||
| 125 | |||
| 126 | let data = Data { | ||
| 127 | config, | ||
| 128 | db_pool, | ||
| 129 | voice_manager: vm, | ||
| 130 | copyparty, | ||
| 131 | http_client, | ||
| 132 | volume: std::sync::atomic::AtomicU64::new(1.0f64.to_bits()), | ||
| 133 | triggers: tokio::sync::Mutex::new(triggers), | ||
| 134 | ac_cache: tokio::sync::Mutex::new(HashMap::new()), | ||
| 135 | music_tracks: tokio::sync::Mutex::new(HashMap::new()), | ||
| 136 | }; | ||
| 137 | |||
| 138 | // Spawn cron scheduler loop. | ||
| 139 | { | ||
| 140 | let cron_ctx = crate::tasks::CronContext { | ||
| 141 | pool: data.db_pool.clone(), | ||
| 142 | http: ctx.http.clone(), | ||
| 143 | voice_manager: data.voice_manager.clone(), | ||
| 144 | copyparty: data.copyparty.clone(), | ||
| 145 | }; | ||
| 146 | tokio::spawn(async move { | ||
| 147 | crate::tasks::cron_scheduler_loop(cron_ctx).await; | ||
| 148 | }); | ||
| 149 | } | ||
| 150 | |||
| 151 | poise::builtins::register_globally(ctx, &framework.options().commands).await?; | ||
| 152 | |||
| 153 | events::handle_ready(ctx).await.map_err(|e| { | ||
| 154 | tracing::error!("on_ready handler: {e}"); | ||
| 155 | e | ||
| 156 | })?; | ||
| 157 | |||
| 158 | Ok(data) | ||
| 159 | }) | ||
| 160 | }) | ||
| 161 | .build(); | ||
| 162 | |||
| 163 | let mut client = | ||
| 164 | serenity::ClientBuilder::new(&config_for_client.token, serenity::GatewayIntents::all()) | ||
| 165 | .framework(framework) | ||
| 166 | .register_songbird() | ||
| 167 | .await | ||
| 168 | .expect("Error creating client"); | ||
| 169 | |||
| 170 | client.start().await.expect("Error starting client"); | ||
| 171 | } | ||
| 172 | |||
| 173 | async fn handle_event(ctx: &serenity::Context, event: &serenity::FullEvent, user_data: &Data) { | ||
| 174 | match event { | ||
| 175 | serenity::FullEvent::Ready { | ||
| 176 | data_about_bot: _, .. | ||
| 177 | } => {} | ||
| 178 | serenity::FullEvent::VoiceStateUpdate { old, new } => { | ||
| 179 | events::handle_voice_state_update(ctx, user_data, old.clone(), new) | ||
| 180 | .await | ||
| 181 | .unwrap_or_else(|e| tracing::error!("{}: {e}", lang::L_EV_VOICE_STATE_ERROR)); | ||
| 182 | } | ||
| 183 | serenity::FullEvent::Message { new_message } => { | ||
| 184 | events::handle_message(ctx, user_data, new_message) | ||
| 185 | .await | ||
| 186 | .unwrap_or_else(|e| tracing::error!("{}: {e}", lang::L_EV_MESSAGE_HANDLER_ERROR)); | ||
| 187 | } | ||
| 188 | _ => {} | ||
| 189 | } | ||
| 190 | } | ||
