-- SPDX-FileCopyrightText: 2022 Gregor Kleen <gregor.kleen@ifi.lmu.de>,Steffen Jost <jost@tcs.ifi.lmu.de>
--
-- SPDX-License-Identifier: AGPL-3.0-or-later

module Utils.DB where

import ClassyPrelude.Yesod

import qualified Data.List as List
import qualified Data.Map as Map
import qualified Data.Set as Set
import qualified Database.Esqueleto.Legacy as E
-- import Database.Persist -- currently not needed here

import Utils
import Control.Lens
import Control.Lens.Extras (is)

import Control.Monad.Catch hiding (bracket)

import qualified Utils.Pool as Custom

import Database.Persist.Sql (runSqlConn)

import GHC.Stack (HasCallStack, CallStack, callStack)

-- import Control.Monad.Fix (MonadFix)
-- import Control.Monad.Fail (MonadFail)

-- import Control.Monad.Trans.Reader (withReaderT)


emptyOrIn :: PersistField typ
          => E.SqlExpr (E.Value typ) -> Set typ -> E.SqlExpr (E.Value Bool)
emptyOrIn criterion testSet
  | Set.null testSet = E.val True
  | otherwise        = criterion `E.in_` E.valList (Set.toList testSet)

entities2map :: PersistEntity record => [Entity record] -> Map (Key record) record
entities2map = foldl' (\m entity -> Map.insert (entityKey entity) (entityVal entity) m) Map.empty

getJustBy :: (PersistRecordBackend record backend, PersistUniqueRead backend, MonadIO m, MonadThrow m, Show (Unique record))
          => Unique record -> ReaderT backend m (Entity record)
getJustBy u = getBy u >>= maybe
  (throwM . PersistForeignConstraintUnmet $ tshow u)
  return

getKeyBy :: (PersistRecordBackend record backend, PersistUniqueRead backend, MonadIO m)
         => Unique record -> ReaderT backend m (Maybe (Key record))
getKeyBy u = fmap entityKey <$> getBy u -- TODO optimize this, so that DB does not deliver entire record!

getKeyJustBy :: (PersistRecordBackend record backend, PersistUniqueRead backend, MonadIO m, MonadThrow m, Show (Unique record))
             => Unique record -> ReaderT backend m (Key record)
getKeyJustBy u = getKeyBy u >>= maybe
  (throwM . PersistForeignConstraintUnmet $ tshow u)
  return

getKeyBy404 :: (PersistRecordBackend record backend, PersistUniqueRead backend, MonadHandler m)
            => Unique record -> ReaderT backend m (Key record)
getKeyBy404 u = getKeyBy u >>= maybe notFound return

getEntity404 :: (PersistStoreRead backend, PersistRecordBackend record backend, MonadHandler m)
             => Key record -> ReaderT backend m (Entity record)
getEntity404 k = Entity k <$> get404 k

existsBy :: (PersistRecordBackend record backend, PersistUniqueRead backend, MonadIO m)
         => Unique record -> ReaderT backend m Bool
existsBy = fmap (is _Just) . getKeyBy

existsBy404 :: (PersistRecordBackend record backend, PersistUniqueRead backend, MonadHandler m)
            => Unique record -> ReaderT backend m ()
existsBy404 = bool notFound (return ()) <=< fmap (is _Just) . getKeyBy

existsKey :: (PersistRecordBackend record backend, PersistQueryRead backend, MonadIO m)
          => Key record -> ReaderT backend m Bool
existsKey = exists . pure . (persistIdField ==.)

-- -- Available in persistent since 2.11.0.0
-- exists :: (PersistRecordBackend record backend, PersistQueryRead backend, MonadIO m)
--        => [Filter record] -> ReaderT backend m Bool
-- exists = fmap (not . null) . flip selectKeysList [LimitTo 1]

exists404 :: (PersistRecordBackend record backend, PersistQueryRead backend, MonadHandler m)
          => [Filter record] -> ReaderT backend m ()
exists404 = bool (return ()) notFound <=< fmap null . flip selectKeysList [LimitTo 1]

