// Fault tolerant rate limiting backed by Postgres.

use anyhow::Result;
use chrono::{DateTime, TimeDelta, Utc};
use diesel::{
    dsl::{auto_type, insert_into, AsSelect},
    pg::Pg,
    prelude::*,
    sql_types::Timestamptz,
};
use uuid::Uuid;

use crate::schema::{governor_entries, governors};

define_sql_function! {
    fn greatest(a: diesel::sql_types::Integer, b: diesel::sql_types::Integer) -> Integer
}

#[derive(Clone, Debug, Identifiable, Insertable, Queryable, Selectable)]
#[diesel(table_name = governors)]
pub struct Governor {
    pub id: Uuid,
    pub team_id: Uuid,
    pub project_id: Option<Uuid>,
    pub window_size: TimeDelta,
    pub max_count: i32,
    pub rolling_count: i32,
}

impl Governor {
    pub fn insert_new<'a>(
        db_conn: &mut diesel::PgConnection,
        team_id: &'a Uuid,
        project_id: Option<&'a Uuid>,
        window_size: &'a TimeDelta,
        max_count: i32,
    ) -> Result<Self> {
        let id: Uuid = Uuid::now_v7();
        Ok(insert_into(governors::table)
            .values((
                governors::team_id.eq(team_id),
                governors::id.eq(id),
                governors::project_id.eq(project_id),
                governors::window_size.eq(window_size),
                governors::max_count.eq(max_count),
            ))
            .get_result(db_conn)?)
    }

    #[auto_type(no_type_alias)]
    pub fn all() -> _ {
        let select: AsSelect<Governor, Pg> = Governor::as_select();
        governors::table.select(select)
    }

    #[auto_type(no_type_alias)]
    pub fn with_id(governor_id: Uuid) -> _ {
        governors::id.eq(governor_id)
    }

    #[auto_type(no_type_alias)]
    pub fn with_team(team_id: Uuid) -> _ {
        governors::team_id.eq(team_id)
    }

    #[auto_type(no_type_alias)]
    pub fn with_project(project_id: Option<Uuid>) -> _ {
        governors::project_id.is_not_distinct_from(project_id)
    }

    // TODO: return a custom result enum instead of a Result<Option>, for
    // better readability
    /**
     * Attempt to increment the rolling count. If the governor is not full,
     * returns a GovernorEntry which can be used to cancel the operation and
     * restore the rolling count. If governor is full, returns None.
     */
    pub fn create_entry(&self, conn: &mut diesel::PgConnection) -> Result<Option<GovernorEntry>> {
        let entry = diesel::insert_into(governor_entries::table)
            .values((
                governor_entries::id.eq(Uuid::now_v7()),
                governor_entries::governor_id.eq(self.id),
            ))
            .get_result(conn)?;
        let n_rows = diesel::update(
            governors::table
                .filter(governors::id.eq(self.id))
                .filter(governors::rolling_count.lt(self.max_count)),
        )
        .set(governors::rolling_count.eq(governors::rolling_count + 1))
        .execute(conn)?;
        assert!(n_rows < 2);
        if n_rows == 1 {
            Ok(Some(entry))
        } else {
            // Clean up unused entry, or else it will artificially decrement
            // rolling count when it expires
            diesel::delete(governor_entries::table.filter(GovernorEntry::with_id(entry.id)))
                .execute(conn)?;
            Ok(None)
        }
    }

    /**
     * Governors work by continually incrementing a counter and then
     * periodically decrementing it as entries fall out of the current window of
     * time. This function performs the latter part of the cycle, sweeping out
     * expired entries and adjusting the counter accordingly.
     */
    pub fn reclaim(&self, conn: &mut diesel::PgConnection) -> Result<()> {
        let n_expired_entries: i32 = diesel::delete(
            GovernorEntry::belonging_to(self).filter(
                governor_entries::timestamp
                    .lt(diesel::dsl::now.into_sql::<Timestamptz>() - self.window_size),
            ),
        )
        .execute(conn)?
        .try_into()
        .expect("a governor should never have been allowed enough entries to overflow an i32");
        // Clamp rolling_count >= 0
        diesel::update(governors::table.filter(Self::with_id(self.id.clone())))
            .set(
                governors::rolling_count
                    .eq(greatest(governors::rolling_count - n_expired_entries, 0)),
            )
            .execute(conn)?;
        Ok(())
    }

    pub fn reclaim_all(conn: &mut diesel::PgConnection) -> Result<()> {
        let applicable_governors = governors::table
            .inner_join(governor_entries::table)
            .filter(
                governor_entries::timestamp
                    .lt(diesel::dsl::now.into_sql::<Timestamptz>() - governors::window_size),
            )
            .select(Self::as_select())
            .group_by(governors::id)
            .load(conn)?;
        tracing::info!(
            "reclaiming counts for {} governors",
            applicable_governors.len()
        );
        for governor in applicable_governors {
            governor.reclaim(conn)?;
        }
        Ok(())
    }

    /**
     * Reset all governors to a count of 0, to fix any accumulated error between
     * rolling counts and number of entries.
     */
    pub fn reset_all(conn: &mut diesel::PgConnection) -> Result<()> {
        // Delete entries and then reset counts, not vice-versa; otherwise
        // concurrent inserts could result in rolling counts getting stuck
        // higher than they should be
        diesel::delete(governor_entries::table).execute(conn)?;
        diesel::update(governors::table)
            .set(governors::rolling_count.eq(0))
            .execute(conn)?;
        Ok(())
    }
}

#[derive(Associations, Clone, Debug, Identifiable, Insertable, Queryable, Selectable)]
#[diesel(table_name = governor_entries)]
#[diesel(belongs_to(Governor))]
pub struct GovernorEntry {
    pub id: Uuid,
    pub governor_id: Uuid,
    pub timestamp: DateTime<Utc>,
}

impl GovernorEntry {
    #[auto_type(no_type_alias)]
    pub fn with_id(entry_id: Uuid) -> _ {
        governor_entries::id.eq(entry_id)
    }

    /**
     * Removes this entry from the governor and decrements the overall rolling
     * count by 1.
     */
    pub fn cancel(&self, conn: &mut diesel::PgConnection) -> Result<()> {
        let entry_filter = Self::with_id(self.id.clone());
        let governor_filter = Governor::with_id(self.governor_id.clone());
        diesel::update(governors::table.filter(governor_filter))
            .set(governors::rolling_count.eq(greatest(governors::rolling_count - 1, 0)))
            .execute(conn)?;
        diesel::delete(governor_entries::table.filter(entry_filter)).execute(conn)?;
        Ok(())
    }
}
