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

module UnliftIO.Async.Utils
  ( allocateAsync, allocateLinkedAsync
  , allocateAsyncWithUnmask, allocateLinkedAsyncWithUnmask
  , allocateAsyncMasked, allocateLinkedAsyncMasked
  ) where

import ClassyPrelude hiding (cancel, async, link, finally, mask)
import Control.Lens
import Control.Lens.Extras (is)

import qualified UnliftIO.Async as UnliftIO
import qualified Control.Concurrent.Async as A

import Control.Monad.Trans.Resource
import qualified Control.Monad.Trans.Resource.Internal as ResourceT.Internal
import Data.Acquire

import Control.Monad.Catch


withReference :: forall m a. (MonadUnliftIO m, MonadResource m) => ((IO (), IO ()) -> m a) -> m a
withReference act = do
  releaseAct <- newEmptyTMVarIO

  let doAlloc = do
        iSt <- liftResourceT getInternalState
        liftIO $ mask $ \_ -> do
          ResourceT.Internal.stateAlloc iSt
          atomically $ putTMVar releaseAct ()
        return iSt
      doRelease iSt eCase = liftIO . whenM (atomically $ is _Just <$> tryTakeTMVar releaseAct) $ do
        flip ResourceT.Internal.stateCleanup iSt $ case eCase of
          ExitCaseSuccess _ -> ReleaseNormal
          ExitCaseException _ -> ReleaseException
          ExitCaseAbort -> ReleaseEarly

  withRunInIO $ \run -> 
    fmap fst . generalBracket (run doAlloc) doRelease $ \iSt -> do
      res <- run $ act
        ( atomically $ takeTMVar releaseAct
        , ResourceT.Internal.stateCleanup ReleaseNormal iSt
        )
      atomically $ guard =<< isEmptyTMVar releaseAct
      return res


allocateAsync :: forall m a.
                 ( MonadUnliftIO m, MonadResource m )
              => m a -> m (Async a)
allocateAsync act = withReference $ \(signalReady, releaseRef) -> withRunInIO $ \run -> run . fmap (view _2) . flip allocate A.cancel . A.async . flip finally releaseRef $ signalReady >> run act

allocateLinkedAsync :: forall m a. (MonadUnliftIO m, MonadResource m) => m a -> m (Async a)
allocateLinkedAsync = uncurry (<$) . (id &&& UnliftIO.link) <=< allocateAsync


allocateAsyncWithUnmask :: forall m a.
                           ( MonadUnliftIO m, MonadResource m)
                        => ((forall b. m b -> m b) -> m a) -> m (Async a)
allocateAsyncWithUnmask act = withReference $ \(signalReady, releaseRef) -> withRunInIO $ \run -> run . fmap (view _2) . flip allocate A.cancel $ A.asyncWithUnmask $ \unmask -> flip finally releaseRef $ signalReady >> run (act $ liftIO . unmask . run)

allocateLinkedAsyncWithUnmask :: forall m a. (MonadUnliftIO m, MonadResource m) => ((forall b. m b -> m b) -> m a) -> m (Async a)
allocateLinkedAsyncWithUnmask act = uncurry (<$) . (id &&& UnliftIO.link) =<< allocateAsyncWithUnmask act


allocateAsyncMasked :: forall m a.
                       ( MonadUnliftIO m, MonadResource m)
                    => m a -> m (Async a)
allocateAsyncMasked act = allocateAsyncWithUnmask (const act)

allocateLinkedAsyncMasked :: forall m a. (MonadUnliftIO m, MonadResource m) => m a -> m (Async a)
allocateLinkedAsyncMasked act = allocateLinkedAsyncWithUnmask (const act)