siren/src/main.rs

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<commands::oai::OAI>,
  pool: Pool<ConnectionManager<PgConnection>>
}

#[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::<AudioConfigs>().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<serenity::model::id::UserId> = 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::<AudioConfigs>(Arc::new(RwLock::new(HashMap::default())));
  }

  if let Err(why) = client.start_autosharded().await {
    error!("An error occurred while running the client: {:?}", why);
  }
}