module Foundation.DB
  ( runDBRead, runDBRead'
  , runSqlPoolRetry, runSqlPoolRetry'
  , dbPoolPressured
  , runDBInternal, runDBInternal'
  ) where

import Import.NoFoundation hiding (runDB, getDBRunner)

import Foundation.Type

import qualified Control.Retry as Retry
import GHC.IO.Exception (IOErrorType(OtherError))
  
import Database.Persist.Sql (SqlReadBackend(..))
import Database.Persist.Sql.Raw.QQ (executeQQ)

import qualified Utils.Pool as Custom

  
runSqlPoolRetry :: forall m a backend c.
                   ( HasCallStack
                   , MonadUnliftIO m, BackendCompatible SqlBackend backend
                   , MonadLogger m, MonadMask m
                   )
                => ReaderT backend m a
                -> Custom.Pool' m DBConnLabel c backend
                -> m a
runSqlPoolRetry action pool = runSqlPoolRetry' action pool callStack
  
runSqlPoolRetry' :: forall m a backend c.
                    ( MonadUnliftIO m, BackendCompatible SqlBackend backend
                    , MonadLogger m, MonadMask m
                    )
                 => ReaderT backend m a
                 -> Custom.Pool' m DBConnLabel c backend
                 -> CallStack
                 -> m a
runSqlPoolRetry' action pool lbl = do
  let policy = Retry.fullJitterBackoff 1e3 & Retry.limitRetriesByCumulativeDelay 10e6
      handlers = Retry.skipAsyncExceptions `snoc` Retry.logRetries suggestRetry logRetry
        where suggestRetry :: IOException -> m Bool
              suggestRetry ioExc = return $
                   ioeGetErrorType ioExc == OtherError
                && ioeGetLocation ioExc == "libpq"
              logRetry :: forall e.
                          Exception e
                       => Bool -- ^ Will retry
                       -> e
                       -> Retry.RetryStatus
                       -> m ()
              logRetry shouldRetry@False err status = $logErrorS "runSqlPoolRetry" . pack $ Retry.defaultLogMsg shouldRetry err status
              logRetry shouldRetry@True  err status = $logWarnS "runSqlPoolRetry" . pack $ Retry.defaultLogMsg shouldRetry err status

  Retry.recovering policy handlers $ \Retry.RetryStatus{..} -> do
    $logDebugS "runSqlPoolRetry" $ "rsIterNumber = " <> tshow rsIterNumber
    customRunSqlPool' action pool lbl

runDBRead :: HasCallStack => ReaderT SqlReadBackend (HandlerFor UniWorX) a -> (HandlerFor UniWorX) a
runDBRead = runDBRead' callStack

runDBRead' :: CallStack -> ReaderT SqlReadBackend (HandlerFor UniWorX) a -> (HandlerFor UniWorX) a
runDBRead' lbl action = do
  $logDebugS "YesodPersist" "runDBRead"
  flip (runSqlPoolRetry' . withReaderT SqlReadBackend $ [executeQQ|SET TRANSACTION READ ONLY|] *> action) lbl . appConnPool =<< getYesod
  
runDBInternal :: HasCallStack
              => ReaderT SqlBackend (HandlerFor UniWorX) a -> HandlerFor UniWorX a
runDBInternal = runDBInternal' callStack

runDBInternal' :: CallStack -> ReaderT SqlBackend (HandlerFor UniWorX) a -> HandlerFor UniWorX a
runDBInternal' lbl action = do
  $logDebugS "YesodPersist" "runDBInternal"
  flip (runSqlPoolRetry' action) lbl . appConnPool =<< getYesod

dbPoolPressured :: ( MonadHandler m
                   , HandlerSite m ~ UniWorX
                   )
                => m Bool
dbPoolPressured = do
  connPool <- getsYesod @_ @(Custom.Pool' IO _ _ _) appConnPool
  case Custom.getPoolMaxAvailable connPool of
    Nothing -> return False
    Just lim -> atomically $ (>= lim) <$> Custom.getPoolInUseCount connPool