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

-- also see Utils.Persist

module Utils.DB where

import ClassyPrelude.Yesod hiding (addMessageI)

import qualified Data.Monoid as Monoid (First())
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 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) -- , updateWhereCount)

import GHC.Stack (HasCallStack, CallStack, callStack)

-- import Language.Haskell.TH.Lift

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

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

-- | Obtain a record projection from an EntityField
getFieldEnt  :: PersistEntity record => EntityField record typ -> Entity record -> typ
getFieldEnt = view . fieldLens

getField     :: PersistEntity record => EntityField record typ -> record -> typ
getField = view . fieldLensVal

-- | Obtain a lens from an EntityField
fieldLensVal :: PersistEntity record => EntityField record typ -> Lens' record typ
fieldLensVal f = entityLens . fieldLens f
  where
    entityLens :: Lens' record (Entity record)
    entityLens = lens getVal setVal
    getVal :: record -> Entity record
    getVal   = Entity (error "fieldLensVal unexpectectly required an entity key") -- this is safe, since the lens is only used locally
    setVal :: record -> Entity record -> record
    setVal _ = entityVal


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

notExists :: (PersistRecordBackend record backend, PersistQueryRead backend, MonadIO m) => [Filter record] -> ReaderT backend m Bool
notExists = fmap not . exists

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

-- | given filter criteria like `selectList` this function returns Just if and only if there is precisely one result
--   getByPeseudoUnique
getByFilter ::  (PersistRecordBackend record backend, PersistQueryRead backend, MonadIO m)
            =>  [Filter record] -> ReaderT backend m (Maybe (Entity record))
getByFilter crit =
  selectList crit [LimitTo 2] <&> \case
    [singleEntity] -> Just singleEntity
    _              -> Nothing -- not existing or not unique

getKeyByFilter ::  (PersistRecordBackend record backend, PersistQueryRead backend, MonadIO m)
            =>  [Filter record] -> ReaderT backend m (Maybe (Key record))
getKeyByFilter crit =
  selectKeysList crit [LimitTo 2] <&> \case
    [singleKey] -> Just singleKey
    _           -> Nothing -- not existing or not unique


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

-- | update and retrieve an entity. Will throw an error if the key is updaded
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

-- Notes on upsertBy:
--  * Unique denotes old record
--  * Changes to fields involved in uniqueness work, but may throw an error if updated record already exists
--  * Use Database.Esqueleto.PostgreSQL.upsertBy for more elaborate conflict updates

-- | Safe version of upsertBy which does nothing if the new or updated record would violate a uniqueness constraint
upsertBySafe :: ( MonadIO m
                , PersistEntity record
                , PersistUniqueWrite backend
                , PersistEntityBackend record ~ BaseBackend backend
                )
             => Unique record -> record -> (record -> record) -> ReaderT backend m (Maybe (Key record))
upsertBySafe uniq newr upd = maybeM (insertUnique newr) do_upd (getBy uniq)
  where
    do_upd Entity{entityKey = oid, entityVal = oldr} = do
      delete oid
      insertUnique $ upd oldr

upsertBy_ :: ( MonadIO m
             , PersistEntity record
             , PersistUniqueWrite backend
             , PersistEntityBackend record ~ BaseBackend backend
            )
          => Unique record -> record -> [Update record] -> ReaderT backend m ()
upsertBy_ = ((void .) .) . upsertBy

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)

-- Backport from version persistent-2.14.6.3
insertUnique_ :: ( MonadIO m
             , PersistEntity record
             , PersistUniqueWrite backend
             , PersistEntityBackend record ~ BaseBackend backend
            )
    => record -> ReaderT backend m (Maybe ())
insertUnique_ datum = do
    conflict <- checkUnique datum
    case conflict of
        Nothing -> Just <$> insert_ datum
        Just _ -> return Nothing

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

-- | Deprecated, use selectFirst instead.
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

