// Copyright 2026 Leo Cheng
// SPDX-License-Identifier: Apache-2.0

///|
/// A token bucket: `rate` tokens accrue per second up to `burst`, and each request
/// takes one.
///
/// It is the rate limiter that allows a burst after a quiet period, which is what
/// a limit on a human-facing API wants — nobody should be refused for making two
/// requests in the same second after an hour of making none.
///
/// Time is nanoseconds from the caller's clock, as everywhere in this package.
pub struct Bucket {
  rate : Double
  burst : Double
  mut state : Tokens
}

///|
/// What a bucket holds at one instant: the tokens, and when they were last
/// brought up to date.
///
/// This is the bucket apart from its policy and its clock, which is what two
/// processes sharing one allowance have to agree on. See [`Bucket::step`].
#derive(Eq, Debug)
pub(all) struct Tokens {
  /// How many tokens are in the bucket.
  tokens : Double
  /// The clock reading `tokens` was brought up to date at.
  at : Int64
}

///|
pub extend Tokens with Debug::{to_repr}

///|
pub extend Tokens with Eq::{not_equal, equal}

///|
/// What one call gets: whether it may go ahead, what the bucket holds afterwards,
/// and how long a refused call has to wait.
#derive(Eq, Debug)
pub(all) struct Grant {
  /// Whether the call may go ahead.
  granted : Bool
  /// The state to keep, whether the call went through or not.
  state : Tokens
  /// Zero when granted; otherwise how long until the same call would be granted.
  wait : @moondate.Span
}

///|
pub extend Grant with Debug::{to_repr}

///|
pub extend Grant with Eq::{not_equal, equal}

///|
/// A bucket that starts full, which is what lets the first burst through.
///
/// `rate` is tokens per second and `burst` is the most it will ever hold. A burst
/// smaller than one means nothing ever passes, so it is raised to one.
///
/// # Parameters
/// - `rate` : tokens accrued per second
/// - `burst` : the most the bucket will ever hold; with none given, the rate, and never under one
/// - `now` : the clock reading the bucket starts at, `0L` by default
pub fn Bucket::new(
  rate : Double,
  burst~ : Double = 0.0,
  now~ : Int64 = 0L,
) -> Bucket {
  let cap = if burst > 0.0 { burst } else if rate > 1.0 { rate } else { 1.0 }
  { rate, burst: cap, state: { tokens: cap, at: now, }, }
}

///|
/// The state a bucket that has seen no traffic is in at `now`: full, so the first
/// burst goes through.
///
/// This is what a shared store writes when the key is missing, which is the same
/// state [`Bucket::new`] starts from.
///
/// # Parameters
/// - `now` : the caller's clock reading, in nanoseconds
pub fn Bucket::opening(self : Bucket, now : Int64) -> Tokens {
  { tokens: self.burst, at: now, }
}

///|
/// Whether `n` tokens are available now, taking them if they are.
///
/// # Parameters
/// - `now` : the caller's clock reading, in nanoseconds
/// - `n` : how many tokens the call costs, one by default; a negative cost is taken as zero
pub fn Bucket::take(self : Bucket, now : Int64, n~ : Double = 1.0) -> Bool {
  let grant = self.step(self.state, now, n~)
  self.state = grant.state
  grant.granted
}

///|
/// How long until `n` tokens would be available, without taking them.
///
/// Zero when they are available now. This is what a caller sleeps for rather than
/// spinning, and what a `Retry-After` header is computed from.
///
/// # Parameters
/// - `now` : the caller's clock reading, in nanoseconds
/// - `n` : how many tokens are wanted, one by default
pub fn Bucket::next(
  self : Bucket,
  now : Int64,
  n~ : Double = 1.0,
) -> @moondate.Span {
  self.step(self.state, now, n~).wait
}

///|
/// How many tokens are in the bucket, for a caller that reports its own headroom.
///
/// # Parameters
/// - `now` : the caller's clock reading, in nanoseconds
pub fn Bucket::tokens(self : Bucket, now : Int64) -> Double {
  self.step(self.state, now, n=0.0).state.tokens
}

