-- SPDX-FileCopyrightText: 2021 Oxhead Alpha
-- SPDX-License-Identifier: LicenseRef-MIT-OA

-- | This module defines common exception types and exception-related functions
-- used throughout the "Test.Cleveland" modules.
module Test.Cleveland.Internal.Exceptions
  ( WithCallStack(..)
  , addCallStack
  , throwWithCallStack
  , tryWithCallStack
  , catchWithCallStack
  ) where

import Data.Typeable (cast)
import Fmt (Buildable(..), pretty, unlinesF)

----------------------------------------------------------------------------
-- WithCallStack
----------------------------------------------------------------------------

-- | Wraps an exception and adds some callstack information.
data WithCallStack where
  WithCallStack :: CallStack -> SomeException -> WithCallStack

deriving stock instance Show WithCallStack

instance Buildable WithCallStack where
  build (WithCallStack cstack e) = unlinesF
    [ prettyCallStack cstack
    , displayException e
    ]

instance Exception WithCallStack where
  displayException = pretty

-- | Wrap any exceptions thrown by some monadic action with 'WithCallStack'.
--
-- If the action throws an exception already wrapped in 'WithCallStack',
-- that callstack will be overriden with the current one.
addCallStack :: forall m a. (HasCallStack, MonadCatch m) => m a -> m a
addCallStack ma =
  ma `catch` \se@(SomeException ex) ->
    case fromException @WithCallStack se of
      Just (WithCallStack _ innerEx) -> throwWithCallStack callStack innerEx
      Nothing -> throwWithCallStack callStack ex

throwWithCallStack :: forall e a m. (MonadThrow m, Exception e) => CallStack -> e -> m a
throwWithCallStack cstack ex = throwM (WithCallStack cstack (toException ex))


-- | Similar to 'catch', but also catches exceptions of type @e@ wrapped in 'WithCallStack'.
catchWithCallStack
  :: forall e a m. (Exception e, MonadCatch m)
  => m a -> (Maybe CallStack -> e -> m a) -> m a
catchWithCallStack ma f =
  ma `catch` \(se :: SomeException) ->
    case fromExceptionWithCallStack @e se of
      Just (cstackMb, e) -> f cstackMb e
      Nothing -> throwM se

-- | Attempts to run the given action.
--
-- If it:
--
-- * suceeds, the value will be returned in a 'Right'.
-- * throws an exception of the given type @e@ (or an exception of
--   type @e@ wrapped in 'WithCallStack'), it will be returned in a 'Left'.
-- * throws an exception of any other type, it'll be rethrown
--   (retaining the original callstack, if any).
tryWithCallStack :: forall e a m. (MonadCatch m, Exception e) => m a -> m (Either (Maybe CallStack, e) a)
tryWithCallStack ma =
  try @_ @SomeException ma >>= \case
    Right a -> pure (Right a)
    Left se ->
      case fromExceptionWithCallStack @e se of
        Just e -> pure (Left e)
        Nothing -> throwM se

fromExceptionWithCallStack :: forall e. Exception e => SomeException -> Maybe (Maybe CallStack, e)
fromExceptionWithCallStack se =
  (do
    WithCallStack cstack e <- fromException @WithCallStack se
    e' <- cast @_ @e e
    Just (Just cstack, e')
  )
  <|>
  (fromException @e se <&> (Nothing,))