glimit

//// This module provides a rate limiter that can be used to limit the number of
//// requests or function calls per second for a given identifier.
////
//// A single rate limiter actor stores all token bucket state. Each hit is a single
//// message to the rate limiter, which performs the Token Bucket calculation inline.
//// A periodic sweep removes full or idle buckets to reduce memory usage. The
//// idle threshold defaults to 60 seconds and can be configured via `max_idle`.
//// The rate limiter fails open — if the rate limiter actor is unavailable,
//// requests are allowed through.
////
//// The rate limits are configured using the following two options:
////
//// - `per_second`: The rate of new available tokens per second. Think of this
////   as the steady state rate limit.
//// - `burst_limit`: The maximum number of available tokens. Think of this as
////   the burst rate limit. The default value is the `per_second` rate limit.
////
//// The rate limiter can be applied to a function or handler using the `apply`
//// function, which returns a new function that checks the rate limit before
//// calling the original function.
////
//// # Example
////
//// ```gleam
//// import glimit
////
//// let limiter =
////   glimit.new()
////   |> glimit.per_second(10)
////   |> glimit.burst_limit(100)
////   |> glimit.identifier(fn(request) { request.ip })
////   |> glimit.on_limit_exceeded(fn(_request) { "Rate limit reached" })
////
//// let handler =
////   fn(_request) { "Hello, world!" }
////   |> glimit.apply(limiter)
//// ```
////
//// # Multi-argument functions
////
//// `apply` wraps a single-argument function `fn(a) -> b`. To rate-limit a
//// function with multiple arguments, use `apply2`, `apply3`, or `apply4`:
////
//// ```gleam
//// let limiter =
////   glimit.new()
////   |> glimit.per_second(10)
////   |> glimit.identifier(fn(args: #(String, String)) { args.0 })
////   |> glimit.on_limit_exceeded(fn(_args) { too_many_requests() })
////
//// let limited_handle =
////   handle
////   |> glimit.apply2(limiter)
////
//// limited_handle("user_123", "upload")
//// ```
////
//// # Pluggable store backend
////
//// By default, rate limit state is stored in-memory using an OTP actor. For
//// distributed rate limiting (e.g. across multiple nodes), you can provide a
//// custom `Store` that persists bucket state externally (Redis, Postgres, etc.).
////
//// All token bucket logic stays in glimit — adapters only implement simple
//// get/set/lock/unlock operations. The `glimit/bucket` module is public and
//// provides `to_pairs`/`from_pairs` helpers for serialization.
////
//// ```gleam
//// import glimit
//// import glimit/bucket
////
//// // Redis adapter example (using radish):
//// let store = glimit.Store(
////   get: fn(key) {
////     case radish.execute(client, ["HGETALL", key], 1000) {
////       Ok(fields) -> Ok(bucket.from_pairs(parse_hgetall_response(fields)))
////       Error(_) -> Error(Nil)
////     }
////   },
////   set: fn(key, state, ttl) {
////     let pairs = bucket.to_pairs(state) |> list.flat_map(fn(p) { [p.0, p.1] })
////     let _ = radish.execute(client, ["HSET", key, ..pairs], 1000)
////     let _ = radish.execute(client, ["EXPIRE", key, int.to_string(ttl)], 1000)
////     Ok(Nil)
////   },
////   lock: fn(key) {
////     case radish.execute(client, ["SET", key <> ":lock", "1", "NX", "EX", "5"], 1000) {
////       Ok(_) -> Ok(Nil)
////       Error(_) -> Error(Nil)
////     }
////   },
////   unlock: fn(key) {
////     let _ = radish.execute(client, ["DEL", key <> ":lock"], 1000)
////     Ok(Nil)
////   },
//// )
////
//// glimit.new()
//// |> glimit.per_second(10)
//// |> glimit.store(store)
//// |> glimit.identifier(fn(req) { req.ip })
//// |> glimit.on_limit_exceeded(fn(_) { "Rate limited" })
//// |> glimit.apply(handler)
//// ```
////

import gleam/option.{type Option, None, Some}
import gleam/result
import gleam/string
import glimit/bucket
import glimit/memory_store.{type MemoryStore}
import glimit/rate_limiter

/// A pluggable storage backend for distributed rate limiting.
/// See `glimit/bucket.Store` for full documentation.
///
pub type Store =
  bucket.Store

/// Error type returned when a rate limit check fails.
///
pub type HitError =
  rate_limiter.HitError

/// A rate limiter.
///
pub type RateLimiter(a, b, id) {
  RateLimiter(
    rate_limiter_actor: rate_limiter.RateLimiterActor(id),
    on_limit_exceeded: fn(a) -> b,
    identifier: fn(a) -> id,
    memory_store: Option(MemoryStore),
  )
}

