module Utils.PersistentTokenBucket
  ( TokenBucketSettings(..)
  , persistentTokenBucketTryAlloc'
  , persistentTokenBucketTryAlloc
  ) where

import Import.NoFoundation


data TokenBucketSettings = TokenBucketSettings
  { tbsIdent :: TokenBucketIdent
  , tbsDepth :: Word64
  , tbsInvRate :: NominalDiffTime
  , tbsInitialValue :: Int64
  }


persistentTokenBucketTryAlloc' :: (MonadHandler m, HasAppSettings (HandlerSite m), Integral a)
                               => TokenBucketIdent
                               -> a
                               -> SqlPersistT m Bool
persistentTokenBucketTryAlloc' tbsIdent tokens = do
  TokenBucketConf{..} <- getsYesod $ views _appPersistentTokenBuckets ($ tbsIdent)
  flip persistentTokenBucketTryAlloc tokens TokenBucketSettings 
    { tbsIdent
    , tbsDepth = tokenBucketDepth
    , tbsInvRate = tokenBucketInvRate
    , tbsInitialValue = tokenBucketInitialValue
    }


persistentTokenBucketTryAlloc :: (MonadIO m, Integral a) => TokenBucketSettings -> a -> SqlPersistT m Bool
persistentTokenBucketTryAlloc TokenBucketSettings{..} (fromIntegral -> tokens) = do
  now <- liftIO getCurrentTime
  TokenBucket{..} <- do
    existingBucket <- get $ TokenBucketKey tbsIdent
    case existingBucket of
      Just bkt -> return bkt
      Nothing -> do
        let bkt = TokenBucket
              { tokenBucketIdent = tbsIdent
              , tokenBucketLastValue = tbsInitialValue
              , tokenBucketLastAccess = now
              }
        insert_ bkt
        return bkt
  let currentValue = fromIntegral tbsDepth `min` tokenBucketLastValue + tokenIncrease
      deltaT = now `diffUTCTime` tokenBucketLastAccess
      (tokenIncrease, deltaT')
        | n < 0 = (pred n, (1 + f) * tbsInvRate)
        | otherwise = (n, f * tbsInvRate)
        where (n, f) = properFraction $ deltaT / tbsInvRate
  if | currentValue < 0 -> return False
     | otherwise -> do
         update (TokenBucketKey tbsIdent) [ TokenBucketLastValue =. currentValue - tokens, TokenBucketLastAccess =. addUTCTime (- deltaT') now ]
         return True