siren/service/src/main.rs

extern crate diesel;
#[macro_use]
extern crate diesel_migrations;

use std::env;
use std::collections::HashSet;
use std::sync::Arc;
use log::{error, warn, info};
use serenity::client::Cache;
use serenity::framework::StandardFramework;
use serenity::http::Http;
use serenity::prelude::*;
use songbird::{SerenityInit, Songbird};

use actix_cors::Cors;
use actix_web::{HttpServer, App, web};
use crate::bot::handler::Handler;

mod auth;
mod dnd;
mod bot;
mod storage;
mod users;

#[actix_web::main]
async fn main() -> std::io::Result<()> {
  dotenv::dotenv().ok();
  env_logger::init_from_env(env_logger::Env::default().filter_or("RUST_LOG", "warn,siren=info"));
  storage::init().await;
  match env::var("DATA_DIR_PATH") {
    Ok(data_dir_path) => dnd::load_data(&data_dir_path),
    Err(err) => warn!("Unable to load initial database data: {}", err)
  };

  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 handler = match env::var("OPENAI_API_KEY") {
    Ok(token) => {
      info!("Loaded OpenAI token");
      let default_model = env::var("OPENAI_API_MODEL").unwrap_or("gpt-3.5-turbo".to_string());
      Handler {
        oai: Some(bot::oai::OAI {
          client: reqwest::Client::new(),
          base_url: "https://api.openai.com/v1".to_string(),
          service_url: "http://localhost:5000".to_string(),
          max_attempts: 5,
          token,
          max_context_questions: 30,
          max_tokens: 2048,
          default_model,
        })
      }
    }
    Err(err) => {
      warn!("Could not load OpenAI token: {}", err);
      Handler { oai: None }
    }
  };

  let songbird = Songbird::serenity();

  let mut client = Client::builder(token, intents)
    .event_handler(handler)
    .framework(StandardFramework::new()
    .configure(|c| c.owners(owners)))
    .register_songbird_with(Arc::clone(&songbird))
    .await
    .expect("Error creating client");

  let http = Arc::clone(&client.cache_and_http.http);
  let cache = Arc::clone(&client.cache_and_http.cache);

  let app_data = Arc::new(AppState {
    http,
    cache,
    songbird: Arc::clone(&songbird)
  });

  
  let shard_manager = Arc::clone(&client.shard_manager);

  tokio::spawn(async move {
    tokio::signal::ctrl_c().await.expect("Could not register ctrl+c handler");
    shard_manager.lock().await.shutdown_all().await;
  });

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

  let host = env::var("SERVICE_HOST").unwrap_or("localhost".to_string());
  let port = env::var("SERVICE_PORT").unwrap_or("5000".to_string());

  let server = match HttpServer::new(move || {
    let cors = Cors::default()
      .allow_any_origin()
      .allow_any_method()
      .allow_any_header()
      .supports_credentials()
      .max_age(3600);
    App::new()
      .wrap(cors)
      .app_data(web::Data::new(Arc::clone(&app_data)))
      .configure(crate::auth::init_routes)
      .configure(crate::users::init_routes)
      .configure(crate::dnd::spells::init_routes)
      .configure(crate::bot::guilds::init_routes)
      .configure(crate::bot::messages::init_routes)
  })
  .bind(format!("{}:{}", host, port)) {
    Ok(b) => {
      info!("Binding server to {}:{}", host, port);
      b
    },
    Err(err) => {
      error!("Could not bind server: {}", err);
      return Err(err);
    }
  };

  server.run()
  .await
}

pub struct AppState {
  pub http: Arc<Http>,
  pub cache: Arc<Cache>,
  pub songbird: Arc<Songbird>
}