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

-- | Integration with integrational testing engine (pun intended).
module Test.Cleveland.Internal.Pure
  ( PureM(..)
  , runEmulatedT

  -- * Capability implementations
  , PureState
  , TestError(..)
  , emulatedImpl
  , clevelandOpsImpl
  , clevelandMiscImpl

  -- * Initial environment for Emulated tests
  , initEnv

  -- * Support functions
  , failedInsideBranch
  , moneybagAlias
  , emptyScenarioBranch

  -- * Optics
  , psSecretKeys
  , psDefaultAliasesCounter
  , psRefillableAddresses
  , psNow
  , psLevel
  , psMinBlockTime
  , psGState
  ) where

import Control.Lens (assign, at, each, makeLenses, modifying, to, (%=), (.=), (?=), (?~))
import Control.Lens.Unsound (lensProduct)
import Control.Monad.Catch.Pure (CatchT, runCatchT)
import Control.Monad.Writer (MonadWriter, WriterT, listen, runWriterT, tell)
import Data.Constraint (Dict(..), withDict, (\\))
import Data.Default (def)
import Data.Map qualified as Map
import Data.Ratio ((%))
import Data.Set qualified as Set
import Fmt (Buildable(..), Builder, build, pretty, unlinesF, (+|), (|+))
import Time (Second, toNum, toUnit)

import Lorentz (Mutez, NiceComparable, pattern DefEpName)
import Lorentz qualified as L
import Lorentz.Entrypoints (HasEntrypointArg, TrustEpName(..), useHasEntrypointArg)
import Morley.AsRPC (HasRPCRepr(AsRPC), notesAsRPC, rpcStorageScopeEvi, valueAsRPC)
import Morley.Client (OperationInfo(..))
import Morley.Michelson.Interpret
  (InterpretError(..), InterpretResult(..), MichelsonFailed(..), MichelsonFailureWithStack(..))
import Morley.Michelson.Runtime hiding (ExecutorOp(..), transfer)
import Morley.Michelson.Runtime qualified as Runtime (ExecutorOp(..))
import Morley.Michelson.Runtime.Dummy
  (dummyLevel, dummyMaxSteps, dummyMinBlockTime, dummyNow, dummyOrigination)
import Morley.Michelson.Runtime.GState
  (GState(..), genesisAddress, genesisSecretKey, gsChainIdL, gsContractAddressAliasesL,
  gsContractAddressesL, gsCounterL, gsImplicitAddressAliasesL, gsVotingPowersL, initGState,
  lookupBalance)
import Morley.Michelson.TypeCheck
  (BigMapFinder, TCError(..), TypeCheckOptions(..), typeCheckContractAndStorage,
  typeCheckValueRunCodeCompat, typeCheckingWith)
