use std::collections::{HashSet, HashMap}; use std::env; use std::sync::Arc; use commands::audio::{create_response, AudioConfig, AudioConfigs}; use diesel::r2d2::{Pool, ConnectionManager}; use diesel::pg::PgConnection; use dotenv::dotenv; use log::{error, warn, info}; use serenity::async_trait; use serenity::framework::StandardFramework; use serenity::model::application::interaction::Interaction; use serenity::model::gateway::Ready; use serenity::model::channel::Message; use serenity::http::Http; use serenity::prelude::*; use songbird::SerenityInit; mod commands; mod database; struct Handler { // Open AI Config oai: Option, pool: Pool> } #[async_trait] impl EventHandler for Handler { async fn message(&self, ctx: Context, msg: Message) { // Ignore messages from bots if msg.author.bot { return; } match &self.oai { Some(oai) => { match msg.mentions_me(&ctx.http).await { Ok(mentioned) => { let bot_in_thread = match msg.channel_id.get_thread_members(&ctx.http).await { Ok(t) => { match t.iter().find(|t| t.user_id.unwrap().0 == ctx.cache.current_user_id().0) { Some(_) => true, None => false } } Err(_) => false }; // let has_bot = msg.channel_id.get_thread_members(&ctx.http).await.unwrap().contains(ctx.cache.current_user_id().0); if mentioned || bot_in_thread { commands::oai::generate_response(&ctx, &msg, oai, &self.pool).await; } } Err(why) => warn!("Could not check mentions: {:?}", why) }; } None => {} } } async fn interaction_create(&self, ctx: Context, interaction: Interaction) { if let Interaction::ApplicationCommand(command) = interaction { match command.data.name.as_str() { "play" => commands::audio::play::run(&ctx, &command).await, "stop" => commands::audio::stop::run(&ctx, &command).await, "pause" => commands::audio::pause::run(&ctx, &command).await, "resume" => commands::audio::resume::run(&ctx, &command).await, "skip" => commands::audio::skip::run(&ctx, &command).await, "volume" => commands::audio::volume::run(&ctx, &command).await, _ => { let content: String = match command.data.name.as_str() { "ping" => commands::ping::run(&command.data.options), _ => "Unknown command".to_string() }; if let Err(why) = create_response(&ctx, &command, content).await { warn!("Cannot respond to slash command: {}", why); } } } } } async fn ready(&self, ctx: Context, ready: Ready) { if ready.guilds.is_empty() { warn!("No ready guilds found"); } for guild in ready.guilds { let audio_config_lock = { let data_read = ctx.data.read().await; data_read.get::().expect("Expected AudioConfigs in TypeMap.").clone() }; { let mut audio_configs = audio_config_lock.write().await; let _ = audio_configs.insert(guild.id, AudioConfig { volume: 1.0 }); } let commands = guild.id.set_application_commands(&ctx.http, |commands| { commands.create_application_command(|command: &mut serenity::builder::CreateApplicationCommand| { commands::ping::register(command) }) .create_application_command(|command: &mut serenity::builder::CreateApplicationCommand| { commands::audio::play::register(command) }) .create_application_command(|command: &mut serenity::builder::CreateApplicationCommand| { commands::audio::stop::register(command) }) .create_application_command(|command: &mut serenity::builder::CreateApplicationCommand| { commands::audio::pause::register(command) }) .create_application_command(|command: &mut serenity::builder::CreateApplicationCommand| { commands::audio::resume::register(command) }) .create_application_command(|command: &mut serenity::builder::CreateApplicationCommand| { commands::audio::skip::register(command) }) .create_application_command(|command: &mut serenity::builder::CreateApplicationCommand| { commands::audio::volume::register(command) }) }).await; match commands { Ok(c) => info!("Registered {} commands for guild {}", c.len(), guild.id.0), Err(why) => error!("Could not register commands for guild {}: {:?}", guild.id.0, why) }; } } } #[tokio::main] async fn main() { dotenv().ok(); env_logger::init_from_env(env_logger::Env::default().filter_or("RUST_LOG", "warn,siren=info")); let token: String = env::var("DISCORD_TOKEN").expect("Expected a token in the environment"); let intents: GatewayIntents = GatewayIntents::all(); let http: Http = Http::new(&token); let (owners, _bot_id) = match http.get_current_application_info().await { Ok(info) => { let mut owners: HashSet = HashSet::new(); if let Some(team) = info.team { owners.insert(team.owner_user_id); } else { owners.insert(info.owner.id); } match http.get_current_user().await { Ok(bot) => (owners, bot.id), Err(why) => panic!("Could not access the bot id: {:?}", why) } }, Err(why) => panic!("Could not access application info: {:?}", why) }; let pool = database::establish_connection(); database::run_migrations(&pool); let handler = match env::var("OPENAI_API_KEY") { Ok(token) => { info!("Loaded OpenAI token"); Handler { oai: Some(commands::oai::OAI { client: reqwest::Client::new(), base_url: "https://api.openai.com/v1".to_string(), max_attempts: 5, token , max_context_questions: 15 }), pool } } Err(err) => { warn!("Could not load OpenAI token: {}", err); Handler { oai: None, pool } } }; let mut client = Client::builder(token, intents) .event_handler(handler) .framework(StandardFramework::new() .configure(|c| c.owners(owners))) .register_songbird() .await .expect("Error creating client"); { let mut data = client.data.write().await; data.insert::(Arc::new(RwLock::new(HashMap::default()))); } if let Err(why) = client.start_autosharded().await { error!("An error occurred while running the client: {:?}", why); } }