badgey/src/badgey/models/xp.rs

use built::chrono;
use diesel::prelude::*;
use diesel::insert_into;
use rand::rngs::SmallRng;
use rand::{Rng, SeedableRng};

use manifold::{Db, ManifoldData};
use manifold::error::{ManifoldError, ManifoldResult};
use manifold::models::user::UserInfo;
use manifold::schema::userinfo;
use poise::FrameworkContext;
use poise::serenity_prelude::{Context, Mentionable, Message, RoleId};
use crate::badgey::models::quarantine_channel::QuarantineChannel;
use crate::badgey::schema::xp as xp_table;
use crate::badgey::schema::*;

#[derive(Queryable, Selectable, Identifiable, Insertable, AsChangeset, Associations, Debug, PartialEq, Clone)]
#[diesel(belongs_to(UserInfo, foreign_key = user_id))]
#[diesel(table_name = xp_table)]
#[diesel(primary_key(user_id))]
#[diesel(treat_none_as_null = true)]
pub struct Xp {
    pub user_id: i64,
    pub user_current_level: i64,
    pub xp_value: i64,
    pub last_given_xp: Option<i64>,
    pub rank_track: i64,
    pub rank_track_last_changed: Option<i64>,
    pub freeze_rank: Option<i64>,
    pub freeze_rank_last_changed: Option<i64>
}

#[derive(Queryable, Selectable, Identifiable, Debug, Clone)]
#[diesel(primary_key(role_id))]
pub struct Rank {
    pub role_id: i64,
    pub required_level: i64,
    pub rank_track: i64,
    pub rank_name: String,
}

#[derive(Queryable, Selectable, Identifiable, Debug, Clone, Default)]
#[diesel(primary_key(track_id))]
pub struct Track {
    pub track_id: i64,
    pub track_name: String,
    pub xp_level_constant: f64,
    pub xp_award_range_min: i32,
    pub xp_award_range_max: i32,
}

impl Xp {
    pub fn new(user_id: &i64) -> Self {
        Self {
            user_id: user_id.clone(),
            user_current_level: 0,
            xp_value: 0,
            last_given_xp: None,
            rank_track: 0,
            rank_track_last_changed: None,
            freeze_rank: None,
            freeze_rank_last_changed: None,
        }
    }

    pub fn get(conn: &Db, needle: &i64) -> ManifoldResult<Self> {
        let user = userinfo::table
            .filter(userinfo::user_id.eq(needle))
            .select(UserInfo::as_select())
            .get_result(&mut conn.get()?)?;

        Ok(Xp::belonging_to(&user)
            .select(Xp::as_select())
            .get_result(&mut conn.get()?)?)
    }

    pub fn insert(&self, conn: &Db) -> ManifoldResult<usize> {
        insert_into(xp_table::dsl::xp)
            .values(self)
            .on_conflict(xp_table::user_id)
            .do_update()
            .set(self)
            .execute(&mut conn.get()?)
            .map_err(|e| ManifoldError::from(e))
    }

    pub fn get_level_from_xp(&self, conn: &Db, constant: Option<&f64>) -> i64 {
        let c = match constant {
            Some(c) => *c,
            None => Track::get_track_by_id(conn, &self.rank_track).unwrap_or_default().xp_level_constant
        };
        (c * f64::sqrt(self.xp_value.clone() as f64)) as i64
    }

    pub fn get_xp_to_next_level(&self, conn: &Db) -> i64 {
        let constant = Track::get_track_by_id(conn, &self.rank_track).unwrap_or_default().xp_level_constant;
        let current_level = self.get_level_from_xp(conn, Some(&constant));
        let target_level = current_level + 1;

        (target_level as f64 / constant).powi(2) as i64 - &self.xp_value
    }

    pub fn get_xp_to_next_rank(&self, conn: &Db) -> ManifoldResult<i64> {
        let constant = Track::get_track_by_id(conn, &self.rank_track)?.xp_level_constant;
        let current_level = self.get_level_from_xp(conn, Some(&constant));
        let target_level = Rank::get_next_rank_for_level(conn, &current_level, &self.rank_track)?.required_level;

        Ok((target_level as f64 / constant).powi(2) as i64 - &self.xp_value)
    }

    pub fn check_timeout(value: &Option<i64>, timeout: &i64) -> ManifoldResult<()> {
        if let Some(v) = value {
            if v > &(chrono::Utc::now().timestamp() - timeout) {
                let timeuntil = v - (chrono::Utc::now().timestamp() - timeout);
                Err(t!("models.xp.cooldown", timeuntil = timeuntil))?;
            }
        }

        Ok(())
    }