-- Could be used at Handler.Admin.postAdminProblemsR, but not yet elsewhere, thus inlined for now, as it may be too special:
-- updateWithMessage
--   :: ( YesodPersist site, PersistEntity val, BackendCompatible SqlBackend (YesodPersistBackend site), PersistEntityBackend val ~ SqlBackend
--       , Num a, Ord a, RenderMessage site msg, RedirectUrl site (url,[(Text,Text)]))
--   => url                 -- where to redirect, if changes were mage
--   -> [Filter val]        -- update filter
--   -> [Update val]        -- actual update
--   -> a                   -- expected updates
--   -> (a -> msg)          -- message to add with number of actual changes
--   -> HandlerFor site ()
-- updateWithMessage route flt upd no_req msg = do
--   (fromIntegral -> oks) <- runDB $ updateWhereCount flt upd
--   let mkind = if oks < no_req || no_req <= 0 then Warning else Success
--   addMessageI mkind $ msg oks
--   when (oks > 0) $ do -- reload to ensure updates are displayed
--     getps <- reqGetParams <$> getRequest
--     redirect (route, getps)


-- 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



-- A datatype for a specific heterogeneous list to compute DB updates, consisting of a persistent record field and a fitting lens
data CheckUpdate record iraw =
    forall typ. (Eq typ, PersistField typ) =>
    CheckUpdate     (EntityField record typ) (Getting               typ  iraw typ) -- A persistent record field and fitting getting (also use for typ ~ Maybe typ')
  | forall typ. (Eq typ, PersistField typ) =>
    CheckUpdateMay  (EntityField record (Maybe typ)) (Getting (Maybe typ) iraw (Maybe typ)) -- Special case, when `typ` is optional everywhere, forces update of Nothing to Just values
  | forall typ. (Eq typ, PersistField typ) =>
    CheckUpdateOpt  (EntityField record typ) (Getting (Monoid.First typ) iraw typ) -- Special case, when `typ` is optional for the lens, but not optional in DB.

-- deriving instance Lift (CheckUpdate record iraw) -- not possible, seee Handler.Utils.AvsUpdate for a workaround
-- instance Lift (CheckUpdate record iraw) where
--   lift = $(makeLift ''CheckUpdate)

-- | checks if an update would be performed, if a new different value would be presented. Should agree with `mkUpdate` familiy of functions
mayUpdate :: PersistEntity record => record -> Maybe iraw -> CheckUpdate record iraw -> Bool
mayUpdate ent (Just old) (CheckUpdate up l)
  | let oldval = old ^. l
  , let entval = ent ^. fieldLensVal up
    = oldval == entval
mayUpdate ent (Just old) (CheckUpdateMay up l)
  | let oldval = old ^. l
  , let entval = ent ^. fieldLensVal up
    = isNothing entval || oldval == entval
mayUpdate ent (Just old) (CheckUpdateOpt up l)
  | Just oldval <- old ^? l
  , let entval = ent ^. fieldLensVal up
    = oldval == entval
mayUpdate _ _ _ = False

-- | Compute necessary updates. Given a database record, the new and old raw data, and a pair consisting of a getter from raw data to a value and an EntityField of the same value,
--   an update is returned, if the current value is identical to the old value, which changed in the new raw data
mkUpdate :: PersistEntity record => record -> iraw -> Maybe iraw -> CheckUpdate record iraw -> Maybe (Update record)
mkUpdate ent new (Just old) (CheckUpdate up l)
  | let newval = new ^. l
  , let oldval = old ^. l
  , let entval = ent ^. fieldLensVal up
  , newval /= entval
  , oldval == entval
    = Just (up =. newval)
mkUpdate ent new (Just old) (CheckUpdateMay up l)
  | let newval = new ^. l
  , let oldval = old ^. l
  , let entval = ent ^. fieldLensVal up
  , (isNothing entval && isJust newval) || (newval /= entval && oldval == entval)
    = Just (up =. newval)
mkUpdate ent new (Just old) (CheckUpdateOpt up l)
  | Just newval <- new ^? l
  , Just oldval <- old ^? l
  , let entval = ent ^. fieldLensVal up
  , newval /= entval
  , oldval == entval
    = Just (up =. newval)
mkUpdate _ _ _ _ = Nothing

-- | Like `mkUpdate` but performs the update even if there was no old value to check if the value had been edited
mkUpdate' :: PersistEntity record => record -> iraw -> Maybe iraw -> CheckUpdate record iraw -> Maybe (Update record)
mkUpdate' ent new Nothing = mkUpdateDirect ent new
mkUpdate' ent new just    = mkUpdate       ent new just

-- | Like `mkUpdate` but performs the update without comparison to a previous older value, whenever current entity value and new value are different
mkUpdateDirect :: PersistEntity record => record -> iraw -> CheckUpdate record iraw -> Maybe (Update record)
mkUpdateDirect ent new (CheckUpdate up l)
  | let newval = new ^. l
  , let entval = ent ^. fieldLensVal up
  , newval /= entval
    = Just (up =. newval)
mkUpdateDirect ent new (CheckUpdateMay up l)
  | let newval = new ^. l
  , let entval = ent ^. fieldLensVal up
  , newval /= entval
    = Just (up =. newval)
mkUpdateDirect ent new (CheckUpdateOpt up l)
  | Just newval <- new ^? l
  , let entval = ent ^. fieldLensVal up
  , newval /= entval
    = Just (up =. newval)
mkUpdateDirect _ _ _ = Nothing

-- | Unconditionally update a record through CheckUpdate
updateRecord :: PersistEntity record => record -> iraw -> CheckUpdate record iraw -> record
updateRecord ent new (CheckUpdate up l) =
  let newval  = new ^. l
      lensRec = fieldLensVal up
  in  ent & lensRec .~ newval
updateRecord ent new (CheckUpdateMay up l) =
  let newval  = new ^. l
      lensRec = fieldLensVal up
  in  ent & lensRec .~ newval
updateRecord ent new (CheckUpdateOpt up l)
  | Just newval <- new ^? l
    = ent & fieldLensVal up .~ newval
  | otherwise
    = ent

-- | like mkUpdate' but only returns the update if the new value would be unique
-- mkUpdateCheckUnique' :: PersistEntity record => record -> iraw -> Maybe iraw -> CheckUpdate record iraw -> DB (Maybe (Update record))
mkUpdateCheckUnique' :: (MonadIO m, PersistQueryRead backend, PersistEntity record, PersistEntityBackend record ~ BaseBackend backend)
                    => record -> a -> Maybe a -> CheckUpdate record a -> ReaderT backend m (Maybe (Update record))
mkUpdateCheckUnique' ent new Nothing (CheckUpdate up l)
  | let newval = new ^. l
  , let entval = ent ^. fieldLensVal up
  , newval /= entval
    = do
        newval_exists <- exists [up ==. newval]
        return $ toMaybe (not newval_exists) (up =. newval)
mkUpdateCheckUnique' ent new Nothing (CheckUpdateMay up l)
  | let newval = new ^. l
  , let entval = ent ^. fieldLensVal up
  , newval /= entval
    = do
        newval_exists <- exists [up ==. newval]
        return $ toMaybe (not newval_exists) (up =. newval)
mkUpdateCheckUnique' ent new Nothing (CheckUpdateOpt up l)
  | Just newval <- new ^? l
  , let entval = ent ^. fieldLensVal up
  , newval /= entval
    = do
        newval_exists <- exists [up ==. newval]
        return $ toMaybe (not newval_exists) (up =. newval)
mkUpdateCheckUnique' ent new (Just old) (CheckUpdate up l)
  | let newval = new ^. l
  , let oldval = old ^. l
  , let entval = ent ^. fieldLensVal up
  , newval /= entval
  , oldval == entval
    = do
        newval_exists <- exists [up ==. newval]
        return $ toMaybe (not newval_exists) (up =. newval)
mkUpdateCheckUnique' ent new (Just old) (CheckUpdateMay up l)
  | let newval = new ^. l
  , let oldval = old ^. l
  , let entval = ent ^. fieldLensVal up
  , (isNothing entval && isJust newval) || (newval /= entval && oldval == entval)
    = do
        newval_exists <- exists [up ==. newval]
        return $ toMaybe (not newval_exists) (up =. newval)
mkUpdateCheckUnique' ent new (Just old) (CheckUpdateOpt up l)
  | Just newval <- new ^? l
  , Just oldval <- old ^? l
  , let entval = ent ^. fieldLensVal up
  , newval /= entval
  , oldval == entval
    = do
        newval_exists <- exists [up ==. newval]
        return $ toMaybe (not newval_exists) (up =. newval)
mkUpdateCheckUnique' _ _ _ _ = return Nothing
