{-# Language DerivingVia #-}
{-# Language OverloadedStrings #-}
module HomeAssistant.Runtime.RateLimit where
import Control.Monad.IO.Class (MonadIO, liftIO)
import Data.Time (NominalDiffTime)
import Control.Concurrent.STM (newTVarIO, atomically, readTVar, writeTVar)
import Data.Time.Clock.POSIX (getPOSIXTime)
import Katip.Monadic (KatipContext(..))
import Control.Retry (fullJitterBackoff, limitRetries, retrying, RetryStatus (..))
import Data.Maybe (isNothing)
import Control.Exception.Annotated (Exception, checkpointCallStack, throw)
import Control.Monad.Catch (MonadCatch)
import Katip (Katip, Severity (WarningS), logFM, ls)
import Control.Monad (when)
import System.Metrics.Counter (Counter)
import qualified System.Metrics.Counter as Counter
import System.Metrics (Store, createCounter)

data RateLimitMetrics = RateLimitMetrics
  { rlRetries :: Counter
  , rlRejects :: Counter
  }


registerRateLimitMetrics :: MonadIO m => Store -> m RateLimitMetrics
registerRateLimitMetrics store = liftIO $ RateLimitMetrics
  <$> createCounter "ratelimit.retry" store
  <*> createCounter "ratelimit.exhausted" store

newtype RateLimited m a = RateLimited (m a)
  deriving (Functor, Applicative, Monad, MonadIO, KatipContext, Katip) via m


data RateLimiter = RateLimiter
  { metrics :: RateLimitMetrics
  , limiter :: forall m a. (KatipContext m, MonadIO m) => RateLimited m a -> m (Maybe a)
  }


slidingWindowLimiter
  :: RateLimitMetrics
  -> NominalDiffTime -- ^ Within how many seconds
  -> Int -- ^ How many requests we can make
  -> IO RateLimiter
slidingWindowLimiter metrics timeWindow size = do
  window <- newTVarIO []
  return $ RateLimiter metrics $ \(RateLimited action) -> do
    now <- liftIO getPOSIXTime
    allowed <- liftIO $ atomically $ do
      elements <- takeWhile (\x -> x > now - timeWindow) <$> readTVar window
      if length elements >= size
         then pure False
         else True <$ writeTVar window (now : elements)
    if allowed
       then Just <$> action
       else pure Nothing

data RateLimitException = RateLimitExhausted
  deriving Show

instance Exception RateLimitException

runRateLimited :: (MonadIO m, KatipContext m, MonadCatch m) => RateLimiter -> RateLimited m a -> m a
runRateLimited (RateLimiter metrics run) action = checkpointCallStack $ do
  x <- retrying retryPolicy (const $ pure . isNothing) $ \RetryStatus{rsIterNumber} -> do
    when (rsIterNumber > 0) $ do
      logFM WarningS $ ls $ "Rate limited, attempt no. " <> show rsIterNumber
      incRetry
    run action
  maybe (incReject >> throw RateLimitExhausted) pure x
  where
    retryPolicy = fullJitterBackoff 500_000 <> limitRetries 5
    incRetry = liftIO $ Counter.inc (rlRetries metrics)
    incReject = liftIO $ Counter.inc (rlRejects metrics)