existsKey404 :: (PersistRecordBackend record backend, PersistQueryRead backend, MonadHandler m)
             => Key record -> ReaderT backend m ()
existsKey404 = bool notFound (return ()) <=< existsKey

updateBy :: (PersistUniqueRead backend, PersistStoreWrite backend, MonadIO m, PersistRecordBackend record backend )
         => Unique record -> [Update record] -> ReaderT backend m ()
updateBy uniq updates = do
  key <- getKeyBy uniq
  for_ key $ flip update updates

updateGetEntity :: (PersistStoreWrite backend, MonadIO m, PersistRecordBackend record backend) => Key record -> [Update record] -> ReaderT backend m (Entity record)
updateGetEntity k = fmap (Entity k) . updateGet k

-- | insert or replace a record based on a single uniqueness constraint
--   this function was meant to be supplied with the uniqueness constraint, but it would be unsafe if the uniqueness constraint would not match the supplied record
replaceBy :: (PersistUniqueWrite backend, MonadIO m, OnlyOneUniqueKey record, PersistEntityBackend record ~ BaseBackend backend) 
              => record -> ReaderT backend m ()
replaceBy r = do 
  u <- onlyUnique r
  deleteBy u
  insert_ r

-- | Like 'myReplaceUnique' or 'replaceUnique' but with reversed result: returns 'Nothing' if the replacement was not possible,
--   and 'Just key' for the successfully replaced record
uniqueReplace :: ( MonadIO m
                 , Eq (Unique record)
                 , PersistRecordBackend record backend
                 , PersistUniqueWrite backend
                 )
              => Key record -> record -> ReaderT backend m (Maybe (Key record))
uniqueReplace key datumNew = flipMaybe key <$> myReplaceUnique key datumNew

-- | Identical to 'Database.Persist.Class', except for the better type signature (original requires Eq  record which is not needed anyway)
myReplaceUnique :: ( MonadIO m
                   , Eq (Unique record)
                   , PersistRecordBackend record backend
                   , PersistUniqueWrite backend
                   )
                => Key record -> record -> ReaderT backend m (Maybe (Unique record))
myReplaceUnique key datumNew = getJust key >>= replaceOriginal
  where
    uniqueKeysNew = persistUniqueKeys datumNew
    replaceOriginal original = do
        conflict <- checkUniqueKeys changedKeys
        case conflict of
            Nothing -> replace key datumNew >> return Nothing
            (Just conflictingKey) -> return $ Just conflictingKey
      where
        changedKeys = uniqueKeysNew List.\\ uniqueKeysOriginal
        uniqueKeysOriginal = persistUniqueKeys original

replaceEntity :: ( MonadIO m
                 , PersistRecordBackend record backend
                 , PersistStoreWrite backend
                 )
              => Entity record -> ReaderT backend m ()
replaceEntity Entity{..} = replace entityKey entityVal

checkUniqueKeys :: ( MonadIO m
                   , PersistUniqueRead backend
                   , PersistRecordBackend record backend
                   )
                => [Unique record] -> ReaderT backend m (Maybe (Unique record))
checkUniqueKeys [] = return Nothing
checkUniqueKeys (x:xs) = do
    y <- getBy x
    case y of
        Nothing -> checkUniqueKeys xs
        Just _ -> return (Just x)


put :: ( MonadIO m
       , PersistUniqueWrite backend
       , PersistRecordBackend record backend
       )
    => record -> ReaderT backend m (Key record)
-- ^ `insert`, but remove all records with matching uniqueness constraints first
put v = do
  forM_ (persistUniqueKeys v) deleteBy
  insert v

selectMaybe :: forall record backend m.
               ( MonadIO m
               , PersistQueryRead backend
               , PersistRecordBackend record backend
               )
            => [Filter record] -> [SelectOpt record]
            -> ReaderT backend m (Maybe (Entity record))
