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

{-# OPTIONS_HADDOCK not-home #-}

-- | This module defines annotated exceptions and machinery to work with them.
module Test.Cleveland.Internal.Exceptions.Annotated
  ( module Test.Cleveland.Internal.Exceptions.Annotated
  ) where

import Data.Dependent.Map qualified as Map
import Data.GADT.Compare (GCompare(..), GEq(..), GOrdering(..))
import Data.Typeable (cast, eqT)
import Fmt (Buildable(..), Builder, build, pretty)
import Text.Show qualified as Show (Show(..))
import Type.Reflection (typeRep)

-- $setup
-- >>> import Test.Cleveland.Internal.Abstract
-- >>> import Test.Cleveland.Internal.Exceptions.CallStack

----------------------------------------------------------------------------
-- AnnotatedException
----------------------------------------------------------------------------

-- | 'SomeException' augmented with arbitrary annotations.
data AnnotatedException = AnnotatedException AnnotationMap SomeException
  deriving stock Show

instance Exception AnnotatedException where
  displayException (AnnotatedException annmap e) =
    pretty $ printAnnotations annmap (displayException e)

-- | Insert a new annotation into 'AnnotatedException'.
insertAnnEx'
  :: (Semigroup ann, ExceptionAnnotation ann)
  => ann -> AnnotatedException -> AnnotatedException
insertAnnEx' ann (AnnotatedException anns ex) = AnnotatedException (insertAnn ann anns) ex

-- | Insert a new annotation into 'SomeException'.
--
-- If it's already an 'AnnotatedException', simply adds a new annotation.
-- Otherwise, wraps the exception into 'AnnotatedException' first.
insertAnnEx
  :: (Semigroup ann, ExceptionAnnotation ann)
  => ann -> SomeException -> SomeException
insertAnnEx ann se = toException $! case fromException @AnnotatedException se of
  Just ex -> insertAnnEx' ann ex
  Nothing -> AnnotatedException (singletonAnn ann) se

-- | Remove an annotation of the given type from 'AnnotatedException'.
removeAnnEx' :: forall ann. ExceptionAnnotation ann => AnnotatedException -> AnnotatedException
removeAnnEx' (AnnotatedException anns ex) = AnnotatedException (removeAnn @ann anns) ex

-- | Remove an annotation of the given type from 'SomeException'. If it's not an
-- 'AnnotatedException', this is a no-op.
removeAnnEx :: forall ann. ExceptionAnnotation ann => SomeException -> SomeException
removeAnnEx se = case fromException @AnnotatedException se of
  Just ex -> toException $ removeAnnEx' @ann ex
  Nothing -> se

-- | Extract the given annotation type from 'AnnotatedException' if possible.
lookupAnnEx' :: forall ann. (ExceptionAnnotation ann) => AnnotatedException -> Maybe ann
lookupAnnEx' (AnnotatedException anns _) = lookupAnn @ann anns

-- | Extract the given annotation type from 'SomeException' if possible.
lookupAnnEx :: forall ann. (ExceptionAnnotation ann) => SomeException -> Maybe ann
lookupAnnEx se = lookupAnnEx' =<< fromException @AnnotatedException se

{- | Try to extract a given exception type, skipping over annotations. This
essentially makes 'AnnotatedException' transparent for this function.

This is intended to be used as the default implementation of 'fromException' for
exception types that are expected to be annotated, as to make
'AnnotatedException' transparent for 'catch', 'try', etc.

For example:

@
instance Exception GenericTestError where
  displayException = pretty
  fromException = fromPossiblyAnnotatedException
@

Now, when catching @GenericTestError@, it'll work even if wrapped in
'AnnotatedException':

>>> try @_ @GenericTestError (addCallStack $ throwM UnexpectedSuccess)
Left UnexpectedSuccess

When catching exceptions not explicitly set up like this, this trick
unfortunately won't work, so using this function explicitly is required:

>>> data ExceptionToCatch = ExceptionToCatch deriving stock Show
>>> instance Exception ExceptionToCatch
>>> action = addCallStack $ throwM ExceptionToCatch
>>> try @_ @ExceptionToCatch action
*** Exception: AnnotatedException ...
...ExceptionToCatch
>>> :{
action `catch` \(err :: SomeException) -> case fromPossiblyAnnotatedException err of
  Just (err' :: ExceptionToCatch) -> pass
  Nothing -> throwM err
:}
-}
fromPossiblyAnnotatedException :: Exception e => SomeException -> Maybe e
fromPossiblyAnnotatedException se =
  (do
    AnnotatedException _ e <- fromException se
    fromException e
  ) <|> fromExceptionDefault se
  where
    fromExceptionDefault (SomeException e) = cast e

----------------------------------------------------------------------------
-- AnnotationMap
----------------------------------------------------------------------------

data Key ann where
  Key :: ExceptionAnnotation a => Key a

instance GEq Key where
  geq (Key :: Key a) (Key :: Key b) = eqT @a @b

instance GCompare Key where
  gcompare (Key :: Key a) (Key :: Key b) =
    case compare (annotationPriority @a) (annotationPriority @b) of
      LT -> GLT
      GT -> GGT
      EQ -> gcompare (typeRep @a) (typeRep @b)

-- | A map from 'Key' to corresponding annotation.
newtype AnnotationMap = AnnotationMap { unAnnotationMap :: Map.DMap Key Identity }

instance Show AnnotationMap where
  show = Map.showTreeWith (\Key -> Show.show) True False . unAnnotationMap

-- | Print all annotations from an 'AnnotationMap', given the initial message.
--
-- Essentially right-fold using 'displayAnnotation'.
printAnnotations :: Buildable a => AnnotationMap -> a -> Builder
printAnnotations (AnnotationMap annmap) = flip (Map.foldrWithKey go) annmap . build
  where
    go :: Key t -> Identity t -> Builder -> Builder
    go Key (Identity a) acc = displayAnnotation a acc

-- | Try to find the given type @t@ in an 'AnnotationMap'.
lookupAnn :: ExceptionAnnotation t => AnnotationMap -> Maybe t
lookupAnn (AnnotationMap anns) = runIdentity <$> Map.lookup Key anns

-- | Construct a singleton 'AnnotationMap'.
singletonAnn :: ExceptionAnnotation t => t -> AnnotationMap
singletonAnn = AnnotationMap . Map.singleton Key . Identity

-- | Insert a new annotation into an 'AnnotationMap'. If an annotation of this
-- type already exists, merge the annotations using the 'Semigroup' instance.
insertAnn :: (Semigroup t, ExceptionAnnotation t) => t -> AnnotationMap -> AnnotationMap
insertAnn x (AnnotationMap anns) = AnnotationMap $ Map.insertWith (<>) Key (Identity x) anns

-- | Remove an annotation of a given type @t@ from 'AnnotationMap'.
removeAnn :: forall t. ExceptionAnnotation t => AnnotationMap -> AnnotationMap
removeAnn (AnnotationMap am) = AnnotationMap $ Map.delete (Key @t) am

----------------------------------------------------------------------------
-- ExceptionAnnotation
----------------------------------------------------------------------------

-- | Type class for exception annotations.
class (Show ann, Typeable ann) => ExceptionAnnotation ann where
  -- | Given an annotation and the error message (as 'Builder'), produce the
  -- text annotated with the annotation.
  displayAnnotation :: ann -> Builder -> Builder

  -- | Relative priority for sorting annotations. Annotations with higher
  -- priority will be applied first. Default is @0@.
  annotationPriority :: Int
  annotationPriority = 0

-- | Add an annotation to exceptions thrown by a monadic action.
annotateExceptions
  :: forall m a ann. (MonadCatch m, Semigroup ann, ExceptionAnnotation ann)
  => ann -> m a -> m a
annotateExceptions ann ma = ma `catch` (throwM . insertAnnEx ann)