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

{-# OPTIONS_HADDOCK not-home #-}

-- | Exception handling
module Test.Cleveland.Internal.Actions.ExceptionHandling
  ( module Test.Cleveland.Internal.Actions.ExceptionHandling
  ) where

import Data.List.NonEmpty qualified as NE
import Fmt (Doc, build, indentF, unlinesF)

import Lorentz
  (CustomError(..), ErrorTagMap, IsError, Label, MText, MustHaveErrorArg, errorTagToMText)
import Lorentz.Constraints
import Morley.Micheline.Class (fromExpression)
import Morley.Michelson.Untyped qualified as U
import Test.Cleveland.Internal.Abstract
import Test.Cleveland.Internal.Actions.Assertions
import Test.Cleveland.Internal.Actions.Helpers
import Test.Cleveland.Internal.Actions.Misc
import Test.Cleveland.Internal.Actions.TransferFailurePredicate
import Test.Cleveland.Internal.Exceptions

-- | Attempt to run an action and return its result or, if interpretation fails, an error.
attempt
  :: forall e caps m a. (HasCallStack, MonadCleveland caps m, Exception e)
  => m a -> m (Either e a)
attempt action = do
  caps <- ask
  let action' :: ClevelandBaseMonad caps a = runReaderT action caps
  lift $ cmiAttempt (getMiscCap caps) action'

-- | Asserts that a transfer should fail, and returns the exception.
catchTransferFailure :: (HasCallStack, MonadCleveland caps m) => m a -> m TransferFailure
catchTransferFailure action =
  attempt action >>= \case
    Left err -> return err
    Right _ -> runIO $ throwM UnexpectedSuccess

-- | Asserts that a transfer should fail, and runs some 'TransferFailurePredicate's over the
-- exception.
--
-- > expectTransferFailure (failedWith (constant @MText "NOT_ADMIN")) $
-- >   call contractAddr (Call @"Ep") arg
--
-- > call contractAddr (Call @"Ep") arg & expectTransferFailure
-- >   ( failedWith (customError #tag 3) &&
-- >     addressIs contractAddr
-- >   )
expectTransferFailure :: (HasCallStack, MonadCleveland caps m) => TransferFailurePredicate -> m a -> m ()
expectTransferFailure predicate act = do
  err <- catchTransferFailure act
  checkTransferFailure err predicate

-- | Check whether a given predicate holds for a given 'TransferFailure'.
checkTransferFailure :: (HasCallStack, MonadCleveland caps m) => TransferFailure -> TransferFailurePredicate -> m ()
checkTransferFailure err predicate =
  case go predicate of
    Nothing -> pass
    Just expectedOutcome -> failure $ unlinesF
      [ "Expected transfer to fail with an error such that:"
      , ""
      , indentF 2 expectedOutcome
      , ""
      , "But these conditions were not met."
      , "Actual transfer error:"
      , indentF 2 $ build err
      ]
  where
    -- Collect descriptions of all failed predicates
    -- Note that 'Nothing' signifies success here, and Just is a failure description.
    go :: TransferFailurePredicate -> Maybe Doc
    go = \case
      AndPredicate ps -> fmap (fmtExpectedOutcomes "AND") . nonEmpty . mapMaybe go $ toList ps
        -- if all results are successful, i.e. 'Nothing', the result is 'Nothing'
      OrPredicate ps -> fmtExpectedOutcomes "OR" <$> traverse go ps
        -- if there is at least one success, i.e. 'Nothing', the result is 'Nothing'
      TransferFailurePredicate desc p
        | p err -> Nothing
        | otherwise -> Just $ build desc

    fmtExpectedOutcomes :: Doc -> NonEmpty Doc -> Doc
    fmtExpectedOutcomes delimiter = \case
      expectedOutcome :| [] -> expectedOutcome
      expectedOutcomes ->
        unlinesF
          [ "("
          , indentF 2 $ unlinesF $ NE.intersperse delimiter expectedOutcomes
          , ")"
          ]

-- | Asserts that interpretation of a contract ended with @FAILWITH@, returning the given constant
-- value.
expectFailedWith
  :: forall err a caps m
   . (HasCallStack, MonadCleveland caps m, NiceConstant err)
  => err -> m a -> m ()
expectFailedWith err = expectTransferFailure $ failedWith (constant err)

-- | Asserts that interpretation of a contract ended with @FAILWITH@, returning the given lorentz
-- error.
expectError
  :: forall err a caps m
   . (HasCallStack, MonadCleveland caps m, IsError err)
  => err -> m a -> m ()
expectError err = expectTransferFailure $ failedWith (lerror err)

-- | Asserts that interpretation of a contract ended with @FAILWITH@, returning the given custom
-- lorentz error.
expectCustomError
  :: forall arg a tag caps m
   . ( HasCallStack, MonadCleveland caps m
     , IsError (CustomError tag)
     , MustHaveErrorArg tag (MText, arg)
     )
  => Label tag -> arg -> m a -> m ()
expectCustomError tag arg = expectTransferFailure $ failedWith (customError tag arg)

-- | Version of 'expectCustomError' for error with @unit@ argument.
expectCustomError_
  :: ( HasCallStack, MonadCleveland caps m
     , IsError (CustomError tag)
     , MustHaveErrorArg tag (MText, ())
     )
  => Label tag -> m a -> m ()
expectCustomError_ tag = expectCustomError tag ()

-- | Version of 'expectCustomError' specialized for expecting @NoErrorArg@s.
expectCustomErrorNoArg
  :: ( HasCallStack, MonadCleveland caps m
     , IsError (CustomError tag)
     , MustHaveErrorArg tag MText
     )
  => Label tag -> m a -> m ()
expectCustomErrorNoArg tag = expectTransferFailure $ failedWith (customErrorNoArg tag)

-- | Version of 'expectCustomError' that ignores the argument (or whether it
-- even exists) and only checks the tag.
expectCustomErrorAnyArg
  :: ( HasCallStack, MonadCleveland caps m )
  => Label tag -> m a -> m ()
expectCustomErrorAnyArg tag = expectTransferFailure $ failedWithPredicate
  ("Contract failed with any custom error tagged with \"" <> build expectedTag <> "\"")
  \expr ->
    -- the more "proper" way to handle this would be to go over all possible cases
    -- of 'CustomErrorArgRep', but it's polymorphic, and we don't have a wildcard
    -- in the typed representation, so we'd have to try to run 'fromExpression's
    -- for all possible argument types, and that's just not worth it.  So we use
    -- our knowledge of the untyped representation instead. -- @lierdakil
    (Just expectedTag ==) $ fromExpression @U.Value expr & rightToMaybe >>= \case
      U.ValueString text -> Just text
      U.ValuePair (U.ValueString text) _ -> Just text
      -- right combs can be represented as sequences of at least two elements
      U.ValueSeq (U.ValueString text :| _ : _) -> Just text
      _ -> Nothing
  where
    expectedTag = errorTagToMText tag

-- | Asserts that interpretation of a contract ended with @FAILWITH@, returning the given lorentz
-- numeric error.
expectNumericError
  :: forall err a caps m
   . (HasCallStack, MonadCleveland caps m, IsError err)
  => ErrorTagMap -> err -> m a -> m ()
expectNumericError tagMap err = expectTransferFailure $ failedWith (numericError tagMap err)

-- | Prefix error messages potentially thrown from the given code block.
--
-- The prefix will be put at a separate line before the main text, if text is multiline, otherwise
-- it will be separated from the main text with @: @.
--
-- This affects errors produced by functions like 'failure', 'assert', '@==', etc.
-- Errors related to events in the chain will not be touched.
--
-- Example:
--
-- > for [1..10] \i -> clarifyErrors ("For i=" +| i |+ "") $
-- >   askContract i @@== i * 2
clarifyErrors :: forall caps m a. (MonadCleveland caps m)
              => Doc -> m a -> m a
clarifyErrors message action =
  attempt action >>= \case
    Left e -> withCap getMiscCap \cap -> cmiThrow cap $
      insertAnnEx (ErrorsClarification [message]) e
    Right val -> pure val