path: root/src/bot.rs

use std::{
    collections::HashSet,
    fs::File,
    future::Future,
    path::PathBuf,
    pin::Pin,
    str::FromStr,
    sync::Arc,
};

use chrono::Datelike;
use fnv::{
    FnvHashMap,
    FnvHashSet,
};
use lazy_static::lazy_static;
use log::{
    debug,
    error,
    info,
    trace,
    warn,
};
use poise::{
    BoxFuture,
    FrameworkError,
    PrefixContext,
};
use serenity::{
    all::{
        GuildId,
        ReactionType,
    },
    builder::CreateMessage,
    model::{
        event::ResumedEvent,
        gateway::Ready,
        id::{
            ChannelId,
            MessageId,
        },
    },
    prelude::*,
};
use songbird::{
    Call,
    Event,
    EventContext,
    SerenityInit,
    TrackEvent,
};
use tokio::sync::Mutex;

use crate::{
    commands,
    config::CONFIG,
    err_msg,
    util,
    util::OAUTH_URL,
    PoiseContext,
    PoiseData,
};

pub struct HttpKey;

impl TypeMapKey for HttpKey {
    type Value = reqwest::Client;
}

struct Handler;

#[serenity::async_trait]
impl EventHandler for Handler {
    async fn ready(&self, ctx: Context, r: Ready) {
        let guild = r.guilds.iter().find(|g| g.id == CONFIG.discord.guild());

        if guild.is_none() {
            info!("bot isn't in configured guild. join here: {:?}", OAUTH_URL.as_str());
            return;
        }

        info!("connected");

        #[cfg(debug_assertions)]
        let botname = "thulani (dev)";

        #[cfg(not(debug_assertions))]
        let botname = "thulani";

        if let Some(guild) = guild {
            if let Err(e) = guild.id.edit_nickname(&ctx, Some(botname)).await {
                error!("changing nickname: {:?}", e);
            }
        }

        let sb = songbird::get(&ctx).await.unwrap();

        let c = sb.get_or_insert(CONFIG.discord.guild());
        let mut call = c.lock().await;

        call.remove_all_global_events();

        call.add_global_event(Event::Track(TrackEvent::End), SongbirdHandler(c.clone()));
    }

    async fn resume(&self, _ctx: Context, _resume: ResumedEvent) {
        info!("reconnected to discord");
    }

    async fn message_delete(
        &self,
        _ctx: Context,
        _: ChannelId,
        deleted_message_id: MessageId,
        _: Option<GuildId>,
    ) {
        MESSAGE_WATCH.lock().await.remove(&deleted_message_id);
    }
}

struct SongbirdHandler(Arc<Mutex<Call>>);

#[serenity::async_trait]
impl songbird::events::EventHandler for SongbirdHandler {
    async fn act(&self, _ctx: &EventContext<'_>) -> Option<Event> {
        let mut call = self.0.lock().await;
        if call.queue().is_empty() {
            let _ = call.leave().await;
        }

        None
    }
}

lazy_static! {
    static ref MESSAGE_WATCH: Mutex<FnvHashMap<MessageId, MessageId>> =
        Mutex::new(FnvHashMap::default());
    static ref PREFIXES: Vec<&'static str> =
        vec!["!thulani ", "!thulan ", "!thulando madando ", "!thulando "];
    static ref RESTRICTED_PREFIXES: Vec<&'static str> = vec!["!todd ", "!toddbert ", "!toddlani "];
    static ref ALL_PREFIXES: Vec<&'static str> = {
        let mut all_prefixes: Vec<&'static str> = vec![];
        all_prefixes.extend(PREFIXES.iter());
        all_prefixes.extend(RESTRICTED_PREFIXES.iter());
        all_prefixes
    };
    static ref RESTRICT_IDS: FnvHashSet<u64> = {
        let default_path = PathBuf::from_str("restrict.json").unwrap();
        let restrict_path = CONFIG.restrict.as_ref().unwrap_or(&default_path);

        let restrict_ids = File::open(restrict_path)
            .map_err(anyhow::Error::from)
            .and_then(|f| serde_json::from_reader::<_, Vec<u64>>(f).map_err(anyhow::Error::from));

        if let Err(ref e) = restrict_ids {
            warn!("opening restrict file: {}", e);
        }

        let result = restrict_ids.unwrap_or_default().into_iter().collect::<FnvHashSet<_>>();

        info!("restricted ids: {result:?}");

        result
    };
}

