module Webserver.Data.Sql.Cache.RateLimit where

----------------------------------------------------------------------------------------------------

import           Data.Functor.Contravariant (contramap)
import           Data.Int                   (Int32)
import qualified Data.Time.Clock            as Clock
import           Data.Tuple                 (fst, snd)
import           GHC.Real                   (fromIntegral)
import qualified Hasql.Decoders             as Decoders
import qualified Hasql.Encoders             as Encoders
import qualified Hasql.Pool                 as Pool
import qualified Hasql.Session              as Session
import qualified Hasql.Statement            as Statement
import           Webserver.Control.Nil      (Nil (..))
import qualified Webserver.Control.Sql      as Sql

----------------------------------------------------------------------------------------------------

$(Sql.generateTable "cache" "rate_limits")

----------------------------------------------------------------------------------------------------

instance Nil RateLimits where
  nil = RateLimits {id = nil, key = nil, hits = nil, windowStart = nil}

----------------------------------------------------------------------------------------------------

{-  Atomic UPSERT that increments hits within a window or resets if the window has expired.
    Returns the current hit count after the upsert.
    Parameters: $1 = key (Text), $2 = window duration (Interval).
-}
checkRateLimitStatement :: Statement.Statement (Text, Clock.DiffTime) Int32
checkRateLimitStatement =
  Statement.Statement
    "INSERT INTO cache.rate_limits (key, hits, window_start) \
    \VALUES ($1, 1, NOW()) \
    \ON CONFLICT (key) DO UPDATE \
    \SET hits = CASE \
    \    WHEN cache.rate_limits.window_start + $2 < NOW() THEN 1 \
    \    ELSE cache.rate_limits.hits + 1 \
    \    END, \
    \    window_start = CASE \
    \    WHEN cache.rate_limits.window_start + $2 < NOW() THEN NOW() \
    \    ELSE cache.rate_limits.window_start \
    \    END \
    \RETURNING hits"
    ( contramap fst (Encoders.param (Encoders.nonNullable Encoders.text))
        <> contramap snd (Encoders.param (Encoders.nonNullable Encoders.interval))
    )
    (Decoders.singleRow (Decoders.column (Decoders.nonNullable Decoders.int4)))
    True

----------------------------------------------------------------------------------------------------

{-  Check if a request identified by the given key is within the rate limit.
    Returns True if the request is allowed, False if rate limited.
-}
checkRateLimit :: Pool.Pool -> Text -> Int -> Clock.DiffTime -> IO Bool
checkRateLimit pool key_ maxHits window = do
  hits_ <-
    Sql.runSession pool
      <| Session.statement (key_, window) checkRateLimitStatement
  return <| fromIntegral hits_ <= maxHits