content: -- |
  -- Module      : Streamly.Internal.Data.Fold.Exception
  -- Copyright   : (c) 2025 Composewell Technologies
  -- License     : BSD-3-Clause
  -- Maintainer  : streamly@composewell.com
  -- Stability   : experimental
  -- Portability : GHC
  --
  module Streamly.Internal.Data.Fold.Exception
      (
      -- * Resources
        before
      , bracketIO
      , finallyIO
  
      -- * Exceptions
      , onException
      )
  where
  
  ------------------------------------------------------------------------------
  -- Imports
  ------------------------------------------------------------------------------
  
  import Streamly.Internal.Data.Tuple.Strict (Tuple'(..))
  import Control.Monad.IO.Class (MonadIO(..))
  import Control.Monad.Catch (MonadCatch)
  import Streamly.Internal.Data.IOFinalizer (newIOFinalizer, runIOFinalizer)
  
  import qualified Control.Monad.Catch as MC
  
  import Streamly.Internal.Data.Fold.Step
  import Streamly.Internal.Data.Fold.Type
  
  ------------------------------------------------------------------------------
  -- Exceptions
  ------------------------------------------------------------------------------
  
  {-
  
  -- | Exception handling states of a fold
  data HandleExc s f1 f2 = InitDone !s | InitFailed !f1 | StepFailed !f2
  
  -- | @handle initHandler stepHandler fold@ produces a new fold from a given
  -- fold.  The new fold executes the original @fold@, if an exception occurs
  -- when initializing the fold then @initHandler@ is executed and fold resulting
  -- from that starts execution.  If an exception occurs while executing the
  -- @step@ function of a fold then the @stephandler@ is executed and we start
  -- executing the fold resulting from that.
  --
  -- The exception is caught and handled, not rethrown. If the exception handler
  -- itself throws an exception that exception is thrown.
  --
  -- /Internal/
  --
  {-# INLINE handle #-}
  handle :: (MonadCatch m, Exception e)
      => (e -> m (Fold m a b))
      -> (e -> Fold m a b -> m (Fold m a b))
      -> Fold m a b
      -> Fold m a b
  handle initH stepH (Fold step1 initial1 extract1) = Fold step initial extract
  
      where
  
      initial = fmap InitDone initial1 `MC.catch` (fmap InitFailed . initH)
  
      step (InitDone s) a =
          let f = Fold step1 (return s) extract1
           in fmap InitDone (step1 s a)
                  `MC.catch` (\e -> fmap StepFailed (stepH e f))
      step (InitFailed (Fold step2 initial2 extract2)) a = do
          s <- initial2
          s1 <- step2 s a
          return $ InitFailed $ Fold step2 (return s1) extract2
      step (StepFailed (Fold step2 initial2 extract2)) a = do
          s <- initial2
          s1 <- step2 s a
          return $ StepFailed $ Fold step2 (return s1) extract2
  
      extract (InitDone s) = extract1 s
      extract (InitFailed (Fold _ initial2 extract2)) = initial2 >>= extract2
      extract (StepFailed (Fold _ initial2 extract2)) = initial2 >>= extract2
  
  -}
  
  -- | @onException action fold@ runs @action@ whenever the fold throws an
  -- exception.  The action is executed on any exception whether it is in
  -- initial, step or extract action of the fold.
  --
  -- The exception is not caught, simply rethrown. If the @action@ itself
  -- throws an exception that exception is thrown instead of the original
  -- exception.
  --
  -- /Internal/
  --
  {-# INLINE onException #-}
  onException :: MonadCatch m => m x -> Fold m a b -> Fold m a b
  onException action (Fold step1 initial1 extract1 final1) =
      Fold step initial extract final
  
      where
  
      initial = initial1 `MC.onException` action
      step s a = step1 s a `MC.onException` action
      extract s = extract1 s `MC.onException` action
      final s = final1 s `MC.onException` action
  
  -- | @bracketIO before after between@ runs @before@ and invokes @between@ using
  -- its output, then runs the fold generated by @between@.  If the fold ends
  -- normally, due to an exception or if it is garbage collected prematurely then
  -- @after@ is run with the output of @before@ as argument.
  --
  -- If @before@ or @after@ throw an exception that exception is thrown.
  --
  {-# INLINE bracketIO #-}
  bracketIO :: (MonadIO m, MonadCatch m)
      => IO x -> (x -> IO c) -> (x -> Fold m a b) -> Fold m a b
  bracketIO bef aft bet = Fold step initial extract final
  
      where
  
      initial = do
          r <- liftIO bef
          ref <- liftIO $ newIOFinalizer (aft r)
          case bet r of
              Fold step1 initial1 extract1 final1 -> do
                  res <- initial1 `MC.onException` liftIO (runIOFinalizer ref)
                  case res of
                      Partial s -> do
                          let fld1 = Fold step1 (pure (Partial s)) extract1 final1
                          pure $ Partial $ Tuple' ref fld1
                      Done b -> do
                          liftIO $ runIOFinalizer ref
                          pure $ Done b
  
      step (Tuple' ref (Fold step1 initial1 extract1 final1)) a = do
          res <- initial1
          case res of
              Partial s -> do
                  s1 <- step1 s a `MC.onException` liftIO (runIOFinalizer ref)
                  let fld1 = Fold step1 (pure s1) extract1 final1
                  pure $ Partial $ Tuple' ref fld1
              Done b -> do
                  liftIO $ runIOFinalizer ref
                  pure $ Done b
  
      extract (Tuple' ref (Fold _ initial1 extract1 _)) = do
          res <- initial1
          case res of
              Partial s -> extract1 s `MC.onException` liftIO (runIOFinalizer ref)
              Done b -> pure b
  
      final (Tuple' ref (Fold _ initial1 _ final1)) = do
          res <- initial1
          case res of
              Partial s -> do
                  val <- final1 s `MC.onException` liftIO (runIOFinalizer ref)
                  runIOFinalizer ref
                  pure val
              Done b -> pure b
  
  -- | Run a side effect whenever the fold stops normally, aborts due to an
  -- exception or is garbage collected.
  --
  {-# INLINE finallyIO #-}
  finallyIO :: (MonadIO m, MonadCatch m) => IO b -> Fold m a b -> Fold m a b
  finallyIO aft (Fold step1 initial1 extract1 final1) =
      Fold step initial extract final
  
      where
  
      initial = do
          ref <- liftIO $ newIOFinalizer aft
          res <- initial1 `MC.onException` liftIO (runIOFinalizer ref)
          pure $ case res of
              Done b -> Done b
              Partial s -> Partial $ Tuple' ref s
  
      step (Tuple' ref s) a = do
          res <- step1 s a `MC.onException` liftIO (runIOFinalizer ref)
          pure $ case res of
              Done b -> Done b
              Partial s1 -> Partial $ Tuple' ref s1
  
      extract (Tuple' ref s) =
          extract1 s `MC.onException` liftIO (runIOFinalizer ref)
  
      final (Tuple' ref s) = do
          res <- final1 s `MC.onException` liftIO (runIOFinalizer ref)
          liftIO $ runIOFinalizer ref
          pure res
  
  
  -- | Run a side effect before the initialization of the fold.
  --
  {-# INLINE before #-}
  before :: Monad m => m x -> Fold m a b -> Fold m a b
  before effect (Fold s i e f) = Fold s (effect *> i) e f
  