{-# LANGUAGE UndecidableInstances #-}

-- | Internal module for 'Scoped', 'ScopedResource' & co.
--
-- Only import this if you need to wrap an otherwise unsafe interface around resources
--
-- @since 0.1.0.0
module Control.Monad.Scoped.Internal
  ( -- * definitions of 'Scoped' and 'ScopedResource' and functions to work with them
    Scoped (..)
  , ScopedResource (..)
  , Scoping (..)
  , registerHandler
  , (:<)

    -- ** Helpers to create your own 'Scoped' wrappers around resources
  , bracketScoped
  )
where

import Control.Applicative (Alternative (empty, (<|>)))
import Control.Exception qualified
import Control.Monad (MonadPlus (mplus), ap, mzero)
import Control.Monad.IO.Unlift (MonadIO (liftIO), MonadUnliftIO (..))
import Control.Monad.Trans.Class (MonadTrans (lift))
import Data.Kind (Constraint, Type)
import Data.Type.Equality (type (~~))
import GHC.Exts (RuntimeRep, TYPE)

type role Scoped nominal representational representational

-- | The 'Scoped' monad that provides the possibility to safely scope the allocation of a resource
--
-- It is used to abstract over all of the CPS style withSomething functions, like 'System.IO.withFile'
--
-- Be sure to properly mask handlers if you are using 'UnsafeMkScoped'. Use safe helper functions like
-- 'registerHandler' or 'bracketScoped' where possible.
--
-- Scoped also works for wrapping unboxed and unlifted monad transformers.
--
-- @since 0.1.0.0
type Scoped :: forall {k} {rep :: RuntimeRep}. [Type] -> (k -> TYPE rep) -> Type -> Type
newtype Scoped s m a = UnsafeMkScoped
  { unsafeRunScoped :: forall b. (a -> m b) -> m b
  -- ^ Unsafely runs a scoped block. Use 'scoped' instead, otherwise resources might escape
  }

type role ScopedResource nominal representational

-- | A scoped resource with token @s@ belonging to a 'Scoped' block with the same token.
--
-- If you are creating a 'ScopedResource', make sure the resource is deallocated properly
-- when the 'Scoped' block is exited.
--
-- @since 0.1.0.0
type ScopedResource :: Type -> Type -> Type
newtype ScopedResource s a = UnsafeMkScopedResource
  { unsafeUnwrapScopedResource :: a
  -- ^ Unsafely runs a scoped resource. It forgets the scope of the resource and may now be escaped incorrectly
  }
  deriving stock (Eq, Ord, Show)

-- | when using a resource, all that matters is that the resource can only be used in the scope that it was created in
--   or any scope that is farther in than that scope
--
--   This constraint has to be put to connect the resource and the scope that it was created in
--
-- @since 0.1.0.0
type (:<) :: Type -> [Type] -> Constraint
class s :< ss

instance s :< '[s]

instance {-# INCOHERENT #-} s :< (s : s'' : ss)

instance s :< (s'' : ss) => s :< (s' : s'' : ss)

-- | the 'Scoping' class is there to give overloading to blocks, s.t. we don't have to run different functions
-- depending on whether we run a final block or not.
--
-- This type class is internal since there should not be any more instances and since it is expected that the contraint
-- on 'scoped' is immediately discharged
--
-- @since 0.1.0.0
type Scoping :: [Type] -> (k -> TYPE r) -> (Type -> Type) -> Constraint
class Scoping ss m n | n -> m ss where
  -- | Run a 'Scoped' block safely, making sure that none of the safely allocated resources can escape it, using
  -- the same trick as 'Control.Monad.ST.ST'
  --
  -- All of the allocated resources will live until the end of the block is reached
  --
  -- @since 0.1.0.0
  scoped :: (forall s. Scoped (s : ss) m a) -> n a

instance (Applicative m, m ~ n, l ~ '[]) => Scoping l m n where
  scoped act = unsafeRunScoped act pure

instance {-# OVERLAPPING #-} Monad m => Scoping (s : ss) m (Scoped (s : ss) m) where
  scoped act = lift (unsafeRunScoped act pure)

-- | Run a handler masked for async exception when the 'Scoped' block ends
--
-- You can register a handler wherever in your 'Scoped' block you want, but it will nonetheless be run
-- in reverse order that the handlers have been registered, after the scoped block's actions have been finished
--
-- Mind that this uses 'Control.Exception.finally' under the hood and thus does not mask the handler with an uninterruptible mask
--
-- @since 0.1.0.0
registerHandler
  :: MonadUnliftIO m
  => m a
  -- ^ the handler to be registered
  -> Scoped ss m ()
registerHandler hdl = UnsafeMkScoped \k -> k () `finally` hdl

finally :: MonadUnliftIO m => m a -> m b -> m a
finally act hdl = withRunInIO \inIO -> Control.Exception.finally (inIO act) (inIO hdl)
{-# INLINE finally #-}

-- | A wrapper around 'Control.Exception.bracket' to allocate a resource safely in a 'Scoped' block
--
-- It returns a 'ScopedResource' that belongs to the 'Scoped' block it was allocated in
--
-- Note that this uses 'Control.Exception.bracket' internally and thus uses an interruptible mask for the cleanup continuation
--
-- @since 0.1.0.0
bracketScoped
  :: MonadUnliftIO m
  => m a
  -- ^ an action that allocates a resource of type @a@
  -> (a -> m b)
  -- ^ an action that deallocates a resource of type @a@
  -> Scoped (s : ss) m (ScopedResource s a)
bracketScoped act kfail = UnsafeMkScoped \k -> bracket act kfail (k . UnsafeMkScopedResource)

bracket :: MonadUnliftIO m => m a -> (a -> m b) -> (a -> m c) -> m c
bracket act yes no = withRunInIO \inIO -> Control.Exception.bracket (inIO act) (inIO . yes) (inIO . no)
{-# INLINE bracket #-}

-- | @since 0.1.0.0
instance Functor (Scoped s m) where
  fmap f (UnsafeMkScoped m) = UnsafeMkScoped \k -> m (k . f)
  {-# INLINE fmap #-}

-- | @since 0.1.0.0
instance Applicative (Scoped s m) where
  pure a = UnsafeMkScoped \k -> k a
  {-# INLINE pure #-}
  (<*>) = ap
  {-# INLINE (<*>) #-}

-- | @since 0.1.0.0
instance (Alternative m', m' ~~ m) => Alternative (Scoped s m) where
  empty = UnsafeMkScoped (const empty)
  {-# INLINE empty #-}
  UnsafeMkScoped a <|> UnsafeMkScoped b = UnsafeMkScoped (\k -> a k <|> b k)
  {-# INLINE (<|>) #-}

-- | @since 0.1.0.0
instance Monad (Scoped s m) where
  UnsafeMkScoped m >>= k = UnsafeMkScoped \k' -> m \a -> unsafeRunScoped (k a) k'
  {-# INLINE (>>=) #-}

-- | @since 0.1.0.0
instance (Alternative m', m' ~~ m) => MonadPlus (Scoped s m) where
  mplus = (<|>)
  {-# INLINE mplus #-}
  mzero = empty
  {-# INLINE mzero #-}

-- | You can perform 'IO' in a scoped block, but it does not inherit its safety guarantees
--
-- @since 0.1.0.0
instance (MonadIO m', m' ~~ m) => MonadIO (Scoped s m) where
  liftIO = lift . liftIO
  {-# INLINE liftIO #-}

-- | You can 'fail' in a Scoped block
--
-- @since 0.1.0.0
instance (MonadFail m', m' ~~ m) => MonadFail (Scoped s m) where
  fail = lift . fail
  {-# INLINE fail #-}

-- | You can use all the actions you can use in the underlying monad @m@ also in the 'Scoped' monad by 'lift'ing into it.
--
-- @since 0.1.0.0
instance MonadTrans (Scoped s) where
  lift m = UnsafeMkScoped (m >>=)
  {-# INLINE lift #-}