crabbucket_redis

import gleam/erlang/process.{type Subject}
import gleam/int
import gleam/io
import gleam/result
import radish.{type Message}
import radish/error.{type Error}
import radish/hash
import radish/resp
import radish/utils

@external(erlang, "os", "system_time")
fn system_time(second_division: Int) -> Int

/// `window_duration_milliseconds` describes the length of the time window
/// valid for a number of tokens.
/// For instance, a value of 1,000 with a default token count of 100 would mean
/// an action could occur 100 times per second.
pub type RedisTokenBucket {
  RedisTokenBucket(
    redis_connection: Subject(Message),
    window_duration_milliseconds: Int,
  )
}

pub type RemainingTokenCountSuccess {
  HasRemainingTokens(remaining_tokens: Int, next_reset_timestamp: Int)
}

pub type RemainingTokenCountFailure {
  MustWaitUntil(next_reset_timestamp: Int)
  RedisError(error: Error)
}

/// Logic:
/// - Try to get the current value for this token bucket key
/// - If there is none, the action should be allowed,
///   with the current time and remaining tokens being cached
/// - If there was a value, check if its time window has passed,
///   in which case we can 'reset' it to a fresh value, like above
/// - If no tokens remain in the current entry, return -1
/// - Else, decrement remaining tokens, saving and returned that value
/// - All of the above return values are part of a tuple:
///   {remaining_tokens, timestamp_of_next_time_window}
/// - An expiration time is set on the key, but this is more for garbage collection,
///   and is not relied on
const script = "local key = KEYS[1]
local window_start_arg = tonumber(ARGV[1])
local window_duration_arg = tonumber(ARGV[2])
local tokens_arg = tonumber(ARGV[3])
local expiration_seconds = math.floor(window_duration_arg / 1000) + 1
local window_start = tonumber(redis.call('HGET', key, 'window_start'))
local tokens = tonumber(redis.call('HGET', key, 'tokens'))
if window_start == nil or tokens == nil then
    redis.call('HSET', key, 'window_start', window_start_arg, 'tokens', tokens_arg)
    redis.call('EXPIRE', key, expiration_seconds)
    return {tokens_arg, window_start_arg + window_duration_arg}
end
local time_arr = redis.call('TIME')
local time_ms = (time_arr[1] * 1000) + (math.floor(time_arr[2] / 1000))
if (window_start + window_duration_arg) < time_ms then
    redis.call('HSET', key, 'window_start', window_start_arg, 'tokens', tokens_arg)
    redis.call('EXPIRE', key, expiration_seconds)
    return {tokens_arg, window_start_arg + window_duration_arg}
end
if tokens <= 0 then
    return {-1, window_start + window_duration_arg}
end
redis.call('HSET', key, 'tokens', tokens - 1)
return {tokens - 1, window_start + window_duration_arg}"

/// Takes an arbitrary string key, inserting a record if non-existing.
/// 
/// Suggestion: you may wish to format you string key such that it indicates
/// usage, purpose, value type, and value.
/// For instance, if you're limiting a specific endpoint for a given user,
/// the key might look something like:
/// `limit:some_endpoint_name:user:12345`
/// 
/// Return should be either `HasRemainingTokens(remaining_tokens: Int)`,
/// which indicates that an action may proceed and contains how many more times
/// the action may occur with the current window,
/// or it may be `MustWaitFor(milliseconds: Int)`,
/// which indicates that no tokens remain for this keep and how many
/// milliseconds remain until the end of the current window.
pub fn remaining_tokens_for_key(
  bucket: RedisTokenBucket,
  key: String,
  default_token_count: Int,
) -> Result(RemainingTokenCountSuccess, RemainingTokenCountFailure) {
  let cmd =
    utils.prepare([
      "EVAL",
      script,
      "1",
      key,
      system_time(1000) |> int.to_string(),
      bucket.window_duration_milliseconds |> int.to_string(),
      default_token_count - 1 |> int.to_string(),
    ])

  use results <- result.try(
    utils.execute(bucket.redis_connection, cmd, 500)
    |> result.map_error(fn(e) { RedisError(e) }),
  )

  // Implementation detail:
  // Lua script executed by Redis returns _how many tokens are remaining_,
  // so we can't check for tokens_remaining <= 0, which would cause issues
  // in the case of having just used the last token in a window.
  // Therefore, a negative remaining token value is used to indicate the case
  // of the cache entry not having remaining tokens.
  case results {
    [resp.Array([resp.Integer(tokens_remaining), resp.Integer(next_reset)])] -> {
      case tokens_remaining < 0 {
        True -> {
          Error(MustWaitUntil(next_reset))
        }
        False -> Ok(HasRemainingTokens(tokens_remaining, next_reset))
      }
    }
    other -> {
      io.debug(other)
      panic as "Should be unreachable"
    }
  }
}

/// Returns True if record was deleted, False if record didn't exist
pub fn clear_key(bucket: RedisTokenBucket, key: String) -> Result(Bool, Error) {
  use result <- result.try({
    hash.del(bucket.redis_connection, key, ["window_start", "tokens"], 500)
  })

  Ok(result > 0)
}