content: {-# LANGUAGE TemplateHaskell #-}
  
  -- | Module containing ratelimit stuff
  module Calamity.HTTP.Internal.Ratelimit (
    newRateLimitState,
    doRequest,
    RatelimitEff (..),
    getRatelimitState,
  ) where
  
  import Calamity.HTTP.Internal.Route
  import Calamity.HTTP.Internal.Types
  import Calamity.Internal.Utils
  import Calamity.Types.LogEff
  import Calamity.Types.TokenEff
  import Control.Applicative
  import Control.Concurrent
  import Control.Concurrent.Event (Event)
  import Control.Concurrent.Event qualified as E
  import Control.Concurrent.STM
  import Control.Exception.Safe qualified as Ex
  import Control.Monad
  import Data.Aeson qualified as Aeson
  import Data.Aeson.Optics
  import Data.ByteString qualified as B
  import Data.ByteString.Lazy qualified as LB
  import Data.Maybe
  import Data.Text qualified as T
  import Data.Time
  import Data.Time.Clock.POSIX
  import Network.HTTP.Client (responseStatus)
  import Network.HTTP.Req
  import Network.HTTP.Types
  import Optics
  import Optics.Operators.Unsafe ((^?!))
  import Polysemy (Sem, makeSem)
  import Polysemy qualified as P
  import StmContainers.Map qualified as SC
  import Prelude hiding (error)
  import Prelude qualified
  
  data RatelimitEff m a where
    GetRatelimitState :: RatelimitEff m RateLimitState
  
  makeSem ''RatelimitEff
  
  newRateLimitState :: IO RateLimitState
  newRateLimitState = RateLimitState <$> SC.newIO <*> SC.newIO <*> E.newSet
  
  data Ratelimit
    = KnownRatelimit Bucket
    | UnknownRatelimit RouteKey
  
  getRateLimit :: RateLimitState -> RouteKey -> STM Ratelimit
  getRateLimit s h = do
    bucketKey <- SC.lookup h $ bucketKeys s
    bucket <- join <$> traverse (`SC.lookup` buckets s) bucketKey
    case bucket of
      Just bucket' ->
        pure $ KnownRatelimit bucket'
      Nothing ->
        pure $ UnknownRatelimit h
  
  mergeBucketStates :: BucketState -> BucketState -> BucketState
  mergeBucketStates old new =
    new
      { ongoing = old ^. #ongoing
      , -- we only ignore the previous 'remaining' if we've not reset yet and the
        -- reset time has changed
        remaining =
          if isJust (old ^. #resetTime) && old ^. #resetKey /= new ^. #resetKey
            then min (old ^. #remaining) (new ^. #remaining)
            else new ^. #remaining
      , -- only take the new resetTime if it actually changed
        resetTime =
          if old ^. #resetKey /= new ^. #resetKey
            then new ^. #resetTime
            else old ^. #resetTime
      }
  
  updateKnownBucket :: Bucket -> BucketState -> STM ()
  updateKnownBucket bucket bucketState = modifyTVar' (bucket ^. #state) (`mergeBucketStates` bucketState)
  
  {- | Knowing the bucket for a route, and the ratelimit info, map the route to
   the bucket key and retrieve the bucket
  -}
  updateBucket :: RateLimitState -> RouteKey -> B.ByteString -> BucketState -> STM Bucket
  updateBucket s h b bucketState = do
    bucketKey <- SC.lookup h $ bucketKeys s
    case bucketKey of
      Just bucketKey' -> do
        -- if we know the bucket key here, then the bucket has already been made
        -- if the given bucket key is different than the known bucket key then oops
        bucket <- SC.lookup bucketKey' $ buckets s
        case bucket of
          Just bucket' -> do
            modifyTVar' (bucket' ^. #state) (`mergeBucketStates` bucketState)
            pure bucket'
          Nothing -> Prelude.error "Not possible"
      Nothing -> do
        -- we didn't know the key to this bucket, but we might know the bucket
        -- if we truly don't know the bucket, then make a new one
        bs <- do
          bucket <- SC.lookup b $ buckets s
          case bucket of
            Just bs -> pure bs
            Nothing -> do
              bs <- Bucket <$> newTVar bucketState
              SC.insert bs b $ buckets s
              pure bs
  
        SC.insert b h $ bucketKeys s
        pure bs
  
  resetBucket :: Bucket -> STM ()
  resetBucket bucket =
    modifyTVar'
      (bucket ^. #state)
      ( \bs ->
          bs
            & #remaining
            .~ bs
            ^. #limit
            & #resetTime
            .~ Nothing
      )
  
  canResetBucketNow :: UTCTime -> BucketState -> Bool
  canResetBucketNow _ BucketState {ongoing} | ongoing > 0 = False
  -- don't allow resetting the bucket if there's ongoing requests, we'll wait
  -- until another request finishes and updates the counter
  canResetBucketNow now bs = case bs ^. #resetTime of
    Just rt -> now > rt
    Nothing -> False
  
  -- canResetBucket :: BucketState -> Bool
  -- canResetBucket bs = isNothing $ bs ^. #startedWaitingTime
  
  shouldWaitForUnlock :: BucketState -> Bool
  shouldWaitForUnlock BucketState {remaining = 0, ongoing} = ongoing > 0
  shouldWaitForUnlock _ = False
  
  data WaitDelay
    = WaitUntil UTCTime
    | WaitRetrySoon
    | GoNow
    deriving (Show)
  
  intoWaitDelay :: Maybe UTCTime -> WaitDelay
  intoWaitDelay (Just t) = WaitUntil t
  intoWaitDelay Nothing = WaitRetrySoon
  
  -- | Maybe wait for a bucket, updating its state to say we used it
  useBucketOnce :: Bucket -> IO ()
  useBucketOnce bucket = go 0
    where
      go :: Int -> IO ()
      go tries = do
        now <- getCurrentTime
        mWaitDelay <- atomically $ do
          s <- readTVar $ bucket ^. #state
  
          -- -- [0]
          -- -- if there are ongoing requests, wait for them to finish and deliver
          -- -- truth on the current ratelimit state
          when
            (shouldWaitForUnlock s)
            retry
  
          -- if there are no ongoing requests, and the bucket reset time has lapsed,
          -- we can just reset the bucket.
          --
          -- if we've already reset the bucket then there should be an ongoing
          -- request so we'll just end up waiting for that to finish
          when
            (canResetBucketNow now s)
            (resetBucket bucket)
  
          s <- readTVar $ bucket ^. #state
  
          if s ^. #remaining - s ^. #ongoing > 0
            then do
              -- there are tokens remaining for us to use
              modifyTVar'
                (bucket ^. #state)
                ( (#remaining %~ pred)
                    . (#ongoing %~ succ)
                )
              pure GoNow
            else do
              -- the bucket has expired, there are no ongoing requests because of
              -- [0] wait and then retry after we can unlock the bucket
              pure (intoWaitDelay $ s ^. #resetTime)
  
        case mWaitDelay of
          WaitUntil waitUntil -> do
            if waitUntil < now
              then threadDelayMS 20
              else -- if the reset is in the past, we're fucked
                threadDelayUntil waitUntil
            -- if we needed to sleep, go again so that multiple concurrent requests
            -- don't exceed the bucket, to ensure we don't sit in a loop if a
            -- request dies on us, bail out after 50 loops
            when (tries < 50) $ go (tries + 1) -- print "bailing after number of retries"
          WaitRetrySoon -> do
            threadDelayMS 20
            when (tries < 50) $ go (tries + 1) -- print "bailing after number of retries"
          GoNow -> do
            -- print "ok going forward with request"
            pure ()
  
  doDiscordRequest :: (P.Members '[RatelimitEff, TokenEff, LogEff, P.Embed IO] r) => IO LbsResponse -> Sem r DiscordResponseType
  doDiscordRequest r = do
    r'' <- P.embed $ Ex.catchAny (Right <$> r) (pure . Left . Ex.displayException)
    case r'' of
      Right r' -> do
        let status = responseStatus . toVanillaResponse $ r'
        if
          | statusIsSuccessful status -> do
              let resp = responseBody r'
              debug $ "Got good response from discord: " <> (T.pack . show $ status)
              now <- P.embed getCurrentTime
              let rlHeaders = buildBucketState now r'
              pure $ Good resp rlHeaders
          | status == status429 -> do
              now <- P.embed getCurrentTime
              let resp = responseBody r'
              case (resp ^? _Value, buildBucketState now r') of
                (Just !rv, bs) ->
                  pure $ Ratelimited (parseRetryAfter now rv) (isGlobal rv) bs
                _ ->
                  pure $ ServerError (statusCode status)
          | statusIsClientError status -> do
              let err = responseBody r'
              error . T.pack $ "Something went wrong: " <> show err <> ", response: " <> show r'
              pure $ ClientError (statusCode status) err
          | otherwise -> do
              debug . T.pack $ "Got server error from discord: " <> (show . statusCode $ status)
              pure $ ServerError (statusCode status)
      Left e -> do
        error . T.pack $ "Something went wrong with the http client: " <> e
        pure . InternalResponseError $ T.pack e
  
  -- | Parse a ratelimit header returning when it unlocks
  parseRateLimitHeader :: (HttpResponse r) => UTCTime -> r -> Maybe UTCTime
  parseRateLimitHeader now r = computedEnd <|> end
    where
      computedEnd :: Maybe UTCTime
      computedEnd = flip addUTCTime now <$> resetAfter
  
      resetAfter :: Maybe NominalDiffTime
      resetAfter = realToFrac <$> responseHeader r "X-Ratelimit-Reset-After" ^? _Just % _Double
  
      end :: Maybe UTCTime
      end =
        posixSecondsToUTCTime
          . realToFrac
          <$> responseHeader r "X-Ratelimit-Reset"
            ^? _Just
            % _Double
  
  buildBucketState :: (HttpResponse r) => UTCTime -> r -> Maybe (BucketState, B.ByteString)
  buildBucketState now r = (,) <$> bs <*> bucketKey
    where
      remaining = responseHeader r "X-RateLimit-Remaining" ^? _Just % _Integral
      limit = responseHeader r "X-RateLimit-Limit" ^? _Just % _Integral
      resetKey = ceiling <$> responseHeader r "X-RateLimit-Reset" ^? _Just % _Double
      resetTime = parseRateLimitHeader now r
      bs = BucketState resetTime <$> resetKey <*> remaining <*> limit <*> pure 0
      bucketKey = responseHeader r "X-RateLimit-Bucket"
  
  -- | Parse the retry after field, returning when to retry
  parseRetryAfter :: UTCTime -> Aeson.Value -> UTCTime
  parseRetryAfter now r = addUTCTime retryAfter now
    where
      retryAfter = realToFrac $ r ^?! key "retry_after" % _Double
  
  isGlobal :: Aeson.Value -> Bool
  isGlobal r = r ^? key "global" % _Bool == Just True
  
  -- Either (Either a a) b
  data ShouldRetry a b
    = Retry a
    | RFail a
    | RGood b
  
  retryRequest ::
    (P.Members '[RatelimitEff, TokenEff, LogEff, P.Embed IO] r) =>
    -- | number of retries
    Int ->
    -- | action to perform
    Sem r (ShouldRetry a b) ->
    Sem r (Either a b)
  retryRequest maxRetries action = retryInner 0
    where
      retryInner numRetries = do
        res <- action
        case res of
          Retry r | numRetries > maxRetries -> do
            debug . T.pack $ "Request failed after " <> show maxRetries <> " retries"
            pure $ Left r
          Retry _ ->
            retryInner (numRetries + 1)
          RFail r -> do
            debug "Request failed due to error response"
            pure $ Left r
          RGood r ->
            pure $ Right r
  
  threadDelayMS :: Int -> IO ()
  threadDelayMS ms = threadDelay (1000 * ms)
  
  tenMS :: NominalDiffTime
  tenMS = 0.01
  
  threadDelayUntil :: UTCTime -> IO ()
  threadDelayUntil when = do
    let when' = addUTCTime tenMS when -- lol
    now <- getCurrentTime
    let msUntil = ceiling . (* 1000) . realToFrac @_ @Double $ diffUTCTime when' now
    threadDelayMS msUntil
  
  -- Run a single request
  doSingleRequest ::
    (P.Members '[RatelimitEff, TokenEff, LogEff, P.Embed IO] r) =>
    RateLimitState ->
    Route ->
    -- | Global lock
    Event ->
    -- | Request action
    IO LbsResponse ->
    Sem r (ShouldRetry RestError LB.ByteString)
  doSingleRequest rlstate route gl r = do
    P.embed $ E.wait (globalLock rlstate)
  
    rl <- P.embed . atomically $ getRateLimit rlstate (routeKey route)
  
    case rl of
      KnownRatelimit bucket ->
        P.embed $ useBucketOnce bucket
      _ -> debug "unknown ratelimit"
  
    r' <- doDiscordRequest r
  
    case r' of
      Good v rlHeaders -> do
        void . P.embed . atomically $ do
          case rl of
            KnownRatelimit bucket ->
              modifyTVar' (bucket ^. #state) (#ongoing %~ pred)
            _ -> pure ()
          case (rl, rlHeaders) of
            (KnownRatelimit bucket, Just (bs, _bk)) ->
              updateKnownBucket bucket bs
            (_, Just (bs, bk)) ->
              void $ updateBucket rlstate (routeKey route) bk bs
            (_, Nothing) -> pure ()
        pure $ RGood v
      Ratelimited unlockWhen False (Just (bs, bk)) -> do
        debug . T.pack $ "429 ratelimited on route, retrying at " <> show unlockWhen
  
        P.embed . atomically $ do
          case rl of
            KnownRatelimit bucket -> do
              modifyTVar' (bucket ^. #state) (#ongoing %~ pred)
              updateKnownBucket bucket bs
            _ -> void $ updateBucket rlstate (routeKey route) bk bs
  
        P.embed $ do
          threadDelayUntil unlockWhen
  
        pure $ Retry (HTTPError 429 Nothing)
      Ratelimited unlockWhen False _ -> do
        debug "Internal error (ratelimited but no headers), retrying"
        case rl of
          KnownRatelimit bucket ->
            void . P.embed . atomically $ modifyTVar' (bucket ^. #state) (#ongoing %~ pred)
          _ -> pure ()
  
        P.embed $ threadDelayUntil unlockWhen
        pure $ Retry (HTTPError 429 Nothing)
      Ratelimited unlockWhen True bs -> do
        debug "429 ratelimited globally"
  
        P.embed $ do
          atomically $ do
            case rl of
              KnownRatelimit bucket ->
                modifyTVar' (bucket ^. #state) (#ongoing %~ pred)
              _ -> pure ()
            case bs of
              Just (bs', bk) ->
                void $ updateBucket rlstate (routeKey route) bk bs'
              Nothing ->
                pure ()
  
          E.clear gl
          threadDelayUntil unlockWhen
          E.set gl
        pure $ Retry (HTTPError 429 Nothing)
      ServerError c -> do
        debug "Server failed, retrying"
        case rl of
          KnownRatelimit bucket ->
            P.embed $ useBucketOnce bucket
          _ -> debug "unknown ratelimit"
        pure $ Retry (HTTPError c Nothing)
      InternalResponseError c -> do
        debug "Internal error, retrying"
        case rl of
          KnownRatelimit bucket ->
            P.embed $ useBucketOnce bucket
          _ -> debug "unknown ratelimit"
        pure $ Retry (InternalClientError c)
      ClientError c v -> do
        case rl of
          KnownRatelimit bucket ->
            P.embed $ useBucketOnce bucket
          _ -> debug "unknown ratelimit"
        pure $ RFail (HTTPError c $ Aeson.decode v)
  
  doRequest :: (P.Members '[RatelimitEff, TokenEff, LogEff, P.Embed IO] r) => RateLimitState -> Route -> IO LbsResponse -> Sem r (Either RestError LB.ByteString)
  doRequest rlstate route action =
    retryRequest
      5
      (doSingleRequest rlstate route (globalLock rlstate) action)
  