module Foundation.Yesod.Persist
  ( runDB, getDBRunner
  , runDB', getDBRunner'
  -- , runCachedDBRunner
  -- , runCachedDBRunner'
  , module Foundation.DB
  ) where

import Import.NoFoundation hiding (runDB, getDBRunner)

import Foundation.Type
import Foundation.DB
import Foundation.Authorization

import Database.Persist.Sql (transactionUndo)
import qualified Database.Persist.Sql as SQL
import qualified Database.Persist.SqlBackend.Internal as SQL

import qualified Utils.Pool as Custom

import UnliftIO.Resource (allocate, unprotect)


runDB :: ( YesodPersistBackend UniWorX ~ SqlBackend
         , BearerAuthSite UniWorX
         , HasCallStack
         )
      => YesodDB UniWorX a -> HandlerFor UniWorX a
runDB = runDB' callStack

runDB' :: ( YesodPersistBackend UniWorX ~ SqlBackend
          , BearerAuthSite UniWorX
          )
       => CallStack -> YesodDB UniWorX a -> HandlerFor UniWorX a
runDB' lbl action = do
  $logDebugS "YesodPersist" "runDB"
  let action' = do
        dryRun <- isDryRunDB
        if | dryRun    -> action <* transactionUndo
           | otherwise -> action

  flip (runSqlPoolRetry' action') lbl . appConnPool =<< getYesod

getDBRunner :: ( YesodPersistBackend UniWorX ~ SqlBackend
               , BearerAuthSite UniWorX
               , HasCallStack
               )
            => HandlerFor UniWorX (DBRunner UniWorX, HandlerFor UniWorX ())
getDBRunner = getDBRunner' callStack

getDBRunner' :: ( YesodPersistBackend UniWorX ~ SqlBackend
                , BearerAuthSite UniWorX
                )
             => CallStack -> HandlerFor UniWorX (DBRunner UniWorX, HandlerFor UniWorX ())
getDBRunner' lbl = do
  pool <- getsYesod appConnPool
  UnliftIO{..} <- askUnliftIO
  let withPrep conn f = f (persistBackend conn) (SQL.getStmtConn $ persistBackend conn)
  (relKey, (conn, ident)) <- allocate
    (do
        (conn, ident) <- unliftIO $ Custom.takeResource' pool lbl
        withPrep conn (\c f -> SQL.connBegin c f Nothing)
        return (conn, ident)
    )
    (\(conn, ident) -> do
        withPrep conn SQL.connRollback
        unliftIO $ Custom.releaseResource True pool (conn, ident)
    )

  let cleanup = liftIO $ do
        withPrep conn SQL.connCommit
        unliftIO $ Custom.releaseResource True pool (conn, ident)
        void $ unprotect relKey
      runDBRunner :: forall a. YesodDB UniWorX a -> HandlerFor UniWorX a
      runDBRunner = flip runReaderT conn

  return . (, cleanup) $ DBRunner
    (\action -> do
        let action' = do
              dryRun <- isDryRunDB
              if | dryRun    -> action <* transactionUndo
                 | otherwise -> action
        $logDebugS "YesodPersist" "runDBRunner"
        runDBRunner action'
    )

-- runCachedDBRunner :: ( BackendCompatible backend (YesodPersistBackend UniWorX)
--                      , YesodPersistBackend UniWorX ~ SqlBackend
--                      , BearerAuthSite UniWorX
--                      , HasCallStack
--                      )
--                   => CachedDBRunner backend (HandlerFor UniWorX) a
--                   -> HandlerFor UniWorX a
-- runCachedDBRunner = runCachedDBRunner' callStack

-- runCachedDBRunner' :: ( BackendCompatible backend (YesodPersistBackend UniWorX)
--                       , YesodPersistBackend UniWorX ~ SqlBackend
--                       , BearerAuthSite UniWorX
--                       )
--                    => CallStack
--                    -> CachedDBRunner backend (HandlerFor UniWorX) a
--                    -> HandlerFor UniWorX a
-- runCachedDBRunner' lbl act = do
--   cleanups <- newTVarIO []
--   res <- flip runCachedDBRunnerSTM act $ do
--     (runner, cleanup) <- getDBRunner' lbl
--     atomically . modifyTVar' cleanups $ (:) cleanup
--     return $ fromDBRunner runner
--   mapM_ liftHandler =<< readTVarIO cleanups
--   return res
