module Utils.Sql
  ( setSerializable, setSerializable'
  ) where

import ClassyPrelude.Yesod
import Numeric.Natural
import Settings.Log

import Database.PostgreSQL.Simple (SqlError)
import Database.PostgreSQL.Simple.Errors (isSerializationError)
import Control.Monad.Catch (MonadMask)

import Database.Persist.Sql
import Database.Persist.Sql.Raw.QQ

import Control.Retry

import Control.Lens ((&))


setSerializable :: forall m a. (MonadLogger m, MonadMask m, MonadIO m, ReadLogSettings (SqlPersistT m)) => SqlPersistT m a -> SqlPersistT m a
setSerializable = setSerializable' $ fullJitterBackoff 1e3 & limitRetriesByCumulativeDelay 10e6

setSerializable' :: forall m a. (MonadLogger m, MonadMask m, MonadIO m, ReadLogSettings (SqlPersistT m)) => RetryPolicyM (SqlPersistT m) -> SqlPersistT m a -> ReaderT SqlBackend m a
setSerializable' policy act = do
  LogSettings{logSerializableTransactionRetryLimit} <- readLogSettings
  didCommit <- newTVarIO False
  recovering policy (skipAsyncExceptions `snoc` logRetries suggestRetry (logRetry logSerializableTransactionRetryLimit)) $ act' didCommit
  where
    suggestRetry :: SqlError -> ReaderT SqlBackend m Bool
    suggestRetry = return . isSerializationError

    logRetry :: Maybe Natural
             -> Bool -- ^ Will retry
             -> SqlError
             -> RetryStatus
             -> ReaderT SqlBackend m ()
    logRetry _ shouldRetry@False err status = $logErrorS "SQL.setSerializable" . pack $ defaultLogMsg shouldRetry err status
    logRetry (Just limit) shouldRetry err status
      | fromIntegral limit <= rsIterNumber status = $logInfoS "SQL.setSerializable" . pack $ defaultLogMsg shouldRetry err status
    logRetry _ shouldRetry err status = $logDebugS "SQL.setSerializable" . pack $ defaultLogMsg shouldRetry err status

    act' :: TVar Bool -> RetryStatus -> ReaderT SqlBackend m a
    act' didCommit RetryStatus{..} = do
      prevCommited <- atomically $ swapTVar didCommit False
      $logDebugS "SQL.setSerializable" $ "prevCommited = " <> tshow prevCommited <> "; rsIterNumber = " <> tshow rsIterNumber
      if
        | rsIterNumber == 0 -> [executeQQ|SET TRANSACTION ISOLATION LEVEL SERIALIZABLE|] *> act''
        | prevCommited      -> [executeQQ|BEGIN TRANSACTION ISOLATION LEVEL SERIALIZABLE|] *> act''
        | otherwise         -> transactionUndoWithIsolation Serializable *> act''
      where act'' = do
              res <- act
              atomically $ writeTVar didCommit True
              transactionSaveWithIsolation ReadCommitted
              return res