FrenBot/src/models/task.rs

use crate::config::GlobalData;
use crate::discord::shop::restock_shop;
use crate::error::Error;
use crate::models::birthday::BirthdayEntry;
use crate::models::insult_compliment::{RandomResponseTemplate, ResponseType};
use chrono::{Days, Duration, TimeZone, Timelike, Utc};
use j_db::database::Database;
use j_db::model::JdbModel;
use serde::{Deserialize, Serialize};
use serenity::all::Mentionable;
use serenity::prelude::Context;

#[derive(Debug, Deserialize, Serialize, Clone, Copy, PartialEq)]
pub enum TaskType {
    RemoveRole { user_id: u64, role_id: u64 },
    CheckBirthdays,
    HandleReload,
    RestockShop,
}

impl TaskType {
    pub fn exclusive(&self) -> bool {
        matches!(
            self,
            TaskType::CheckBirthdays | TaskType::HandleReload | TaskType::RestockShop
        )
    }
}

#[derive(Debug, Deserialize, Serialize, Clone)]
pub struct Task {
    id: Option<u64>,
    pub task_type: TaskType,
    pub time: chrono::DateTime<Utc>,
}

impl JdbModel for Task {
    fn id(&self) -> Option<u64> {
        self.id
    }

    fn set_id(&mut self, id: u64) {
        self.id = Some(id)
    }

    fn tree() -> String {
        "tasks".to_string()
    }

    fn check_unique(&self, _other: &Self) -> bool {
        true
    }
}

impl Task {
    pub fn add_task(
        db: &Database,
        task_type: TaskType,
        time: chrono::DateTime<chrono::Utc>,
    ) -> Result<(), Error> {
        if task_type.exclusive() {
            let old_tasks: Vec<Task> = db
                .filter(|_, task: &Task| task.task_type == task_type)?
                .collect();

            for task in old_tasks {
                db.remove::<Task>(task.id().unwrap())?;
            }
        }

        db.insert::<Task>(Task {
            id: None,
            task_type,
            time,
        })?;

        println!("Adding {:?} task to run at {}", task_type, time);

        Ok(())
    }

    pub async fn create_reoccurring_tasks(ctx: &Context) -> Result<(), Error> {
        let mut data = ctx.data.write().await;
        let global_data = data.get_mut::<GlobalData>().unwrap();

        Task::add_task(&global_data.db, TaskType::CheckBirthdays, Utc::now())?;
        Task::add_task(&global_data.db, TaskType::RestockShop, Utc::now())?;
        Task::add_task(
            &global_data.db,
            TaskType::HandleReload,
            Utc::now() + Duration::hours(1),
        )?;

        Ok(())
    }

    pub async fn run_tasks(ctx: &Context) -> Result<(), Error> {
        let mut data = ctx.data.write().await;
        let global_data = data.get_mut::<GlobalData>().unwrap();

        let active_tasks: Vec<Task> = global_data
            .db
            .filter(|_, task: &Task| task.time < chrono::Utc::now())?
            .collect();

        for task in active_tasks {
            match task.task_type {
                TaskType::RemoveRole { user_id, role_id } => {
                    let user = global_data
                        .cfg
                        .guild_id
                        .member(&ctx.http, user_id)
                        .await
                        .unwrap();

                    println!("Removing role {} from {}", role_id, user.display_name());
                    user.remove_role(&ctx.http, role_id).await.unwrap();
                }
                TaskType::CheckBirthdays => {
                    println!("Checking Birthdays");
                    let todays_birthdays = BirthdayEntry::todays_birthdays(
                        &global_data.db,
                        chrono::Utc::now().date_naive(),
                    )
                    .unwrap();

                    for birth in todays_birthdays {
                        if let Ok(user) = global_data
                            .cfg
                            .guild_id
                            .member(&ctx.http, birth.discord_id)
                            .await
                        {
                            global_data
                                .cfg
                                .announcement_channel
                                .say(&ctx.http, format!("Happy birthday {}!", user.mention()))
                                .await
                                .unwrap();

                            let compliment = RandomResponseTemplate::get_random_response(
                                &global_data.db,
                                ResponseType::Compliment,
                                user.display_name(),
                            )
                            .unwrap()
                            .unwrap();
                            global_data
                                .cfg
                                .announcement_channel
                                .say(&ctx.http, compliment)
                                .await
                                .unwrap();
                        }
                    }

                    let next_check = chrono_tz::America::Chicago
                        .from_utc_datetime(&Utc::now().naive_utc())
                        .with_hour(8)
                        .unwrap()
                        .with_minute(0)
                        .unwrap()
                        .with_second(0)
                        .unwrap()
                        .checked_add_days(Days::new(1))
                        .unwrap();

                    Task::add_task(
                        &global_data.db,
                        TaskType::CheckBirthdays,
                        next_check.with_timezone(&Utc),
                    )?;
                }
                TaskType::HandleReload => {
                    println!("Reloading config...");
                    let res = global_data.reload().await;

                    match res {
                        Ok(_) => println!("Finished reloading config!"),
                        Err(err) => println!("Error reloading config: {:?}", err),
                    }

                    Task::add_task(
                        &global_data.db,
                        TaskType::HandleReload,
                        Utc::now() + Duration::hours(1),
                    )?;
                }
                TaskType::RestockShop => {
                    println!("Restocking Shop...");
                    let res = restock_shop(ctx, global_data).await;

                    match res {
                        Ok(_) => println!("Finished restocking shop!"),
                        Err(err) => println!("Error restocking shop: {:?}", err),
                    }

                    Task::add_task(
                        &global_data.db,
                        TaskType::RestockShop,
                        Utc::now() + Duration::hours(1),
                    )?;
                }
            }

            let _ = global_data.db.remove::<Task>(task.id().unwrap()).is_ok();
        }

        Ok(())
    }
}