bingus-bot / src / main.rs

#![allow(unused_features)]
#![feature(iter_map_windows)]
#![feature(iter_intersperse)]
#![feature(test)]

mod brain;
mod cmd;
mod on_message;
mod status;

pub mod prelude {
    pub use anyhow::Context;
    use std::result::Result as StdResult;
    pub type Result<T = (), E = anyhow::Error> = StdResult<T, E>;
}

use std::{
    collections::HashSet,
    fs::File,
    path::{Path, PathBuf},
    sync::{
        Arc,
        atomic::{AtomicBool, Ordering},
    },
};

use brotli::enc::{BrotliEncoderParams, backward_references::BrotliEncoderMode};
use log::{debug, error, info, warn};
use prelude::*;
use tokio::{
    signal::unix::{SignalKind, signal},
    sync::RwLock,
    time::{self, Duration},
};
use twilight_gateway::{
    CloseFrame, Event, EventTypeFlags, Intents, MessageSender, Shard, ShardId, StreamExt,
};
use twilight_http::Client as HttpClient;
use twilight_model::{
    application::interaction::InteractionData,
    id::{
        Id,
        marker::{ApplicationMarker, ChannelMarker, UserMarker},
    },
};

use crate::{
    brain::Brain,
    cmd::{handle_app_command, register_all_commands},
    on_message::handle_discord_message,
    status::update_status,
};

pub type BrainHandle = RwLock<Brain>;

#[derive(Debug)]
pub struct BotContext {
    http: HttpClient,
    self_id: Id<UserMarker>,
    app_id: Id<ApplicationMarker>,
    owners: HashSet<Id<UserMarker>>,
    brain_file_path: PathBuf,
    reply_channels: HashSet<Id<ChannelMarker>>,
    brain_handle: BrainHandle,
    shard_sender: MessageSender,
    pending_save: AtomicBool,
}

async fn handle_discord_event(event: Event, ctx: Arc<BotContext>) -> Result {
    match event {
        Event::MessageCreate(msg) => handle_discord_message(msg, ctx)
            .await
            .context("While handling a new message"),
        Event::InteractionCreate(mut inter) => {
            if let Some(InteractionData::ApplicationCommand(data)) =
                std::mem::take(&mut inter.0.data)
            {
                handle_app_command(*data, ctx, inter.0)
                    .await
                    .context("While handling an app command")
            } else {
                Ok(())
            }
        }
        Event::Ready(ev) => {
            info!("Connected to gateway as {}", ev.user.name);
            let brain = ctx.brain_handle.read().await;
            update_status(&brain, &ctx.shard_sender).context("Failed to update status on ready")
        }
        _ => Ok(()),
    }
}

const BROTLI_BUF_SIZE: usize = 1024 * 1000;
fn get_brotli_params() -> BrotliEncoderParams {
    BrotliEncoderParams {
        quality: 5,
        mode: BrotliEncoderMode::BROTLI_MODE_TEXT,
        ..Default::default()
    }
}

fn load_brain(path: &Path) -> Result<Option<Brain>> {
    if path.exists() {
        let mut file = File::open(path).context("Failed to open brain file")?;
        let mut brotli_stream = brotli::Decompressor::new(&mut file, BROTLI_BUF_SIZE);
        rmp_serde::from_read(&mut brotli_stream)
            .map(Some)
            .context("Failed to decode brain file")
    } else {
        Ok(None)
    }
}

async fn save_brain(ctx: Arc<BotContext>) -> Result {
    let scratch_path = ctx.brain_file_path.with_file_name(format!(
        "~{}",
        ctx.brain_file_path.file_name().unwrap().to_str().unwrap()
    ));
    let mut file = File::create(&scratch_path).context("Failed to open brain file")?;
    let mut brotli_writer =
        brotli::CompressorWriter::with_params(&mut file, BROTLI_BUF_SIZE, &get_brotli_params());

    let brain = ctx.brain_handle.read().await;
    rmp_serde::encode::write(&mut brotli_writer, &*brain)
        .context("Failed to write serialized brain")?;

    std::fs::rename(&scratch_path, &ctx.brain_file_path)
        .context("Failed to override scratch file")?;

    debug!("Saved brain file");
    Ok(())
}