fn on_err(err: FrameworkError<PoiseData, anyhow::Error>) -> BoxFuture<()> {
    Box::pin(async move {
        let Some(msg) = err_msg(&err) else {
            warn!("error handler missing poise context");
            return;
        };

        let ctx = err.serenity_context();

        let text = match err {
            FrameworkError::ArgumentParse {
                ..
            }
            | FrameworkError::SubcommandRequired {
                ..
            } => "format your commands right. fuck you.".to_string(),
            FrameworkError::CooldownHit {
                ..
            } => "slow the fuck down bitch".to_string(),
            FrameworkError::NotAnOwner {
                ..
            } => "who do you think you are?".to_string(),
            FrameworkError::GuildOnly {
                ..
            } => "what in the sam hill are you smoking".to_string(),
            FrameworkError::DmOnly {
                ..
            } => "take that back or i'm revoking your kitten status".to_string(),
            FrameworkError::UnknownCommand {
                ctx,
                msg,
                prefix,
                msg_content,
                trigger,
                invocation_data,
                framework,
                ..
            } => {
                let command = poise::Command {
                    name: "meme".to_owned(),
                    ..Default::default()
                };

                fn noop<U, E>(
                    _ctx: PrefixContext<'_, U, E>,
                ) -> BoxFuture<Result<(), FrameworkError<U, E>>> {
                    Box::pin(async { Ok(()) })
                }

                let ctx = PrefixContext {
                    serenity_context: ctx,
                    prefix,
                    msg,
                    command: &command,
                    trigger,
                    invocation_data,
                    parent_commands: &[],
                    data: &(),
                    invoked_command_name: "",
                    action: noop,
                    args: msg_content,
                    framework,

                    __non_exhaustive: (),
                };

                match util::pop_string(msg_content)
                    .map_err(anyhow::Error::from)
                    .and_then(|(_rest, s)| s.parse().map_err(anyhow::Error::from))
                {
                    Ok(u) => {
                        if let Err(e) = commands::unrecognized(PoiseContext::Prefix(ctx), u).await {
                            error!("processing audio: {e}");
                            "BANIC".to_string()
                        } else {
                            return;
                        }
                    },
                    Err(e) => {
                        error!("processing unrecognized message: {e}");
                        "BANIC".to_string()
                    },
                }
            },
            _ => "BANIC".to_string(),
        };

        error!("error encountered: {err:#?}");
        if let Err(e) = msg.react(ctx, ReactionType::Unicode("❌".to_owned())).await {
            error!("reacting to failed message: {e}");
        }

        let cm = CreateMessage::default().content(text).tts(msg.tts);
        if let Err(e) = msg.channel_id.send_message(ctx, cm).await {
            error!("sending error to chat: {e}");
        }
    })
}

async fn framework() -> poise::Framework<PoiseData, anyhow::Error> {
    let additional_prefixes =
        ALL_PREFIXES.iter().skip(1).map(|x| poise::Prefix::Literal(x.to_owned())).collect();

    let framework = poise::Framework::builder()
        .options(poise::FrameworkOptions {
            pre_command: before_handle,
            post_command: after_handle,
            on_error: on_err,

            command_check: Some(check),

            prefix_options: poise::PrefixFrameworkOptions {
                prefix: ALL_PREFIXES.get(0).map(|&x| x.to_owned()),
                additional_prefixes,
                case_insensitive_commands: true,
                mention_as_prefix: false,
                ignore_bots: true,
                ..Default::default()
            },

            commands: commands::commands(),
            owners: HashSet::from_iter([CONFIG.discord.owner()]),
            initialize_owners: false,
            skip_checks_for_owners: true,

            ..Default::default()
        })
        .setup(|_ctx, _ready, _framework| Box::pin(async move { Ok(()) }))
        .build();

    framework
}

fn check(ctx: PoiseContext) -> BoxFuture<anyhow::Result<bool>> {
    Box::pin(async move {
        if !ctx.guild_id().map_or(false, |x| x == CONFIG.discord.guild()) {
            info!(
                "rejecting command '{}' from user '{}': wrong guild",
                ctx.command().name,
                ctx.author().name
            );
            return Ok(false);
        }

        if ctx.author().id == CONFIG.discord.owner() {
            return Ok(true);
        }

        let restricted_prefix = RESTRICTED_PREFIXES.iter().any(|&prefix| ctx.prefix() == prefix);

        if !restricted_prefix {
            return Ok(true);
        }

        const PERMITTED_WEEKDAY: chrono::Weekday = chrono::Weekday::Tue;

        let user_is_restricted = RESTRICT_IDS.contains(&ctx.author().id.get());
        let restrictions_flipped = chrono::Local::now().weekday() == PERMITTED_WEEKDAY;

        if user_is_restricted == restrictions_flipped {
            return Ok(true);
        }

        let reason = if !restrictions_flipped {
            "restricted prefix".to_owned()
        } else {
            format!("it is {PERMITTED_WEEKDAY:?}")
        };

        info!(
            "rejecting command '{}' from user '{}': {}",
            ctx.command().name,
            ctx.author().name,
            reason
        );

        util::reply(ctx, "no").await?;

        Ok(false)
    })
}

fn before_handle<'fut>(ctx: PoiseContext<'fut>) -> Pin<Box<dyn Future<Output = ()> + Send + 'fut>> {
    debug!(
        "got command '{}' from user '{}' ({})",
        ctx.command().name,
        ctx.author().name,
        ctx.author().id
    );

    Box::pin(async {})
}

fn after_handle(ctx: PoiseContext) -> BoxFuture<()> {
    Box::pin(async move {
        trace!("command '{}' completed successfully", ctx.command().name);
    })
}

pub async fn run() -> anyhow::Result<()> {
    #[cfg(all(windows, feature = "windows_autostart_postgres"))]
    unsafe {
        crate::windows_util::ensure_postgres_started()?;
    };

    let token = &CONFIG.discord.auth.token;

    let sb_config = songbird::Config::default();

    let mut client =
        Client::builder(token, GatewayIntents::non_privileged() | GatewayIntents::MESSAGE_CONTENT)
            .event_handler(Handler)
            .register_songbird_from_config(sb_config)
            .type_map_insert::<HttpKey>(reqwest::Client::new())
            .framework(framework().await)
            .await?;

    let shard_manager = client.shard_manager.clone();

    tokio::spawn(async move {
        tokio::signal::ctrl_c().await.unwrap();
        warn!("got ^C");

        shard_manager.shutdown_all().await;
        info!("shutdown");
    });

    client.start().await?;

    Ok(())
}