{-# LANGUAGE CPP #-}
{-# LANGUAGE DeriveDataTypeable #-}
{-# LANGUAGE FlexibleInstances #-}
{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE MultiParamTypeClasses #-}
{-# LANGUAGE TypeFamilies #-}
{-# LANGUAGE UndecidableInstances #-}
{-# LANGUAGE RankNTypes #-}

module Control.Monad.Trans.Resource.Internal(
    ExceptionT(..)
  , InvalidAccess(..)
  , MonadResource(..)
  , MonadThrow(..)
  , MonadUnsafeIO(..)
  , ReleaseKey(..)
  , ReleaseMap(..)
  , ResIO
  , ResourceT(..)
  , stateAlloc
  , stateCleanup
  , transResourceT
  , Resource (..)
  , Allocated (..)
  , with
  , mkResource
) where

import Control.Exception (throw,Exception,SomeException)
import Control.Applicative (Applicative (..))
import Control.Monad.Trans.Control
    ( MonadTransControl (..), MonadBaseControl (..)
    , ComposeSt, defaultLiftBaseWith, defaultRestoreM, control)
import Control.Monad.Base (MonadBase, liftBase)
import Control.Monad.Trans.Cont     ( ContT  )
import Control.Monad.Cont.Class   ( MonadCont (..) )
import Control.Monad.Error.Class  ( MonadError (..) )
import Control.Monad.RWS.Class    ( MonadRWS )
import Control.Monad.Reader.Class ( MonadReader (..) )
import Control.Monad.State.Class  ( MonadState (..) )
import Control.Monad.Writer.Class ( MonadWriter (..) )

import Control.Monad.Trans.Identity ( IdentityT)
import Control.Monad.Trans.List     ( ListT    )
import Control.Monad.Trans.Maybe    ( MaybeT   )
import Control.Monad.Trans.Error    ( ErrorT, Error)
import Control.Monad.Trans.Reader   ( ReaderT  )
import Control.Monad.Trans.State    ( StateT   )
import Control.Monad.Trans.Writer   ( WriterT  )
import Control.Monad.Trans.RWS      ( RWST     )

import qualified Control.Monad.Trans.RWS.Strict    as Strict ( RWST   )
import qualified Control.Monad.Trans.State.Strict  as Strict ( StateT )
import qualified Control.Monad.Trans.Writer.Strict as Strict ( WriterT )

import Control.Monad.IO.Class (MonadIO (..))
import Control.Monad (liftM, ap)
import qualified Control.Exception as E
import Control.Monad.ST (ST)
import Data.IntMap (IntMap)
import qualified Data.IntMap as IntMap
import qualified Data.IORef as I
import Data.Monoid
import Data.Typeable
import Data.Word(Word)

#if __GLASGOW_HASKELL__ >= 704
import Control.Monad.ST.Unsafe (unsafeIOToST)
#else
import Control.Monad.ST (unsafeIOToST)
#endif

#if __GLASGOW_HASKELL__ >= 704
import qualified Control.Monad.ST.Lazy.Unsafe as LazyUnsafe
#else
import qualified Control.Monad.ST.Lazy as LazyUnsafe
#endif

import qualified Control.Monad.ST.Lazy as Lazy

import Control.Monad.Morph

-- | A @Monad@ which allows for safe resource allocation. In theory, any monad
-- transformer stack included a @ResourceT@ can be an instance of
-- @MonadResource@.
--
-- Note: @runResourceT@ has a requirement for a @MonadBaseControl IO m@ monad,
-- which allows control operations to be lifted. A @MonadResource@ does not
-- have this requirement. This means that transformers such as @ContT@ can be
-- an instance of @MonadResource@. However, the @ContT@ wrapper will need to be
-- unwrapped before calling @runResourceT@.
--
-- Since 0.3.0
class (MonadThrow m, MonadUnsafeIO m, MonadIO m, Applicative m) => MonadResource m where
    -- | Lift a @ResourceT IO@ action into the current @Monad@.
    --
    -- Since 0.4.0
    liftResourceT :: ResourceT IO a -> m a


-- | A lookup key for a specific release action. This value is returned by
-- 'register' and 'allocate', and is passed to 'release'.
--
-- Since 0.3.0
data ReleaseKey = ReleaseKey !(I.IORef ReleaseMap) !Int
    deriving Typeable

type RefCount = Word
type NextKey = Int

data ReleaseMap =
    ReleaseMap !NextKey !RefCount !(IntMap (IO ()))
  | ReleaseMapClosed

-- | Convenient alias for @ResourceT IO@.
type ResIO a = ResourceT IO a


instance MonadCont m => MonadCont (ResourceT m) where
  callCC f = ResourceT $ \i -> callCC $ \c -> unResourceT (f (ResourceT . const . c)) i

instance MonadError e m => MonadError e (ResourceT m) where
  throwError = lift . throwError
  catchError r h = ResourceT $ \i -> unResourceT r i `catchError` \e -> unResourceT (h e) i

instance MonadRWS r w s m => MonadRWS r w s (ResourceT m)

instance MonadReader r m => MonadReader r (ResourceT m) where
  ask = lift ask
  local = mapResourceT . local

mapResourceT :: (m a -> n b) -> ResourceT m a -> ResourceT n b
mapResourceT f = ResourceT . (f .) . unResourceT

instance MonadState s m => MonadState s (ResourceT m) where
  get = lift get
  put = lift . put

instance MonadWriter w m => MonadWriter w (ResourceT m) where
  tell   = lift . tell
  listen = mapResourceT listen
  pass   = mapResourceT pass

-- | A @Monad@ which can throw exceptions. Note that this does not work in a
-- vanilla @ST@ or @Identity@ monad. Instead, you should use the 'ExceptionT'
-- transformer in your stack if you are dealing with a non-@IO@ base monad.
--
-- Since 0.3.0
class Monad m => MonadThrow m where
    monadThrow :: E.Exception e => e -> m a

instance MonadThrow IO where
    monadThrow = E.throwIO

instance MonadThrow Maybe where
    monadThrow _ = Nothing
instance MonadThrow (Either SomeException) where
    monadThrow = Left . E.toException
instance MonadThrow [] where
    monadThrow _ = []

#define GO(T) instance (MonadThrow m) => MonadThrow (T m) where monadThrow = lift . monadThrow
#define GOX(X, T) instance (X, MonadThrow m) => MonadThrow (T m) where monadThrow = lift . monadThrow
GO(IdentityT)
GO(ListT)
GO(MaybeT)
GOX(Error e, ErrorT e)
GO(ReaderT r)
GO(ContT r)
GO(ResourceT)
GO(StateT s)
GOX(Monoid w, WriterT w)
GOX(Monoid w, RWST r w s)
GOX(Monoid w, Strict.RWST r w s)
GO(Strict.StateT s)
GOX(Monoid w, Strict.WriterT w)
#undef GO
#undef GOX

instance (MonadThrow m, MonadUnsafeIO m, MonadIO m, Applicative m) => MonadResource (ResourceT m) where
    liftResourceT = transResourceT liftIO

-- | Transform the monad a @ResourceT@ lives in. This is most often used to
-- strip or add new transformers to a stack, e.g. to run a @ReaderT@.
--
-- Note that this function is a slight generalization of 'hoist'.
--
-- Since 0.3.0
transResourceT :: (m a -> n b)
               -> ResourceT m a
               -> ResourceT n b
transResourceT f (ResourceT mx) = ResourceT (\r -> f (mx r))

-- | Since 0.4.7
instance MFunctor ResourceT where
    hoist f (ResourceT mx) = ResourceT (\r -> f (mx r))
-- | Since 0.4.7
instance MMonad ResourceT where
    embed f m = ResourceT (\i -> unResourceT (f (unResourceT m i)) i)

-- | The Resource transformer. This transformer keeps track of all registered
-- actions, and calls them upon exit (via 'runResourceT'). Actions may be
-- registered via 'register', or resources may be allocated atomically via
-- 'allocate'. @allocate@ corresponds closely to @bracket@.
--
-- Releasing may be performed before exit via the 'release' function. This is a
-- highly recommended optimization, as it will ensure that scarce resources are
-- freed early. Note that calling @release@ will deregister the action, so that
-- a release action will only ever be called once.
--
-- Since 0.3.0
newtype ResourceT m a = ResourceT { unResourceT :: I.IORef ReleaseMap -> m a }
#if __GLASGOW_HASKELL__ >= 707
        deriving Typeable
#else
instance Typeable1 m => Typeable1 (ResourceT m) where
    typeOf1 = goType undefined
      where
        goType :: Typeable1 m => m a -> ResourceT m a -> TypeRep
        goType m _ =
            mkTyConApp
#if __GLASGOW_HASKELL__ >= 704
                (mkTyCon3 "resourcet" "Control.Monad.Trans.Resource" "ResourceT")
#else
                (mkTyCon "Control.Monad.Trans.Resource.ResourceT")
#endif
                [ typeOf1 m
                ]
#endif

-- | Indicates either an error in the library, or misuse of it (e.g., a
-- @ResourceT@'s state is accessed after being released).
--
-- Since 0.3.0
data InvalidAccess = InvalidAccess { functionName :: String }
    deriving Typeable

instance Show InvalidAccess where
    show (InvalidAccess f) = concat
        [ "Control.Monad.Trans.Resource."
        , f
        , ": The mutable state is being accessed after cleanup. Please contact the maintainers."
        ]

instance Exception InvalidAccess

-------- All of our monad et al instances
instance Functor m => Functor (ResourceT m) where
    fmap f (ResourceT m) = ResourceT $ \r -> fmap f (m r)

instance Applicative m => Applicative (ResourceT m) where
    pure = ResourceT . const . pure
    ResourceT mf <*> ResourceT ma = ResourceT $ \r ->
        mf r <*> ma r

instance Monad m => Monad (ResourceT m) where
    return = ResourceT . const . return
    ResourceT ma >>= f = ResourceT $ \r -> do
        a <- ma r
        let ResourceT f' = f a
        f' r

instance MonadTrans ResourceT where
    lift = ResourceT . const

instance MonadIO m => MonadIO (ResourceT m) where
    liftIO = lift . liftIO

instance MonadBase b m => MonadBase b (ResourceT m) where
    liftBase = lift . liftBase

instance MonadTransControl ResourceT where
    newtype StT ResourceT a = StReader {unStReader :: a}
    liftWith f = ResourceT $ \r -> f $ \(ResourceT t) -> liftM StReader $ t r
    restoreT = ResourceT . const . liftM unStReader
    {-# INLINE liftWith #-}
    {-# INLINE restoreT #-}

instance MonadBaseControl b m => MonadBaseControl b (ResourceT m) where
     newtype StM (ResourceT m) a = StMT (StM m a)
     liftBaseWith f = ResourceT $ \reader' ->
         liftBaseWith $ \runInBase ->
             f $ liftM StMT . runInBase . (\(ResourceT r) -> r reader'  )
     restoreM (StMT base) = ResourceT $ const $ restoreM base
instance Monad m => MonadThrow (ExceptionT m) where
    monadThrow = ExceptionT . return . Left . E.toException
instance MonadResource m => MonadResource (ExceptionT m) where
    liftResourceT = lift . liftResourceT
instance MonadIO m => MonadIO (ExceptionT m) where
    liftIO = lift . liftIO

#define GO(T) instance (MonadResource m) => MonadResource (T m) where liftResourceT = lift . liftResourceT
#define GOX(X, T) instance (X, MonadResource m) => MonadResource (T m) where liftResourceT = lift . liftResourceT
GO(IdentityT)
GO(ListT)
GO(MaybeT)
GOX(Error e, ErrorT e)
GO(ReaderT r)
GO(ContT r)
GO(StateT s)
GOX(Monoid w, WriterT w)
GOX(Monoid w, RWST r w s)
GOX(Monoid w, Strict.RWST r w s)
GO(Strict.StateT s)
GOX(Monoid w, Strict.WriterT w)
#undef GO
#undef GOX


-- | The express purpose of this transformer is to allow non-@IO@-based monad
-- stacks to catch exceptions via the 'MonadThrow' typeclass.
--
-- Since 0.3.0
newtype ExceptionT m a = ExceptionT { runExceptionT :: m (Either SomeException a) }

stateAlloc :: I.IORef ReleaseMap -> IO ()
stateAlloc istate = do
    I.atomicModifyIORef istate $ \rm ->
        case rm of
            ReleaseMap nk rf m ->
                (ReleaseMap nk (rf + 1) m, ())
            ReleaseMapClosed -> throw $ InvalidAccess "stateAlloc"

stateCleanup :: I.IORef ReleaseMap -> IO ()
stateCleanup istate = E.mask_ $ do
    mm <- I.atomicModifyIORef istate $ \rm ->
        case rm of
            ReleaseMap nk rf m ->
                let rf' = rf - 1
                 in if rf' == minBound
                        then (ReleaseMapClosed, Just m)
                        else (ReleaseMap nk rf' m, Nothing)
            ReleaseMapClosed -> throw $ InvalidAccess "stateCleanup"
    case mm of
        Just m ->
            mapM_ (\x -> try x >> return ()) $ IntMap.elems m
        Nothing -> return ()
  where
    try :: IO a -> IO (Either SomeException a)
    try = E.try


-- | A @Monad@ based on some monad which allows running of some 'IO' actions,
-- via unsafe calls. This applies to 'IO' and 'ST', for instance.
--
-- Since 0.3.0
class Monad m => MonadUnsafeIO m where
    unsafeLiftIO :: IO a -> m a

instance MonadUnsafeIO IO where
    unsafeLiftIO = id

instance MonadUnsafeIO (ST s) where
    unsafeLiftIO = unsafeIOToST

instance MonadUnsafeIO (Lazy.ST s) where
    unsafeLiftIO = LazyUnsafe.unsafeIOToST

instance (MonadTrans t, MonadUnsafeIO m, Monad (t m)) => MonadUnsafeIO (t m) where
    unsafeLiftIO = lift . unsafeLiftIO

instance Monad m => Functor (ExceptionT m) where
    fmap f = ExceptionT . (liftM . fmap) f . runExceptionT
instance Monad m => Applicative (ExceptionT m) where
    pure = ExceptionT . return . Right
    ExceptionT mf <*> ExceptionT ma = ExceptionT $ do
        ef <- mf
        case ef of
            Left e -> return (Left e)
            Right f -> do
                ea <- ma
                case ea of
                    Left e -> return (Left e)
                    Right x -> return (Right (f x))
instance Monad m => Monad (ExceptionT m) where
    return = pure
    ExceptionT ma >>= f = ExceptionT $ do
        ea <- ma
        case ea of
            Left e -> return (Left e)
            Right a -> runExceptionT (f a)
instance MonadBase b m => MonadBase b (ExceptionT m) where
    liftBase = lift . liftBase
instance MonadTrans ExceptionT where
    lift = ExceptionT . liftM Right
instance MonadTransControl ExceptionT where
    newtype StT ExceptionT a = StExc { unStExc :: Either SomeException a }
    liftWith f = ExceptionT $ liftM return $ f $ liftM StExc . runExceptionT
    restoreT = ExceptionT . liftM unStExc
instance MonadBaseControl b m => MonadBaseControl b (ExceptionT m) where
    newtype StM (ExceptionT m) a = StE { unStE :: ComposeSt ExceptionT m a }
    liftBaseWith = defaultLiftBaseWith StE
    restoreM = defaultRestoreM unStE

instance MonadCont m => MonadCont (ExceptionT m) where
  callCC f = ExceptionT $
    callCC $ \c ->
    runExceptionT (f (\a -> ExceptionT $ c (Right a)))

instance MonadError e m => MonadError e (ExceptionT m) where
  throwError = lift . throwError
  catchError r h = ExceptionT $ runExceptionT r `catchError` (runExceptionT . h)

instance MonadRWS r w s m => MonadRWS r w s (ExceptionT m)

instance MonadReader r m => MonadReader r (ExceptionT m) where
  ask = lift ask
  local = mapExceptionT . local

mapExceptionT :: (m (Either SomeException a) -> n (Either SomeException b)) -> ExceptionT m a -> ExceptionT n b
mapExceptionT f = ExceptionT . f . runExceptionT

instance MonadState s m => MonadState s (ExceptionT m) where
  get = lift get
  put = lift . put

instance MonadWriter w m => MonadWriter w (ExceptionT m) where
  tell   = lift . tell
  listen = mapExceptionT $ \ m -> do
    (a, w) <- listen m
    return $! fmap (\ r -> (r, w)) a
  pass   = mapExceptionT $ \ m -> pass $ do
    a <- m
    return $! case a of
        Left  l      -> (Left  l, id)
        Right (r, f) -> (Right r, f)

data Allocated a = Allocated !a !(IO ())

-- | A method for allocating a scarce resource, providing the means of freeing
-- it when no longer needed. This data type provides
-- @Functor@/@Applicative@/@Monad@ instances for composing different resources
-- together. You can allocate these resources using either the @bracket@
-- pattern (via @with@) or using @ResourceT@ (via @allocateResource@).
--
-- This concept was originally introduced by Gabriel Gonzalez and described at:
-- <http://www.haskellforall.com/2013/06/the-resource-applicative.html>. The
-- implementation in this package is slightly different, due to taking a
-- different approach to async exception safety.
--
-- Since 0.4.10
newtype Resource a = Resource ((forall b. IO b -> IO b) -> IO (Allocated a))
    deriving Typeable

instance Functor Resource where
    fmap = liftM
instance Applicative Resource where
    pure = return
    (<*>) = ap

instance Monad Resource where
    return a = Resource (\_ -> return (Allocated a (return ())))
    Resource f >>= g' = Resource $ \restore -> do
        Allocated x free1 <- f restore
        let Resource g = g' x
        Allocated y free2 <- g restore `E.onException` free1
        return $! Allocated y (free2 `E.finally` free1)

instance MonadIO Resource where
    liftIO f = Resource $ \restore -> do
        x <- restore f
        return $! Allocated x (return ())

instance MonadBase IO Resource where
    liftBase = liftIO

-- | Create a @Resource@ value using the given allocate and free functions.
--
-- Since 0.4.10
mkResource :: IO a -- ^ allocate the resource
           -> (a -> IO ()) -- ^ free the resource
           -> Resource a
mkResource create free = Resource $ \restore -> do
    x <- restore create
    return $! Allocated x (free x)

-- | Allocate the given resource and provide it to the provided function. The
-- resource will be freed as soon as the inner block is exited, whether
-- normally or via an exception. This function is similar in function to
-- @bracket@.
--
-- Since 0.4.10
with :: MonadBaseControl IO m
     => Resource a
     -> (a -> m b)
     -> m b
with (Resource f) g = control $ \run -> E.mask $ \restore -> do
    Allocated x free <- f restore
    run (g x) `E.finally` free