#[tokio::main]
async fn main() -> Result {
    let mut clog = colog::default_builder();
    clog.filter(
        None,
        if cfg!(debug_assertions) {
            log::LevelFilter::Debug
        } else {
            log::LevelFilter::Info
        },
    );
    clog.try_init().context("Failed to initialize colog")?;

    info!("Start of bingus-bot {}", env!("CARGO_PKG_VERSION"));

    // Config
    let token_file = std::env::var("TOKEN_FILE").context("Missing TOKEN_FILE env var")?;
    let reply_channels = std::env::var("REPLY_CHANNELS")
        .context("Missing REPLY_CHANNELS env var")?
        .split(",")
        .filter_map(|s| {
            if s.trim().is_empty() {
                None
            } else {
                Some(s.trim().parse::<u64>().map(|c| Id::new(c)))
            }
        })
        .collect::<Result<_, _>>()
        .context("Invalid channel IDs for REPLY_CHANNELS")?;
    let brain_file_path =
        PathBuf::from(std::env::var("BRAIN_FILE").unwrap_or_else(|_| "brain.msgpackz".to_string()));
    let intents = Intents::GUILD_MESSAGES | Intents::MESSAGE_CONTENT;

    // Read token
    let token = std::fs::read_to_string(token_file).context("Failed to read bot token")?;
    let token = token.trim();

    // Read Brain
    let brain = if let Some(brain) = load_brain(&brain_file_path)? {
        info!("Loading brain from {brain_file_path:?}");
        brain
    } else {
        info!("Creating new brain file at {brain_file_path:?}");
        Brain::default()
    };
    let brain_handle = RwLock::new(brain);

    // Init
    let mut shard = Shard::new(ShardId::ONE, token.to_string(), intents);
    let http = HttpClient::new(token.to_string());

    let app = http
        .current_user_application()
        .await
        .context("Failed to get current App")?
        .model()
        .await
        .context("Failed to deserialize")?;

    let app_id = app.id;

    let self_id = app.bot.context("App is not a bot!")?.id;

    let owners = if let Some(user) = app.owner {
        HashSet::from_iter([user.id])
    } else if let Some(team) = app.team {
        team.members.iter().map(|m| m.user.id).collect()
    } else {
        warn!("No Owner?? Bingus is free!!!");
        HashSet::new()
    };

    let context = Arc::new(BotContext {
        http,
        self_id,
        app_id,
        owners,
        reply_channels,
        brain_file_path,
        brain_handle,
        shard_sender: shard.sender(),
        pending_save: AtomicBool::new(false),
    });

    info!("Registering Commands...");
    register_all_commands(context.clone()).await?;

    let mut interval = time::interval(Duration::from_secs(60));
    interval.tick().await;

    let mut sigterm = signal(SignalKind::terminate()).context("Failed to listen to SIGTERM")?;

    info!("Connecting to gateway...");

    loop {
        tokio::select! {

            biased;

            Ok(()) = tokio::signal::ctrl_c() => {
                info!("SIGINT: Closing connection and saving");
                shard.close(CloseFrame::NORMAL);
            }
            _ = sigterm.recv() => {
                info!("SIGTERM: Closing connection and saving");
                shard.close(CloseFrame::NORMAL);
            }
            _ = interval.tick() => {
                debug!("Save Interval");
                if context.pending_save.load(Ordering::Relaxed) {
                   let ctx = context.clone();
                   tokio::spawn(async move {
                        if let Err(why) = save_brain(ctx.clone()).await {
                            error!("Failed to save brain file:\n{why:?}");
                        }
                        ctx.pending_save.store(false, Ordering::Relaxed);
                   });
                }
            },
            opt = shard.next_event(EventTypeFlags::all()) => {
                match opt {
                    Some(Ok(Event::GatewayClose(_))) | None => {
                        info!("Disconnected from Discord");
                        break;
                    }
                    Some(Ok(event)) => {
                        let ctx = context.clone();
                        tokio::spawn(async move {
                            if let Err(why) = handle_discord_event(event, ctx).await {
                                error!("Error while processing Discord event:\n{why:?}");
                            }
                        });
                    }
                    Some(Err(why)) => {
                        warn!("Failed to receive event:\n{why:?}");
                    }
                }
            }
        }
    }

    if context.pending_save.load(Ordering::Relaxed) {
        save_brain(context)
            .await
            .context("Failed to write brain file on exit")?;
    }

    Ok(())
}