import Morley.Michelson.Typed
  (BigMapId(..), IsoValue, SingI, SomeAnnotatedValue(..), SomeConstrainedValue(SomeValue),
  SomeVBigMap(..), SomeValue, ToT, Value, Value'(..), castM, dfsFoldMapValue, fromVal, toVal)
import Morley.Michelson.Typed qualified as T
import Morley.Michelson.Typed.Operation
  (EmitOperation(..), OriginationOperation(..), TransferOperation(..))
import Morley.Michelson.Untyped qualified as U
import Morley.Tezos.Address
import Morley.Tezos.Address.Alias
import Morley.Tezos.Core (Timestamp, timestampPlusSeconds, unsafeSubMutez, zeroMutez)
import Morley.Tezos.Crypto (SecretKey(..), detSecretKey, sign, toPublic)
import Morley.Util.MismatchError
import Morley.Util.Named

import Morley.Util.Bimap qualified as Bimap
import Test.Cleveland.Internal.Abstract
import Test.Cleveland.Internal.Exceptions (addCallStack, catchWithCallStack, throwWithCallStack)
import Test.Cleveland.Lorentz
import Test.Cleveland.Util (ceilingUnit)

data PureState = PureState
  { _psSecretKeys :: Map ImplicitAddress SecretKey
  , _psDefaultAliasesCounter :: DefaultAliasCounter
  , _psRefillableAddresses :: Set ImplicitAddress
  , _psNow :: Timestamp
  , _psLevel :: Natural
  , _psMinBlockTime :: Natural
  , _psGState :: GState
  }
  deriving stock (Eq, Show)

instance MonadState PureState PureM where
  get = ask >>= readIORef
  put = (ask >>=) . flip writeIORef

newtype PureM a = PureM
  { unPureM :: ReaderT (IORef PureState) (CatchT (WriterT LogsInfo IO)) a
  }
  deriving newtype (Functor, Applicative, Monad, MonadIO, MonadThrow, MonadCatch,
                    MonadReader (IORef PureState), MonadWriter LogsInfo, MonadFail)

data TestError
  = UnexpectedTypeCheckError TCError
  | UnexpectedStorageType (MismatchError T.T)
  | UnexpectedBigMapType Natural (MismatchError T.T)
  | CustomTestError Text
  deriving stock Show

makeLenses ''PureState

instance Buildable TestError where
  build = \case
    UnexpectedTypeCheckError tcErr ->
      "Unexpected type check error. Reason: " +| tcErr |+ ""
    UnexpectedStorageType merr ->
      "Unexpected storage type.\n" +| merr |+ ""
    UnexpectedBigMapType bigMapId mismatchError ->
      unlinesF
        [ "A big_map with the ID " +| bigMapId |+ " was found, but it does not have the expected type."
        , build mismatchError
        ]
    CustomTestError msg -> pretty msg

instance Exception TestError where
  displayException = pretty

moneybagAlias :: ImplicitAlias
moneybagAlias = ImplicitAlias "moneybag"

runEmulatedT :: ImplicitAlias -> EmulatedT PureM a -> IO a
runEmulatedT moneybagAlias' scenario = do
  let clevelandCaps = ClevelandCaps
        { ccSender = Sender genesisAddress
        , ccMoneybag = Moneybag genesisAddress
        , ccMiscCap = clevelandMiscImpl
        , ccOpsCap = clevelandOpsImpl
        }
      caps = EmulatedCaps emulatedImpl clevelandCaps
  let pureM = runReaderT scenario caps
  env <- newIORef (initEnv moneybagAlias')
  (res, _logs) <- runWriterT $ runCatchT $ runReaderT (unPureM pureM) env
  either throwM pure $ res

emulatedImpl :: EmulatedImpl PureM
emulatedImpl =
  EmulatedImpl
    { eiBranchout = \(scenarios :: [(Text, PureM ())]) ->
        forM_ scenarios $ \(name, scenario) -> do
          aliasesState <- get
          newRef <- newIORef aliasesState
          local (\_ -> newRef) scenario `catchWithCallStack` \originalCallStackMb err ->
            maybe throwM throwWithCallStack originalCallStackMb $ failedInsideBranch name err
      , eiGetStorage = addCallStack . exceptionHandler . getStorageImpl
    , eiGetMorleyLogs = getMorleyLogsImpl
    , eiSetVotingPowers = assign (psGState . gsVotingPowersL)
    }

clevelandOpsImpl :: Sender -> ClevelandOpsImpl PureM
clevelandOpsImpl (Sender sender) =
  mapClevelandOpsImplExceptions (addCallStack . exceptionHandler)
    ClevelandOpsImpl
    { coiRunOperationBatch = mapM \case
        OpOriginate UntypedOriginateData{..} -> do
          OpOriginate <$> originate uodContract uodName uodStorage uodBalance
        OpTransfer TransferData{..} -> do
          let fromAddr = #from :! sender
          let toAddr = #to :! toL1Address tdTo
          refillable <- isAddressRefillable sender
          when refillable $ do
            balance <- getBalance sender
            when (balance < tdAmount) $ do
              let moneybag = #from :! genesisAddress
                  toSender = #to :! sender
              void $ transfer @() moneybag toSender (unsafeSubMutez tdAmount balance)
                (TrustEpName DefEpName) ()
          -- Here @toAddr@ is 'Address', so we can not check anything
          -- about it and assume that entrypoint is correct. We pass
          -- unit as contract parameter because it won't be checked
          -- anyway.
          emitOps <- transfer @() fromAddr toAddr tdAmount
            (TrustEpName tdEntrypoint) tdParameter

          return $ OpTransfer $ emitOpToContractEvent <$> emitOps
        OpReveal{} -> do
          -- We do not care about reveals in our Morley runtime
          return $ OpReveal ()
    }

emitOpToContractEvent :: EmitOperation -> ContractEvent
emitOpToContractEvent EmitOperation{eoEmit=T.Emit{..},..} = ContractEvent
  { cePayload = case T.notesT emNotes of
      -- A dirty hack to match reference
      T.TUnit -> Nothing
      _ -> Just $ SomeAnnotatedValue emNotes emValue
  , ceTag = emTag
  , ceSource = eoSource
  }

clevelandMiscImpl :: ClevelandMiscImpl PureM
clevelandMiscImpl =
  mapClevelandMiscImplExceptions (addCallStack . exceptionHandler)
    ClevelandMiscImpl
    { cmiRunIO = \action -> liftIO (try action) >>= \case
        Right res -> pure res
        Left (err :: SomeException) -> throwM err

    , cmiResolveAddress = \case
        a@ImplicitAlias{} -> resolveImplicit a
        a@ContractAlias{} -> resolveContract a

    , cmiSignBytes = \bs addr -> do
        sk <- getSecretKey addr
        liftIO $ sign sk bs

    , cmiGenKey = \sodAlias -> do
      alias <- resolveSpecificOrDefaultAlias sodAlias
      smartGenKey Nothing alias

    , cmiGenFreshKey =
        \sodAlias -> do
          alias <- resolveSpecificOrDefaultAlias sodAlias
          existingAddr <- use (psGState . gsImplicitAddressAliasesL . at alias)
          smartGenKey existingAddr alias

    , cmiOriginateLargeUntyped = originateUntyped

    -- Comments are not supported by integrational testing engine (yet).
    , cmiComment = const pass
    , cmiGetPublicKey = \addr -> do
        toPublic <$> getSecretKey addr
    , cmiGetDelegate = \addr -> do
        ContractState _ _ _ delegate <- contractStorage addr
        pure delegate
    , cmiRegisterDelegate = const pass
    , cmiGetChainId = use $ psGState . gsChainIdL
    , cmiAdvanceTime = \time -> do
        modifying psNow . flip timestampPlusSeconds $
          toNum @Second @Integer $ ceilingUnit $ toUnit @Second time

    , cmiAdvanceToLevel = \fn ->
        -- do not go back in levels
        modifying psLevel (\cl -> max (fn cl) cl)

    , cmiGetNow = use psNow
    , cmiGetLevel = use psLevel
    , cmiGetApproximateBlockInterval = sec . (% 1) <$> use psMinBlockTime
    , cmiAttempt = try
    , cmiThrow = throwM
    , cmiMarkAddressRefillable = setAddressRefillable
    , cmiGetBalance = \(MkConstrainedAddress a) -> getBalance a
    , cmiUnderlyingImpl = pure $ Left emulatedImpl
    , cmiFailure = failure
    , ..
    }
  where
    setAddressRefillable addr = psRefillableAddresses %= Set.insert addr

    originateUntyped :: Sender -> UntypedOriginateData 'IsLarge -> PureM ContractAddress
    originateUntyped _ UntypedOriginateData {..} = do
      originate uodContract uodName uodStorage uodBalance

    cmiGetBigMapValueMaybe
      :: forall k v.
         (NiceComparable k, IsoValue v)
      => BigMapId k v
      -> k
      -> PureM (Maybe v)
    cmiGetBigMapValueMaybe (BigMapId bmId) k = runMaybeT do
      -- The RPC does not distinguish between "the bigmap does not exist"
      -- and "the bigmap exists, but the key doesn't", so we mimic the RPC's behaviour here.
      -- We simply return `Nothing` in both cases.
      VBigMap _ bigMap <- findBigMapByIdMaybe @(ToT k) @(ToT v) bmId
      hoistMaybe $ fromVal @v <$> Map.lookup (toVal k) bigMap

    cmiGetAllBigMapValuesMaybe
      :: forall k v.
         (NiceComparable k, IsoValue v)
      => BigMapId k v
      -> PureM (Maybe [v])
    cmiGetAllBigMapValuesMaybe (BigMapId bmId) = runMaybeT do
      VBigMap _ bigMap <- findBigMapByIdMaybe @(ToT k) @(ToT v) bmId
      pure $ fromVal @v <$> Map.elems bigMap

    -- | In a real chain, when we retrieve a contract's storage via the Tezos RPC,
    -- the storage expression will have all the big_maps replaced with their respective big_map IDs.
    --
    -- Here, we mimic the RPC's behaviour.
    --
    -- We expect all big_maps in the storage to already have an ID.
    -- IDs are assigned to big_maps by the interpreter/runtime when:
    --   * A contract with big_maps in its storage is originated
    --   * A transfer is made and the parameter contains big_maps
    --   * A contract's code is run and it calls `EMPTY_BIG_MAP`, `DUP` or `DUP n`.
    cmiGetSomeStorage :: ContractAddress -> PureM SomeAnnotatedValue
    cmiGetSomeStorage addr = do
      ContractState _ contract (storage :: Value t) _ <- contractStorage addr
      pure $
        SomeAnnotatedValue
          (notesAsRPC $ T.cStoreNotes contract)
          (valueAsRPC storage)
          \\ rpcStorageScopeEvi @t

    -- Generate a fresh address for a given alias.
    --
    -- If this alias is not yet associated with any address,
    -- we use the alias as the seed for generating a brand new address.
    --
    -- If this alias is already associated with an address,
    -- then we have to use a _different_ seed to guarantee we'll get a brand new address.
    -- Therefore, we concatenate the alias with the existing address and use the
    -- result as a seed for generating a brand new address.
    smartGenKey :: Maybe ImplicitAddress -> ImplicitAlias -> PureM ImplicitAddress
    smartGenKey existingAddr alias@(ImplicitAlias aliasTxt) = do
      let
        seed = maybe aliasTxt (mappend aliasTxt . pretty) existingAddr
        sk = detSecretKey (encodeUtf8 seed)
        addr = detGenKeyAddress (encodeUtf8 seed)

      -- Save alias/address association.
      psGState . gsImplicitAddressAliasesL . at alias ?= addr
      -- Save the address's secret key.
      psSecretKeys . at addr .= Just sk
      pure addr

    resolveSpecificOrDefaultAlias (SpecificAlias alias) =
      return alias
    resolveSpecificOrDefaultAlias DefaultAlias = do
      DefaultAliasCounter counter <- use psDefaultAliasesCounter
      psDefaultAliasesCounter %= \(DefaultAliasCounter i) -> DefaultAliasCounter $ i + 1
      return $ mkDefaultAlias counter

    cmiRunCode
      :: forall cp st vd. (HasRPCRepr st, T.IsoValue (AsRPC st))
      => Sender -> RunCode cp st vd -> PureM (AsRPC st)
    cmiRunCode (Sender sender) (RunCode rcContract rcParameter rcStorage rcAmount rcLevel rcNow rcBalance rcSource) = do
      -- Pattern match on the contract constructor to reveal
      -- a proof of `NiceParameter cp` and `NiceStorage st`
      L.Contract{} <- pure rcContract
      param <- typeCheckVal @(ToT cp) rcParameter
      storage <- typeCheckVal @(ToT st) rcStorage

      (now, level) <- use $ psNow `lensProduct` psLevel
      psNow %= maybe id const rcNow
      psLevel %= maybe id const rcLevel
      res <- interpret do
        counter0 <- use $ esGState . gsCounterL
        contractAddr <-
          executeOrigination ! #isGlobalOp True $
            (dummyOrigination storage (L.toMichelsonContract rcContract) counter0) { ooBalance = zeroMutez }
              \\ L.niceStorageEvi @st
              \\ L.niceParameterEvi @cp

        esSourceAddress .= (MkConstrainedAddress <$> rcSource)
        counter1 <- use $ esGState . gsCounterL
        let overrideContractBalance = Just rcBalance
        executeTransfer (#isGlobalOp :! False) overrideContractBalance def $
          TransferOperation
            { toDestination = MkAddress contractAddr
            , toCounter = counter1
            , toTxData = TxData
              { tdSenderAddress = MkConstrainedAddress sender
              , tdParameter = TxTypedParam param \\ L.niceParameterEvi @cp
              , tdEntrypoint = DefEpName
              , tdAmount = rcAmount
              }
            }
        pure contractAddr
      psNow .= now
      psLevel .= level
      case res of
        Left executorError -> throwEE executorError
        Right (executorRes, contractAddr) -> do
          -- Find the storage of the contract and return it in its RPC representation.
          case executorRes ^. erGState . gsContractAddressesL . at contractAddr of
            Nothing ->
              failure $ "Internal error: failed to find contract: '" +| contractAddr |+ "'"
            Just (ContractState _ _ (finalStorage :: Value actualSt) _) -> do
              finalStorage' <- castM @actualSt @(ToT st) finalStorage (throwM . UnexpectedStorageType)
              pure $ fromVal $ valueAsRPC finalStorage'
      where
        typeCheckVal :: forall t. SingI t => U.Value -> PureM (Value t)
        typeCheckVal untypedVal = do
          bigMapFinder <- mkBigMapFinder
          let res =
                typeCheckingWith (TypeCheckOptions False False) $
                  typeCheckValueRunCodeCompat bigMapFinder untypedVal
          case res of
            Right val -> pure val
            Left tcErr -> throwM $ UnexpectedTypeCheckError tcErr

mkBigMapFinder :: PureM BigMapFinder
mkBigMapFinder = do
  pureState <- get

  pure \bigMapId ->
    pureState ^?
      psGState . gsContractAddressesL . each . to getContractStorage . to (getBigMapsWithId bigMapId) . each
  where
    getContractStorage :: ContractState -> SomeValue
    getContractStorage (ContractState _ _ storage _) = SomeValue storage

    getBigMapsWithId :: Natural -> SomeValue -> [SomeVBigMap]
    getBigMapsWithId bigMapId (SomeValue val) =
      dfsFoldMapValue
        (\v -> case v of
            VBigMap (Just bigMapId') _ | bigMapId' == bigMapId -> [SomeVBigMap v]
            _ -> []
        )
        val

-- | Traverse storage values of all contracts and looks for a big_map with the given ID.
findBigMapByIdMaybe
  ::forall k v. (SingI v, SingI k)
  => Natural -> MaybeT PureM (Value ('T.TBigMap k v))
findBigMapByIdMaybe bigMapId = do
  SomeVBigMap (v@VBigMap{} :: Value t) <- MaybeT $ mkBigMapFinder <*> pure bigMapId
  castM @t @('T.TBigMap k v) v $ throwM . UnexpectedBigMapType bigMapId

isAddressRefillable :: ImplicitAddress -> PureM Bool
isAddressRefillable addr = Set.member addr <$> use psRefillableAddresses

getBalance :: L1AddressKind kind => KindedAddress kind -> PureM Mutez
getBalance addr = do
  gs <- use psGState
  pure $ fromMaybe zeroMutez $ lookupBalance addr gs

exceptionHandler :: PureM a -> PureM a
exceptionHandler action = try action >>= \case
  Left err -> exceptionToTransferFailure err >>= throwM
  Right res -> return res
  where
    exceptionToTransferFailure :: ExecutorError' AddressAndAlias -> PureM TransferFailure
    exceptionToTransferFailure err = case err of
      EEZeroTransaction addr -> return $
        TransferFailure addr EmptyTransaction
      EEIllTypedParameter addr _ -> return $
        TransferFailure addr BadParameter
      EEUnexpectedParameterType addr _ -> return $
        TransferFailure addr BadParameter
      EEInterpreterFailed addr (InterpretError (MichelsonFailureWithStack{..}, _)) ->
        case mfwsFailed of
          MichelsonFailedWith val -> return $
            TransferFailure addr $
              FailedWith (EOTVTypedValue val) (Just mfwsErrorSrcPos)
          MichelsonArithError (T.ShiftArithError{}) -> return $
            TransferFailure addr ShiftOverflow
          MichelsonArithError (T.MutezArithError errType _ _) -> return $
            TransferFailure addr $ MutezArithError errType
          MichelsonGasExhaustion -> return $
            TransferFailure addr GasExhaustion
          _ -> throwM err
      _ -> throwM err


getMorleyLogsImpl :: PureM a -> PureM (LogsInfo, a)
getMorleyLogsImpl action = swap <$> listen action

getStorageImpl
  :: forall st addr. (ToStorageType st addr)
  => addr -> PureM st
getStorageImpl addr = do
  withDict (pickNiceStorage @st addr) $ do
    ContractState _ _ (storage :: Value actualT) _ <- contractStorage (toContractAddress addr)
    val <- castM @actualT @(ToT st) storage (throwM . UnexpectedStorageType)
    pure $ T.fromVal val

-- Attempt to retrieve a ContractState given for the given address. Fails if the
-- address is unknown or the address is a simple address (contract without
-- code and storage).
contractStorage :: ContractAddress -> PureM ContractState
contractStorage addr = do
  GState{..} <- use psGState
  case Map.lookup addr gsContractAddresses of
    Just contractState -> pure contractState
    Nothing -> unknownAddress addr

resolveImplicit :: ImplicitAlias -> PureM ImplicitAddress
resolveImplicit alias = do
  use (psGState . gsImplicitAddressAliasesL . at alias) >>= maybe (unknownAlias alias) pure

resolveContract :: ContractAlias -> PureM ContractAddress
resolveContract alias = do
  use (psGState . gsContractAddressAliasesL . at alias) >>= maybe (unknownAlias alias) pure

unknownAddress :: KindedAddress kind -> PureM whatever
unknownAddress =
  throwM . CustomTestError .
  mappend "Unknown address provided: " . pretty

unknownAlias :: Alias kind -> PureM whatever
unknownAlias =
  throwM . CustomTestError .
  mappend "Unknown address alias: " . pretty

-- | Make branch names for a case when we are not within any branch.
emptyScenarioBranch :: ScenarioBranchName
emptyScenarioBranch = ScenarioBranchName []

-- | Add a new branch element to names provided by inner
-- 'Test.Cleveland.branchout' calls.
appendScenarioBranch :: Text -> ScenarioBranchName -> ScenarioBranchName
appendScenarioBranch brName (ScenarioBranchName branches) =
  ScenarioBranchName (brName : branches)

failure :: forall a. Builder -> PureM a
failure = throwM . CustomTestError . pretty

getSecretKey :: ImplicitAddress -> PureM SecretKey
getSecretKey addr = do
  use (psSecretKeys . at addr) >>= \case
    Nothing -> unknownAddress addr
    Just sk -> pure sk

----------------------------------------------------------------------------
-- Support functions
----------------------------------------------------------------------------

initEnv :: ImplicitAlias -> PureState
initEnv alias = PureState
  { _psSecretKeys = one (genesisAddress, genesisSecretKey)
  , _psDefaultAliasesCounter = DefaultAliasCounter 0
  , _psRefillableAddresses = Set.empty
  , _psNow = dummyNow
  , _psLevel = dummyLevel
  , _psGState = initGState & gsImplicitAddressAliasesL . at alias ?~ genesisAddress
  , _psMinBlockTime = dummyMinBlockTime
  }

failedInsideBranch :: Text -> SomeException -> FailedInBranch
failedInsideBranch name err = case fromException @FailedInBranch err of
  Just (FailedInBranch branch exception) ->
    FailedInBranch (appendScenarioBranch name branch) exception
  Nothing ->
    FailedInBranch (appendScenarioBranch name emptyScenarioBranch) err

----------------------------------------------------------------------------
-- Emulator internals
----------------------------------------------------------------------------

transfer
  :: forall cp vd epRef epArg addr.
     (HasEntrypointArg cp epRef epArg, T.IsoValue epArg, L.ToTAddress cp vd addr)
  => "from" :! ImplicitAddress
  -> "to" :! addr
  -> Mutez
  -> epRef
  -> epArg
  -> PureM [EmitOperation]
transfer (arg #from -> from) (arg #to -> toAddr) money epRef param =
  let L.TAddress to' = L.toTAddress @cp @vd toAddr in
  case useHasEntrypointArg @cp @epRef @epArg epRef of
    (Dict, epName) -> registerInterpretation do
      executeGlobalOperations def $ one $
        Runtime.TransferOp $ TransferOperation
          { toDestination = to'
          , toCounter = 0
          , toTxData = TxData
            { tdSenderAddress = MkConstrainedAddress from
            , tdParameter = withDict (T.properParameterEvi @(ToT epArg)) $
                TxTypedParam (T.toVal param)
            , tdEntrypoint = epName
            , tdAmount = money
            }
          }

-- | Originate a contract with given initial storage and balance. Its
-- address is returned.
originate :: U.Contract -> ContractAlias -> U.Value -> Mutez -> PureM ContractAddress
originate uContract alias uStorage balance =
  case typeCheckingWith def $ typeCheckContractAndStorage uContract uStorage of
    Left tcErr -> throwM $ UnexpectedTypeCheckError tcErr
    Right (T.SomeContractAndStorage contract storage) -> do
      counter <- gsCounter <$> use psGState
      registerInterpretation . executeGlobalOrigination $
        (dummyOrigination storage contract counter)
          { ooBalance = balance
          , ooAlias = Just alias
          }

throwEE :: ExecutorError -> PureM a
throwEE err =
  -- Replace all `Address`es with `AddressAndAlias` in the error and throw it.
  traverse addrToAddressAndAlias err >>= throwM
  where
    addrToAddressAndAlias :: Address -> PureM AddressAndAlias
    addrToAddressAndAlias addr@(MkAddress kindedAddr) =
      case kindedAddr of
        ContractAddress{} ->
          use $ psGState . gsContractAddressAliasesL . Bimap.flipped . at kindedAddr . to (AddressAndAlias addr)
        _ ->
          pure $ AddressAndAlias addr Nothing

-- | Runs a set of operations and updates the engine's state.
registerInterpretation :: ExecutorM a -> PureM a
registerInterpretation action = do
  interpretedResult <- interpret action
  tell $ extractLogs interpretedResult

  case interpretedResult of
    Right (executorRes, res) -> do
      psGState .= _erGState executorRes
      pure res
    Left executorError -> throwEE executorError

  where
    extractLogs :: Either ExecutorError (ExecutorRes, a) -> [ScenarioLogs]
    extractLogs = \case
      Left (EEInterpreterFailed addr (InterpretError e)) -> [ScenarioLogs addr $ snd e]
      Right (res, _) -> res ^. erInterpretResults <&> \(addr, InterpretResult{..}) ->
        ScenarioLogs addr iurMorleyLogs
      _ -> []

-- | Interpret an action and return the result _without_ updating the engine's state.
interpret :: ExecutorM a -> PureM (Either ExecutorError (ExecutorRes, a))
interpret action = do
  now <- use psNow
  level <- use psLevel
  gState <- use psGState
  minBlockTime <- use psMinBlockTime
  pure $ runExecutorM now level minBlockTime dummyMaxSteps gState action