///|
/// Answer one call against `state` at `now`, leaving the bucket untouched.
///
/// This is the whole of the arithmetic, and the seam for a bucket shared by
/// several processes: read the state out of the store — [`Bucket::opening`] when
/// the key is missing — ask here, and write `grant.state` back only if the stored
/// value has not changed in between; on a clash, read again and ask again. The
/// compare-and-set is what stops two writers from spending the same token, since
/// an answer is only true of the state it was computed from. Token counts are
/// fractional, so a store that rounds them to an integer loses the point of the
/// bucket.
///
/// # Parameters
/// - `state` : the bucket's state as the store has it
/// - `now` : the caller's clock reading, in nanoseconds
/// - `n` : how many tokens the call costs, one by default; a negative cost is taken as zero
pub fn Bucket::step(
  self : Bucket,
  state : Tokens,
  now : Int64,
  n~ : Double = 1.0,
) -> Grant {
  let n = if n < 0.0 { 0.0 } else { n }
  let mut tokens = state.tokens
  let mut at = state.at
  if now > at {
    let elapsed = (now - at).to_double() / 1.0e9
    tokens = tokens + elapsed * self.rate
    if tokens > self.burst {
      tokens = self.burst
    }
    at = now
  }
  let granted = tokens >= n
  if granted {
    tokens = tokens - n
  }
  let wait = if granted {
    @moondate.Span::new()
  } else if self.rate <= 0.0 {
    // Nothing accrues, so the wait is unbounded; a caller is better told a very
    // long time than a zero it would spin on.
    @moondate.Span::new(days=365L)
  } else {
    { nanos: ((n - tokens) / self.rate * 1.0e9).to_int64(), }
  }
  { granted, state: { tokens, at, }, wait, }
}

///|
/// A sliding window counter: at most `quota` events in any `window` long stretch.
///
/// Where a bucket smooths, a window counts. A quota of "a thousand a day" is a
/// window, not a rate, and answering it with a bucket would let a client spend the
/// whole day's allowance in the first minute.
///
/// The stamps of the events in the window are kept, so the count is exact rather
/// than the approximation a two-bucket scheme gives.
pub struct Sliding {
  quota : Int32
  window : @moondate.Span
  stamps : Array[Int64]
}

///|
/// An empty window.
///
/// # Parameters
/// - `quota` : events allowed in any `window`-long stretch
/// - `window` : the length of that stretch
pub fn Sliding::new(quota : Int32, window : @moondate.Span) -> Sliding {
  { quota, window, stamps: [], }
}

///|
/// Whether an event is within the quota now, counting it if it is.
///
/// # Parameters
/// - `now` : the caller's clock reading, in nanoseconds
pub fn Sliding::take(self : Sliding, now : Int64) -> Bool {
  self.forget(now)
  if self.stamps.length() < self.quota {
    self.stamps.push(now)
    return true
  }
  false
}

///|
/// How many events remain in the current window.
///
/// # Parameters
/// - `now` : the caller's clock reading, in nanoseconds
pub fn Sliding::left(self : Sliding, now : Int64) -> Int32 {
  self.forget(now)
  let used = self.stamps.length()
  if used >= self.quota {
    0
  } else {
    self.quota - used
  }
}

///|
/// When the window will next have room, which is when its oldest event falls out.
///
/// # Parameters
/// - `now` : the caller's clock reading, in nanoseconds
pub fn Sliding::next(self : Sliding, now : Int64) -> @moondate.Span {
  self.forget(now)
  if self.quota < 1 {
    // A quota of zero admits nothing, so no stretch of waiting is enough; a
    // caller is better told a very long time than a zero it would spin on.
    return @moondate.Span::new(days=365L)
  }
  if self.stamps.length() < self.quota {
    return @moondate.Span::new()
  }
  let free = self.stamps[0] + self.window.nanos
  if free <= now {
    @moondate.Span::new()
  } else {
    { nanos: free - now, }
  }
}

///|
/// Drop the events that have fallen out of the window.
fn Sliding::forget(self : Sliding, now : Int64) -> Unit {
  let cut = now - self.window.nanos
  self.stamps.retain(stamp => stamp > cut)
}

///|
/// Where a [`Period`] stands after an event: under the quota, exactly on it, or
/// past it.
///
/// The event that hits the quota is still admitted. It is told apart because it is
/// the last one, which is when a caller would warn that the next will be refused.
#derive(Eq, Debug)
pub(all) enum Quota {
  Under
  Hit
  Over
}

///|
pub extend Quota with Debug::{to_repr}

///|
pub extend Quota with Eq::{not_equal, equal}

///|
/// Where a [`Period`] stands: when the open period started, if one is open, and how
/// many events it has counted. This is the whole of a period's state, and what a
/// store shared by several processes holds.
#derive(Eq, Debug)
pub(all) struct Usage {
  /// When the open period started on the caller's clock, `None` before the first
  /// event and after a period has closed.
  start : Int64?
  /// Events counted in the open period.
  used : Int32
}

///|
pub extend Usage with Debug::{to_repr}

///|
pub extend Usage with Eq::{not_equal, equal}

///|
/// What one event gets from a [`Period`]: where the quota stands after it, and the
/// state to keep.
#derive(Eq, Debug)
pub(all) struct Counted {
  /// Under, on or past the quota once this event is counted.
  quota : Quota
  /// The state to keep, whatever the quota says.
  state : Usage
}

///|
pub extend Counted with Debug::{to_repr}

///|
pub extend Counted with Eq::{not_equal, equal}