/// A builder for configuring the rate limiter.
///
pub type RateLimiterBuilder(a, b, id) {
  RateLimiterBuilder(
    per_second: Option(fn(id) -> Int),
    burst_limit: Option(fn(id) -> Int),
    identifier: Option(fn(a) -> id),
    on_limit_exceeded: Option(fn(a) -> b),
    max_idle_ms: Option(Int),
    store: Option(Store),
  )
}

/// Create a new rate limiter builder.
///
pub fn new() -> RateLimiterBuilder(a, b, id) {
  RateLimiterBuilder(
    per_second: None,
    burst_limit: None,
    identifier: None,
    on_limit_exceeded: None,
    max_idle_ms: Some(60_000),
    store: None,
  )
}

/// Set the rate of new available tokens per second.
///
/// Note that this is not the maximum number of requests that can be made in a single
/// second, but the rate at which tokens are added to the bucket. Think of this as the
/// steady state rate limit, while the `burst_limit` function sets the maximum number of
/// available tokens (or the burst rate limit).
///
/// This value is also used as the default value for the `burst_limit` function.
///
/// # Example
///
/// ```gleam
/// import glimit
///
/// let limiter =
///   glimit.new()
///   |> glimit.per_second(10)
/// ```
///
pub fn per_second(
  limiter: RateLimiterBuilder(a, b, id),
  limit: Int,
) -> RateLimiterBuilder(a, b, id) {
  RateLimiterBuilder(..limiter, per_second: Some(fn(_) { limit }))
}

/// Set the rate limit per second, based on the identifier.
///
/// Note: this function is evaluated once when a bucket is first created for an
/// identifier. If the function returns a different value later, existing buckets
/// are not affected until they are swept (due to idleness or being full) and
/// re-created on the next hit.
///
/// # Example
///
/// ```gleam
/// import glimit
///
/// let limiter =
///   glimit.new()
///   |> glimit.identifier(fn(request) { request.user_id })
///   |> glimit.per_second_fn(fn(user_id) {
///     db.get_rate_limit(user_id)
///   })
/// ```
///
pub fn per_second_fn(
  limiter: RateLimiterBuilder(a, b, id),
  limit_fn: fn(id) -> Int,
) -> RateLimiterBuilder(a, b, id) {
  RateLimiterBuilder(..limiter, per_second: Some(limit_fn))
}

/// Set the maximum number of available tokens.
///
/// The maximum number of available tokens is the maximum number of requests that can be
/// made in a single burst when the bucket is full. The default value is the same as the
/// rate limit per second.
///
/// # Example
///
/// ```gleam
/// import glimit
///
/// let limiter =
///   glimit.new()
///   |> glimit.per_second(10)
///   |> glimit.burst_limit(100)
/// ```
///
pub fn burst_limit(
  limiter: RateLimiterBuilder(a, b, id),
  burst_limit: Int,
) -> RateLimiterBuilder(a, b, id) {
  RateLimiterBuilder(..limiter, burst_limit: Some(fn(_) { burst_limit }))
}

/// Set the maximum number of available tokens, based on the identifier.
///
/// Note: this function is evaluated once when a bucket is first created for an
/// identifier. If the function returns a different value later, existing buckets
/// are not affected until they are swept (due to idleness or being full) and
/// re-created on the next hit.
///
/// # Example
///
/// ```gleam
/// import glimit
///
/// let limiter =
///   glimit.new()
///   |> glimit.identifier(fn(request) { request.user_id })
///   |> glimit.per_second(10)
///   |> glimit.burst_limit_fn(fn(user_id) {
///     db.get_burst_limit(user_id)
///   })
/// ```
///
pub fn burst_limit_fn(
  limiter: RateLimiterBuilder(a, b, id),
  burst_limit_fn: fn(id) -> Int,
) -> RateLimiterBuilder(a, b, id) {
  RateLimiterBuilder(..limiter, burst_limit: Some(burst_limit_fn))
}

/// Set the idle eviction threshold in seconds.
///
/// Buckets that have not been hit for longer than this duration are removed
/// during periodic sweeps. The default is 60 seconds. Set to `0` to disable
/// idle eviction entirely.
///
/// For rate limiters with a high `burst_limit` relative to `per_second`, you
/// may want to increase this value so that partially-refilled buckets are not
/// evicted prematurely. A good rule of thumb is
/// `burst_limit / per_second` seconds.
///
/// # Example
///
/// ```gleam
/// import glimit
///
/// let limiter =
///   glimit.new()
///   |> glimit.per_second(1)
///   |> glimit.burst_limit(1000)
///   |> glimit.max_idle(1000)
/// ```
///
pub fn max_idle(
  limiter: RateLimiterBuilder(a, b, id),
  seconds: Int,
) -> RateLimiterBuilder(a, b, id) {
  case seconds {
    s if s <= 0 -> RateLimiterBuilder(..limiter, max_idle_ms: None)
    s -> RateLimiterBuilder(..limiter, max_idle_ms: Some(s * 1000))
  }
}