selectMaybe fltrs opts = listToMaybe <$> selectList fltrs (LimitTo 1 : opts')
  where opts' = filter (not . isLimit) opts
        isLimit = \case
          LimitTo _ -> True
          _other    -> False


type DBConnLabel = CallStack

customRunSqlPool :: (HasCallStack, MonadUnliftIO m, BackendCompatible SqlBackend backend)
                 => ReaderT backend m a
                 -> Custom.Pool' m DBConnLabel c backend
                 -> m a
customRunSqlPool act p = customRunSqlPool' act p callStack

customRunSqlPool' :: (MonadUnliftIO m, BackendCompatible SqlBackend backend)
                  => ReaderT backend m a
                  -> Custom.Pool' m DBConnLabel c backend
                  -> CallStack
                  -> m a
customRunSqlPool' act p label = Custom.withResource' p label $ runSqlConn act


class WithRunDB backend m' m | m -> backend m' where
  useRunDB :: ReaderT backend m' a -> m a

instance WithRunDB backend m (ReaderT backend m) where
  useRunDB = id

-- newtype DBRunner' backend m = DBRunner' { runDBRunner' :: forall b. ReaderT backend m b -> m b }

-- _DBRunner' :: Iso' (DBRunner site) (DBRunner' (YesodPersistBackend site) (HandlerFor site))
-- _DBRunner' = iso fromDBRunner' toDBRunner
--   where
--     fromDBRunner' :: forall site.
--                      DBRunner site
--                   -> DBRunner' (YesodPersistBackend site) (HandlerFor site)
--     fromDBRunner' DBRunner{..} = DBRunner' runDBRunner
    
--     toDBRunner :: forall site.
--                   DBRunner' (YesodPersistBackend site) (HandlerFor site)
--                -> DBRunner site
--     toDBRunner DBRunner'{..} = DBRunner runDBRunner'

-- fromDBRunner :: BackendCompatible backend (YesodPersistBackend site) => DBRunner site -> DBRunner' backend (HandlerFor site)
-- fromDBRunner DBRunner{..} = DBRunner' (runDBRunner . withReaderT projectBackend)

-- newtype CachedDBRunner backend m a = CachedDBRunner { runCachedDBRunnerUsing :: m (DBRunner' backend m) -> m a }
--   deriving (Functor, Applicative, Monad, MonadFix, MonadFail, Contravariant, MonadIO, Alternative, MonadPlus, MonadUnliftIO, MonadResource, MonadLogger, MonadThrow, MonadCatch, MonadMask) via (ReaderT (m (DBRunner' backend m)) m)

-- instance MonadTrans (CachedDBRunner backend) where
--   lift act = CachedDBRunner (const act)

-- instance MonadHandler m => MonadHandler (CachedDBRunner backend m) where
--   type HandlerSite (CachedDBRunner backend m) = HandlerSite m
--   type SubHandlerSite (CachedDBRunner backend m) = SubHandlerSite m

--   liftHandler = lift . liftHandler
--   liftSubHandler = lift . liftSubHandler

-- instance Monad m => WithRunDB backend m (CachedDBRunner backend m) where
--   useRunDB act = CachedDBRunner (\getRunner -> getRunner >>= \DBRunner'{..} -> runDBRunner' act)

-- runCachedDBRunnerSTM :: MonadUnliftIO m
--                      => m (DBRunner' backend m)
--                      -> CachedDBRunner backend m a
--                      -> m a
-- runCachedDBRunnerSTM doAcquire act = do
--   doAcquireLock <- newTMVarIO ()
--   runnerTMVar <- newEmptyTMVarIO

--   let getRunner = bracket (atomically $ takeTMVar doAcquireLock) (void . atomically . tryPutTMVar doAcquireLock) . const $ do
--         cachedRunner <- atomically $ tryReadTMVar runnerTMVar
--         case cachedRunner of
--           Just cachedRunner' -> return cachedRunner'
--           Nothing -> do
--             runner <- doAcquire
--             void . atomically $ tryPutTMVar runnerTMVar runner
--             return runner
--       getRunnerNoLock = maybe getRunner return =<< atomically (tryReadTMVar runnerTMVar)
  
--   runCachedDBRunnerUsing act getRunnerNoLock