///|
/// A fixed window counter: at most `quota` events per `period`, the count starting
/// again when a period ends.
///
/// Cheaper than [`Sliding`], which keeps every stamp, and coarser: a client can
/// spend one quota at the end of a period and another at the start of the next,
/// so forty events can pass in a second that allows twenty. That is the trade a
/// fixed window makes, and the reason it is the cheap one.
///
/// A period opens on the first event after the last one closed. With `align` it
/// opens on a multiple of `period` on the caller's clock instead, so with a Unix
/// clock a daily quota resets at midnight UTC rather than a day after first use.
///
/// Several processes share one quota by keeping its [`Usage`] in a store and
/// counting with [`Period::step`], as a [`Bucket`] shares its tokens.
pub struct Period {
  quota : Int32
  period : Int64
  align : Bool
  mut state : Usage
}

///|
/// A counter no event has touched. A period shorter than a nanosecond is raised to
/// one.
///
/// # Parameters
/// - `quota` : events allowed per period
/// - `period` : how long a period is; anything under a nanosecond is raised to one
/// - `align` : open periods on multiples of `period` on the caller's clock instead of a period after first use
pub fn Period::new(
  quota : Int32,
  period : @moondate.Span,
  align~ : Bool = false,
) -> Period {
  {
    quota,
    period: if period.nanos < 1L {
      1L
    } else {
      period.nanos
    },
    align,
    state: { start: None, used: 0, },
  }
}

///|
/// The state of a period no event has touched: what a store with no entry for the
/// key stands for.
pub fn Period::opening(self : Period) -> Usage {
  ignore(self)
  { start: None, used: 0, }
}

///|
/// `state` with the period closed if `now` is past it. Only an event opens the
/// next one, so asking how much is left does not move where a period starts. A
/// clock that moves backwards stays in the open period rather than earning a fresh
/// quota.
fn Period::rolled(self : Period, state : Usage, now : Int64) -> Usage {
  match state.start {
    Some(start) if now - start >= self.period => { start: None, used: 0, }
    _ => state
  }
}

///|
/// Count one event at `now` against `state`, leaving the period untouched.
///
/// This is the whole of the arithmetic, and the seam for a quota shared by several
/// processes: read the state out of the store — [`Period::opening`] when the key is
/// missing — count here, and write `state` back only if the stored value has not
/// changed in between; on a clash, read again and count again. The compare-and-set
/// is what stops two writers from both spending the last event of a quota, as with
/// [`Bucket::step`].
///
/// # Parameters
/// - `state` : the period's state as the store has it
/// - `now` : the caller's clock reading, in nanoseconds
pub fn Period::step(self : Period, state : Usage, now : Int64) -> Counted {
  let state = self.rolled(state, now)
  let start = match state.start {
    Some(start) => start
    None if self.align => now - (now % self.period + self.period) % self.period
    None => now
  }
  let used = state.used + 1
  let quota = if used < self.quota {
    Under
  } else if used == self.quota {
    Hit
  } else {
    Over
  }
  { quota, state: { start: Some(start), used, }, }
}

///|
/// Count one event at `now` and say where the quota stands.
///
/// # Parameters
/// - `now` : the caller's clock reading, in nanoseconds
pub fn Period::take(self : Period, now : Int64) -> Quota {
  let counted = self.step(self.state, now)
  self.state = counted.state
  counted.quota
}

///|
/// How many events remain in the current period.
///
/// # Parameters
/// - `now` : the caller's clock reading, in nanoseconds
pub fn Period::left(self : Period, now : Int64) -> Int32 {
  self.state = self.rolled(self.state, now)
  if self.state.used >= self.quota {
    0
  } else {
    self.quota - self.state.used
  }
}

///|
/// How long until the current period ends. Zero when no period is open, because
/// the next event would open one.
///
/// # Parameters
/// - `now` : the caller's clock reading, in nanoseconds
pub fn Period::next(self : Period, now : Int64) -> @moondate.Span {
  self.state = self.rolled(self.state, now)
  match self.state.start {
    Some(start) => { nanos: start + self.period - now, }
    None => @moondate.Span::new()
  }
}

///|
/// A count with a ceiling: the concurrency limit a server puts on itself.
///
/// Unlike a rate, this does not depend on time at all — it is how many things are
/// happening at once. A server past its limit answers rather than queues, which is
/// what keeps a queue from becoming the outage.
pub struct Gate {
  limit : Int32
  mut held : Int32
}

///|
/// A gate that nothing has entered.
///
/// # Parameters
/// - `limit` : how many calls may be inside at once; zero or less means no ceiling
pub fn Gate::new(limit : Int32) -> Gate {
  { limit, held: 0, }
}

///|
/// Whether there is room, taking it if there is.
pub fn Gate::enter(self : Gate) -> Bool {
  if self.limit > 0 && self.held >= self.limit {
    return false
  }
  self.held = self.held + 1
  true
}

///|
/// Give the room back.
pub fn Gate::leave(self : Gate) -> Unit {
  if self.held > 0 {
    self.held = self.held - 1
  }
}

///|
/// How many are inside.
pub fn Gate::held(self : Gate) -> Int32 {
  self.held
}