/// Set a pluggable store backend for distributed rate limiting.
///
/// When a store is configured, bucket state is read from and written to the
/// store on each hit instead of being kept in the actor's in-memory dictionary.
/// The periodic sweep becomes a no-op since external stores handle expiry via TTL.
///
/// # Example
///
/// ```gleam
/// import glimit
///
/// let limiter =
///   glimit.new()
///   |> glimit.per_second(10)
///   |> glimit.store(my_redis_store)
///   |> glimit.identifier(fn(request) { request.ip })
///   |> glimit.on_limit_exceeded(fn(_request) { "Rate limit reached" })
/// ```
///
pub fn store(
  limiter: RateLimiterBuilder(a, b, id),
  store: Store,
) -> RateLimiterBuilder(a, b, id) {
  RateLimiterBuilder(..limiter, store: Some(store))
}

/// Set the handler to be called when the rate limit is reached.
///
/// # Example
///
/// ```gleam
/// import glimit
///
/// let limiter =
///   glimit.new()
///   |> glimit.per_second(10)
///   |> glimit.on_limit_exceeded(fn(_request) { "Rate limit reached" })
/// ```
///
pub fn on_limit_exceeded(
  limiter: RateLimiterBuilder(a, b, id),
  on_limit_exceeded: fn(a) -> b,
) -> RateLimiterBuilder(a, b, id) {
  RateLimiterBuilder(..limiter, on_limit_exceeded: Some(on_limit_exceeded))
}

/// Set the identifier function to be used to identify the rate limit.
///
/// # Example
///
/// ```gleam
/// import glimit
///
/// let limiter =
///   glimit.new()
///   |> glimit.identifier(fn(request) { request.ip })
/// ```
///
pub fn identifier(
  limiter: RateLimiterBuilder(a, b, id),
  identifier: fn(a) -> id,
) -> RateLimiterBuilder(a, b, id) {
  RateLimiterBuilder(..limiter, identifier: Some(identifier))
}

/// Build the rate limiter.
///
/// Note that using `apply` will already build the rate limiter, so this function is
/// only useful if you want to build the rate limiter manually and apply it to multiple
/// functions.
///
/// To apply the resulting rate limiter to a function or handler, use the `apply_built`
/// function.
///
pub fn build(
  config: RateLimiterBuilder(a, b, id),
) -> Result(RateLimiter(a, b, id), String) {
  use per_second <- result.try(case config.per_second {
    Some(per_second) -> Ok(per_second)
    None -> Error("`per_second` rate limit is required")
  })
  let burst_limit = case config.burst_limit {
    Some(burst_limit) -> burst_limit
    None -> per_second
  }
  use identifier <- result.try(case config.identifier {
    Some(identifier) -> Ok(identifier)
    None -> Error("`identifier` function is required")
  })
  use on_limit_exceeded <- result.try(case config.on_limit_exceeded {
    Some(on_limit_exceeded) -> Ok(on_limit_exceeded)
    None -> Error("`on_limit_exceeded` function is required")
  })

  // Resolve the store: use provided store or create an in-memory store
  use #(resolved_store, mem_store) <- result.try(case config.store {
    Some(s) -> Ok(#(s, None))
    None -> {
      case memory_store.new(config.max_idle_ms, 10_000) {
        Ok(#(s, handle)) -> Ok(#(s, Some(handle)))
        Error(_) -> Error("Failed to start memory store")
      }
    }
  })

  use rate_limiter_actor <- result.try(
    rate_limiter.new(per_second, burst_limit, resolved_store)
    |> result.map_error(fn(_) { "Failed to start rate limiter" }),
  )

  Ok(RateLimiter(
    rate_limiter_actor: rate_limiter_actor,
    on_limit_exceeded: on_limit_exceeded,
    identifier: identifier,
    memory_store: mem_store,
  ))
}

/// Apply the rate limiter to a request handler or function.
///
/// Panics if the rate limiter cannot be started or if the `identifier`
/// function or `on_limit_exceeded` function is missing.
///
pub fn apply(
  func: fn(a) -> b,
  config: RateLimiterBuilder(a, b, id),
) -> fn(a) -> b {
  let limiter = case build(config) {
    Ok(limiter) -> limiter
    Error(message) -> panic as message
  }
  apply_built(func, limiter)
}

