{-# LANGUAGE DerivingStrategies #-}
{-# LANGUAGE GeneralizedNewtypeDeriving #-}
{-# LANGUAGE RankNTypes #-}
{-|
Module      : Utils.Monad
Copyright   : (c) UU, 2019
License     : MIT
Maintainer  : Tom Smeding
Stability   : experimental
Portability : POSIX, macOS, Windows

Some useful functions lifted to monadic actions.
-}
module Utils.Monad where

import Control.Monad.IO.Unlift
import Control.Monad.State.Strict (StateT (..), lift, get, put, MonadState, MonadTrans)
import Data.IORef
import Control.Exception (bracket, bracket_)


-- | If the boolean resolves to True, returns the first argument, else the
-- second.
ifM :: Monad m => m Bool -> m a -> m a -> m a
ifM b t e = b >>= \b' -> if b' then t else e

-- | Runs the monadic action only if the boolean resolves to True.
whenM :: Monad m => m Bool -> m () -> m ()
whenM b t = ifM b t (return ())

-- | whenJust m a = when (isJust m) (a (fromJust m)), but with less
-- fromJust.
whenJust :: Monad m => Maybe a -> (a -> m ()) -> m ()
whenJust = flip (maybe (return ()))

-- | 'bracket', but unlifted.
unliftBracket :: MonadUnliftIO m => m a -> (a -> m b) -> (a -> m b) -> m b
unliftBracket f g h =
    withRunInIO $ \run ->
        bracket (run f) (run . g) (run . h)

-- | 'bracket_', but unlifted.
unliftBracket_ :: MonadUnliftIO m => m a -> m b -> m c -> m c
unliftBracket_ f g h =
    withRunInIO $ \run ->
        bracket_ (run f) (run g) (run h)

newtype StateTUnliftIOWrapper s m a = StateTUnliftIOWrapper { runStateTUnliftIOWrapper :: StateT s m a }
  deriving newtype (Functor, Applicative, Monad, MonadIO, MonadState s, MonadTrans)

instance MonadUnliftIO m => MonadUnliftIO (StateTUnliftIOWrapper s m) where
    withRunInIO act = StateTUnliftIOWrapper $ do
        state0 <- get
        ref <- liftIO $ newIORef state0
        let runner mrun (StateT mf) = do
                s1 <- readIORef ref
                (res, s2) <- runStateT (StateT $ mrun . mf) s1
                writeIORef ref s2
                return res
        res <- lift $ withRunInIO $ \mrun -> act (runner mrun . runStateTUnliftIOWrapper)
        liftIO (readIORef ref) >>= put
        return res