    pub async fn award(ctx: &Context, fctx: &FrameworkContext<'_, ManifoldData, ManifoldError>, msg: &Message, db: &Db) -> ManifoldResult<()>{
        if fctx.user_data().await.user_info.lock().await.get_mut(&msg.author.id.as_u64()).is_none() {
            debug!("Tried to add XP to a user we don't know about, aborting.");
            return Ok(())
        }

        let mut rng = SmallRng::from_entropy();
        let xp_reward = rng.gen_range(1..30).clone();
        drop(rng);
        let user_id_i64 = msg.author.id.as_u64().clone() as i64;

        let mut xp = match Xp::get(db, &user_id_i64) {
            Ok(x) => x,
            Err(_) => Xp::new(&user_id_i64)
        };

        let valid = match xp.last_given_xp {
            Some(t) => (chrono::Utc::now().timestamp() - 60) > t && QuarantineChannel::get(&fctx.user_data.database, msg.channel_id.as_u64().clone() as i64).is_err(),
            None => true
        };

        if valid {
            xp.last_given_xp = Some(chrono::Utc::now().timestamp());
            xp.xp_value += &xp_reward;
            let calculated_level = xp.get_level_from_xp(&db, None);

            if xp.freeze_rank.is_some() {
                xp.user_current_level = calculated_level.clone();
                xp.insert(db)?;
                return Ok(())
            }

            if (xp.user_current_level != calculated_level) || xp.user_current_level == 1 {
                if let Some(guild) = msg.guild(&ctx.cache) {
                    let mut member = guild.member(&ctx, msg.author.id).await?;
                    let old_role_id = RoleId::from(Rank::get_rank_for_level(db, &xp.user_current_level, &xp.rank_track).unwrap_or(Rank::new()).role_id as u64);
                    xp.user_current_level = calculated_level;
                    let new_role = RoleId::from(Rank::get_rank_for_level(db, &calculated_level, &xp.rank_track)?.role_id as u64);

                    if (new_role != old_role_id || xp.user_current_level == 1) && msg.author.has_role(ctx, guild.id, new_role).await? == false {
                        member.remove_role(ctx, old_role_id).await?;
                        member.add_role(ctx, new_role).await?;

                        msg.channel_id.say(ctx, format!("{} is now a {}", msg.author.mention(), new_role.mention())).await?;
                    }
                } else {
                    error!("Needed to add level to user, but couldn't get the member from the guild!");
                }
            }

            xp.insert(db)?;

        }

        Ok(())
    }
}

impl Rank {
    pub fn new() -> Self {
        Self {
            role_id: 0,
            required_level: 0,
            rank_track: 0,
            rank_name: t!("models.xp.unranked").to_string(),
        }
    }

    pub fn get_rank_for_level(conn: &Db, level: &i64, track: &i64) -> ManifoldResult<Self> {
        Ok(ranks::table
            .filter(ranks::required_level.le(level))
            .filter(ranks::rank_track.eq(track))
            .order(ranks::required_level.desc())
            .limit(1)
            .get_result(&mut conn.get()?)?
        )
    }

    pub fn get_next_rank_for_level(conn: &Db, level: &i64, track: &i64) -> ManifoldResult<Self> {
        Ok(ranks::table
            .filter(ranks::required_level.gt(level))
            .filter(ranks::rank_track.eq(track))
            .order(ranks::required_level.asc())
            .limit(1)
            .get_result(&mut conn.get()?)?
        )
    }

    pub fn get_rank_by_name(conn: &Db, needle: &String, track: &i64) -> ManifoldResult<Rank> {
        Ok(ranks::table
            .filter(ranks::rank_name.like(needle.trim().to_lowercase()))
            .filter(ranks::rank_track.eq(track))
            .order(ranks::required_level.desc())
            .limit(1)
            .get_result(&mut conn.get()?)?
        )
    }
}

impl Track {
    pub fn get_track_by_name(conn: &Db, needle: &String) -> ManifoldResult<Track> {
        Ok(tracks::table
            .filter(tracks::track_name.like(needle.trim().to_lowercase()))
            .select(Track::as_select())
            .limit(1)
            .get_result(&mut conn.get()?)?
        )
    }

    pub fn get_track_by_id(conn: &Db, needle: &i64) -> ManifoldResult<Track> {
        Ok(tracks::table
            .filter(tracks::track_id.eq(needle))
            .select(Track::as_select())
            .limit(1)
            .get_result(&mut conn.get()?)?
        )
    }
}