/// Apply the rate limiter to a request handler or function.
///
/// This function is useful if you want to build the rate limiter manually using the
/// `build` function.
///
pub fn apply_built(
  func: fn(a) -> b,
  limiter: RateLimiter(a, b, id),
) -> fn(a) -> b {
  fn(input: a) -> b {
    let identifier = limiter.identifier(input)
    case rate_limiter.hit(limiter.rate_limiter_actor, identifier) {
      Ok(Nil) -> func(input)
      Error(rate_limiter.RateLimited) -> limiter.on_limit_exceeded(input)
      Error(rate_limiter.Unavailable) | Error(rate_limiter.StoreLockFailed) ->
        func(input)
    }
  }
}

/// Return the number of tracked identifiers in the in-memory store.
///
/// Returns 0 if the rate limiter uses an external store.
///
pub fn get_count(limiter: RateLimiter(a, b, id)) -> Int {
  case limiter.memory_store {
    Some(ms) -> memory_store.get_count(ms)
    None -> 0
  }
}

/// Remove an identifier from the in-memory store.
///
/// No-op if the rate limiter uses an external store.
///
pub fn remove(limiter: RateLimiter(a, b, id), identifier: id) -> Nil {
  case limiter.memory_store {
    Some(ms) -> {
      let key = "glimit:" <> string.inspect(identifier)
      let _ = memory_store.remove(ms, key)
      Nil
    }
    None -> Nil
  }
}

/// Remove full or idle buckets from the in-memory store synchronously.
///
/// No-op if the rate limiter uses an external store.
///
pub fn sweep(limiter: RateLimiter(a, b, id)) -> Nil {
  case limiter.memory_store {
    Some(_ms) -> Nil
    None -> Nil
  }
}

/// Apply the rate limiter to a 2-argument function.
///
/// The config's `identifier` and `on_limit_exceeded` receive a `#(a, b)` tuple.
///
/// # Example
///
/// ```gleam
/// let limiter =
///   glimit.new()
///   |> glimit.per_second(10)
///   |> glimit.identifier(fn(args: #(String, String)) { args.0 })
///   |> glimit.on_limit_exceeded(fn(_) { "Rate limited" })
///
/// let limited =
///   handle
///   |> glimit.apply2(limiter)
///
/// limited("user_123", "upload")
/// ```
///
pub fn apply2(
  func: fn(a, b) -> c,
  config: RateLimiterBuilder(#(a, b), c, id),
) -> fn(a, b) -> c {
  let wrapped =
    fn(args: #(a, b)) -> c { func(args.0, args.1) }
    |> apply(config)
  fn(a: a, b: b) -> c { wrapped(#(a, b)) }
}

/// Apply the rate limiter to a 3-argument function.
///
/// The config's `identifier` and `on_limit_exceeded` receive a `#(a, b, c)` tuple.
///
/// # Example
///
/// ```gleam
/// let limiter =
///   glimit.new()
///   |> glimit.per_second(10)
///   |> glimit.identifier(fn(args: #(String, String, Int)) { args.0 })
///   |> glimit.on_limit_exceeded(fn(_) { "Rate limited" })
///
/// let limited =
///   handle
///   |> glimit.apply3(limiter)
///
/// limited("user_123", "upload", 42)
/// ```
///
pub fn apply3(
  func: fn(a, b, c) -> d,
  config: RateLimiterBuilder(#(a, b, c), d, id),
) -> fn(a, b, c) -> d {
  let wrapped =
    fn(args: #(a, b, c)) -> d { func(args.0, args.1, args.2) }
    |> apply(config)
  fn(a: a, b: b, c: c) -> d { wrapped(#(a, b, c)) }
}

/// Apply the rate limiter to a 4-argument function.
///
/// The config's `identifier` and `on_limit_exceeded` receive a `#(a, b, c, d)` tuple.
///
/// # Example
///
/// ```gleam
/// let limiter =
///   glimit.new()
///   |> glimit.per_second(10)
///   |> glimit.identifier(fn(args: #(String, String, Int, Bool)) { args.0 })
///   |> glimit.on_limit_exceeded(fn(_) { "Rate limited" })
///
/// let limited =
///   handle
///   |> glimit.apply4(limiter)
///
/// limited("user_123", "upload", 42, True)
/// ```
///
pub fn apply4(
  func: fn(a, b, c, d) -> e,
  config: RateLimiterBuilder(#(a, b, c, d), e, id),
) -> fn(a, b, c, d) -> e {
  let wrapped =
    fn(args: #(a, b, c, d)) -> e { func(args.0, args.1, args.2, args.3) }
    |> apply(config)
  fn(a: a, b: b, c: c, d: d) -> e { wrapped(#(a, b, c, d)) }
}