use fred::prelude::{FredResult, LuaInterface};
use fred::types::Key;
use fred::types::scripts::Script;
use std::sync::LazyLock;
use std::time::Duration;

static RATE_LIMIT_SCRIPT: LazyLock<Script> = LazyLock::new(|| {
    Script::from_lua(include_str!(concat!(
        env!("CARGO_MANIFEST_DIR"),
        "/scripts/rate_limit.lua"
    )))
});

pub(crate) async fn run_script<L: LuaInterface>(
    conn: &L,
    key: Key,
    capacity: u16,
    emission_interval_ms: u32,
) -> FredResult<(bool, u32)> {
    match RATE_LIMIT_SCRIPT
        .evalsha(conn, &key, (capacity, emission_interval_ms))
        .await
    {
        Err(e) if e.details().starts_with("NOSCRIPT") => {
            let _: () = conn
                .script_load(
                    RATE_LIMIT_SCRIPT
                        .lua()
                        .expect("RATE_LIMIT_SCRIPT was initialized incorrectly")
                        .clone(),
                )
                .await?;
            RATE_LIMIT_SCRIPT
                .evalsha(conn, key, (capacity, emission_interval_ms))
                .await
        }
        result => result,
    }
}

/// Manually applies rate-limiting to `key`.
///
/// Useful when the key is derived from parsed request data, such as a body field,
/// where a Tower middleware layer would not yet have access to it.
///
/// - `Ok(None)` --- allowed
/// - `Ok(Some(retry_after))` --- rate-limited; the next slot opens in `retry_after` seconds
pub async fn rate_limit<L>(
    conn: &L,
    key: impl Into<Key>,
    capacity: u16,
    emission_interval: Duration,
) -> FredResult<Option<u32>>
where
    L: LuaInterface,
{
    let emission_interval_ms = emission_interval
        .as_millis()
        .try_into()
        .expect("emission_interval exceeds u32::MAX milliseconds");
    let (allowed, retry_after) =
        run_script(conn, key.into(), capacity, emission_interval_ms).await?;
    Ok((!allowed).then_some(retry_after))
}
