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

{-# OPTIONS_HADDOCK not-home #-}

-- | Implementation that works with real Tezos network, it
-- talks to a Tezos node and uses @octez-client@.

module Test.Cleveland.Internal.Client
  ( module Test.Cleveland.Internal.Client
  ) where

import Control.Lens (_head, each, filtered)
import Data.Aeson.Text qualified as J
import Data.Constraint (withDict, (\\))
import Data.Ratio ((%))
import Data.Set qualified as Set
import Data.Singletons (demote)
import Data.Time (NominalDiffTime, UTCTime, diffUTCTime, secondsToNominalDiffTime)
import Data.Typeable (cast)
import Fmt (Buildable(build), Doc, indentF, pretty, unlinesF, (+|), (|+))
import System.IO (hFlush)
import Time (KnownDivRat, Second, Time, sec, threadDelay, toNum, toUnit)
import Unsafe qualified (fromIntegral)

import Lorentz (NicePackedValue)
import Lorentz qualified as L
import Lorentz.Constraints.Scopes (NiceUnpackedValue)
import Morley.AsRPC (AsRPC, HasRPCRepr(..), TAsRPC, notesAsRPC, rpcSingIEvi)
import Morley.Client
  (AliasBehavior(..), MorleyClientEnv, OperationInfo(..), RevealData(..), disableAlphanetWarning,
  runMorleyClientM)
import Morley.Client qualified as Client
import Morley.Client.Action (Result)
import Morley.Client.Action.Common (DelegationData(..), runErrorsToClientError)
import Morley.Client.App (failOnTimeout, retryOnceOnTimeout)
import Morley.Client.Full qualified as Client
import Morley.Client.Logging (logInfo, logWarning)
import Morley.Client.RPC.Error qualified as RPC
import Morley.Client.RPC.Types
  (AppliedResult(..), BlockConstants(bcHeader), BlockHeaderNoHash(bhnhLevel, bhnhTimestamp),
  BlockId(..), EventOperation(..), OperationHash, OriginationScript(..),
  ProtocolParameters(ProtocolParameters, ppCostPerByte, ppMinimalBlockDelay, ppOriginationSize),
  WithSource(..))
import Morley.Client.RPC.Types qualified as RPC
import Morley.Client.TezosClient.Impl qualified as TezosClient
import Morley.Client.TezosClient.Types (tceMbTezosClientDataDirL)
import Morley.Client.Util qualified as Client
import Morley.Micheline
  (Expression, MichelinePrimitive(..), StringEncode(..), TezosInt64, TezosMutez(..),
  _ExpressionPrim, _ExpressionSeq, fromExpression, mpaArgsL, mpaPrimL, toExpression)
import Morley.Michelson.Text (unMText)
import Morley.Michelson.Typed (BigMapId, SomeAnnotatedValue(..), SomeContractAndStorage(..), toVal)
import Morley.Michelson.Typed qualified as T
import Morley.Michelson.Untyped qualified as U
import Morley.Tezos.Address
import Morley.Tezos.Address.Alias
import Morley.Tezos.Core as Tezos
  (Mutez, Timestamp(..), addMutez, subMutez, timestampFromUTCTime, unsafeAddMutez, unsafeMulMutez,
  unsafeSubMutez)
import Morley.Tezos.Crypto
import Morley.Util.Constrained
import Morley.Util.Named
import Test.Cleveland.Internal.Abstract
import Test.Cleveland.Internal.Common
import Test.Cleveland.Internal.Exceptions
import Test.Cleveland.Lorentz (toL1Address)
import Test.Cleveland.Util (ceilingUnit)

-- | Construct 'Client.MorleyOnlyRpcEnv' from 'NetworkEnv'.
mkMorleyOnlyRpcEnvNetwork
  :: NetworkEnv
  -> [SecretKey]  -- ^ Extra secrets that should be known
  -> Client.MorleyOnlyRpcEnv
mkMorleyOnlyRpcEnvNetwork NetworkEnv{..} extraSecrets =
  Client.MorleyOnlyRpcEnv
  { moreLogAction = Client.mkLogAction neVerbosity
  , moreClientEnv = Client.mceClientEnv neMorleyClientEnv
  , moreSecretKeys = mconcat
      [ one (mkKeyAddress (toPublic sk), sk)
      | sk <- maybe id (:) neSecretKey extraSecrets ]
  }

-- | This error designates that necessary preparations for running tests
-- are not made.
data MoneybagConfigurationException
  = NoMoneybagAddress ImplicitAlias
  | TwoMoneybagKeys ImplicitAlias SecretKey ImplicitAddress
  deriving stock (Generic, Show, Eq)

instance Buildable MoneybagConfigurationException where
  build = \case
    NoMoneybagAddress alias -> unlinesF @_ @Doc
      [ "Moneybag alias is not registered in the tezos node: " <> build alias
      , ""
      , "Cleveland's network tests require a special address with plenty of XTZ for"
      , "originating contracts and performing transfers."
      , ""
      , "By default, Cleveland expects an account with the alias 'moneybag' to already exist."
      , "If no such alias exists, you can choose to either:"
      , "  * Use a different alias, supplied via '--cleveland-moneybag-alias'."
      , "  * Import a moneybag account, by supplying its secret key via '--cleveland-moneybag-secret-key'."
      ]
    TwoMoneybagKeys alias envKey existingAddress -> unlinesF @_ @Doc
      [ "Tried to import the secret key supplied via '--cleveland-moneybag-secret-key' and"
      , "associate it with the alias '" +| alias |+ "', but the alias already exists."
      , ""
      , "  --cleveland-moneybag-secret-key: " <> build envKey
      , "  Existing address             : " <> build existingAddress
      , ""
      , "Possible fix:"
      , "  * If you wish to use the existing address, please remove the '--cleveland-moneybag-secret-key' option."
      , "  * Otherwise, please supply a different alias via '--cleveland-moneybag-alias'."
      ]


instance Exception MoneybagConfigurationException where
  displayException = pretty
  fromException = fromPossiblyAnnotatedException

data ClientState = ClientState
  { csDefaultAliasCounter :: DefaultAliasCounter
  , csRefillableAddresses :: Set ImplicitAddress
  , csMoneybagAddress :: Moneybag
  }

newtype ClientM a = ClientM
  { unClientM :: ReaderT (IORef ClientState) IO a
  }
  deriving newtype (Functor, Applicative, Monad, MonadIO,
                    MonadThrow, MonadCatch, MonadReader (IORef ClientState), MonadFail)

data InternalNetworkScenarioError = TooManyRefillIterations Word ImplicitAddress
  deriving stock (Show)

instance Buildable InternalNetworkScenarioError where
  build (TooManyRefillIterations iter addr) =
    "Too many (" +| iter |+ ") refill iteratons of " +| addr |+ ""

instance Exception InternalNetworkScenarioError where
  displayException = pretty
  fromException = fromPossiblyAnnotatedException

runNetworkT :: NetworkEnv -> NetworkT ClientM a -> IO a
runNetworkT env scenario = do
  disableAlphanetWarning
  moneybagAddr <- setupMoneybagAddress env
  let caps = NetworkCaps
        { ncNetworkEnv = env
        , ncClevelandCaps = ClevelandCaps
            { ccSender = Sender $ unMoneybag moneybagAddr
            , ccMoneybag = moneybagAddr
            , ccMiscCap = networkMiscImpl env
            , ccOpsCap = networkOpsImpl (neMorleyClientEnv env)
            }
        }
  ist <- newIORef ClientState
    { csDefaultAliasCounter = DefaultAliasCounter 0
    , csRefillableAddresses = Set.empty
    , csMoneybagAddress = moneybagAddr
    }
  let clientM = runReaderT scenario caps
  runReaderT (unClientM clientM) ist

-- | Initialize @moneybag@ address by given 'NetworkEnv'
setupMoneybagAddress :: NetworkEnv -> IO Moneybag
setupMoneybagAddress NetworkEnv{..} = do
  let setupEnv = neMorleyClientEnv &
        if neExplicitDataDir
        then id
        else Client.mceTezosClientL . tceMbTezosClientDataDirL .~ Nothing
  storageAddress <- runMorleyClientM setupEnv $
    Client.resolveAddressMaybe (AddressAlias neMoneybagAlias)
  Moneybag <$> case (neSecretKey, storageAddress) of
    (Nothing, Just addr) -> do
      unless neExplicitDataDir do
        ek <- runMorleyClientM setupEnv $
          retryOnceOnTimeout $ TezosClient.getSecretKey neMoneybagAlias
        void $ runMorleyClientM neMorleyClientEnv $
          failOnTimeout ... TezosClient.importKey False neMoneybagAlias ek
      pure $ Client.AddressWithAlias addr neMoneybagAlias
    (Nothing, Nothing) -> throwM $ NoMoneybagAddress neMoneybagAlias
    (Just ek, Just sa)
      | mkKeyAddress (toPublic ek) == sa -> do
          unless neExplicitDataDir $ void $
            runMorleyClientM neMorleyClientEnv $
              failOnTimeout ... TezosClient.importKey False neMoneybagAlias ek
          pure $ Client.AddressWithAlias sa neMoneybagAlias
      | otherwise -> throwM $ TwoMoneybagKeys neMoneybagAlias ek sa
    (Just ek, Nothing) -> do
      runMorleyClientM neMorleyClientEnv $
        failOnTimeout ... TezosClient.importKey False neMoneybagAlias ek
      return $ Client.AddressWithAlias (mkKeyAddress (toPublic ek)) neMoneybagAlias

-- | Implementation that works with real network and uses @octez-node@
-- RPC and @octez-client@.
networkOpsImpl :: MorleyClientEnv -> Sender -> ClevelandOpsImpl ClientM
networkOpsImpl env (Sender sender) =
  mapClevelandOpsImplExceptions (addCallStack . exceptionHandler env)
    ClevelandOpsImpl
    { coiRunOperationBatch = runOperationBatch env sender
    }

networkMiscImpl :: NetworkEnv -> ClevelandMiscImpl ClientM
networkMiscImpl env@NetworkEnv{..} =
  mapClevelandMiscImplExceptions (addCallStack . exceptionHandler neMorleyClientEnv) ClevelandMiscImpl
    { cmiRunIO = liftIO

    , cmiOriginateLargeUntyped =
      \(Sender sender) OriginateData{..} -> do
        let originationScenario :: Client.MorleyClientM (OperationHash, ContractAddress)
            originationScenario = case odContractAndStorage of
              ODContractAndStorageUntyped{..} ->
                Client.originateLargeUntypedContract OverwriteDuplicateAlias odName
                  sender odBalance uodContract uodStorage Nothing odDelegate
              ODContractAndStorageTyped{..} | T.Contract{} <- todContract ->
                Client.originateLargeContract OverwriteDuplicateAlias odName
                  sender odBalance todContract (T.toVal todStorage)
                    Nothing odDelegate
              ODContractAndStorageLorentz{..} | L.Contract{} <- odContract ->
                Client.lOriginateLargeContract OverwriteDuplicateAlias odName
                  sender odBalance odContract odStorage Nothing odDelegate
        -- Note that tezos key reveal operation cost an additional fee
        -- so that's why we reveal keys in origination and transaction
        -- rather than doing it before scenario execution
        (_ :: OperationHash, res) <- liftIO $ do
          revealKeyUnlessRevealed neMorleyClientEnv sender
          runMorleyClientM neMorleyClientEnv originationScenario
        comment $ "Originated large smart contract " +| odName |+ " with address " <> pretty res
        pure res

    , cmiSignBytes = \hash signer -> liftIO $ runMorleyClientM neMorleyClientEnv $
        -- We don't use password protected accounts in cleveland tests
        Client.signBytes signer Nothing hash

    , cmiGenKey = \sodAlias -> do
        alias <- resolveSpecificOrDefaultAlias sodAlias
        liftIO $ runMorleyClientM neMorleyClientEnv $ Client.genKey alias

    , cmiImportKey = \key sodAlias -> do
        alias <- resolveSpecificOrDefaultAlias sodAlias
        liftIO $ runMorleyClientM neMorleyClientEnv $ TezosClient.importKey True alias key

    , cmiGenFreshKey = \sodAlias -> do
        alias <- resolveSpecificOrDefaultAlias sodAlias
        liftIO $ runMorleyClientM neMorleyClientEnv $ Client.genFreshKey alias

    , cmiGetBalance = foldConstrained getBalanceHelper
    , cmiGetChainId = liftIO $ runMorleyClientM neMorleyClientEnv Client.getChainId
    , cmiAttempt = try
    , cmiThrow = throwM
    , cmiMarkAddressRefillable = setAddressRefillable
    , cmiComment = comment
    , cmiUnderlyingImpl = pure $ Right env
    , cmiFailure = clientFailure
    , ..
    }
  where
    cmiTicketBalance
      :: forall t. (T.ForbidOp t, T.SingI t)
      => L1Address -> ContractAddress -> T.Value t -> ClientM Natural
    cmiTicketBalance owner ticketer value = liftIO . runMorleyClientM neMorleyClientEnv $
      Client.getTicketBalance owner Client.GetTicketBalance
        { Client.gtbTicketer = ticketer
        , Client.gtbContent = toExpression $ T.untypeValueHashable value
        , Client.gtbContentType = toExpression $ T.starNotes @t
        }

    cmiAllTicketBalances :: ContractAddress -> ClientM [SomeTicket]
    cmiAllTicketBalances owner = do
      tickets <- liftIO $ runMorleyClientM neMorleyClientEnv $ Client.getAllTicketBalances owner
      forM tickets \Client.GetAllTicketBalancesResponse{..} -> do
        uTy <- fromExpression gatbrContentType & \case
          Right ty -> pure ty
          Left err -> clientFailure $ unlinesF
            [ "Failed to decode ticket value:"
            , "Value:"
            , indentF 2 $ build (J.encodeToTextBuilder gatbrContentType)
            , "Decoding error:"
            , indentF 2 $ build err
            ]
        T.withUType uTy \(_ :: T.Notes t) -> do
          case fromExpression @(T.Value t) gatbrContent of
            Right val ->
              pure @ClientM $ SomeTicket T.Ticket
                { tTicketer = Constrained gatbrTicketer
                , tData = val
                , tAmount = unStringEncode gatbrAmount
                }
            Left err -> clientFailure $ unlinesF
                [ "Failed to decode ticket value:"
                , "Value:"
                , indentF 2 $ build (J.encodeToTextBuilder gatbrContentType)
                , "Decoding error:"
                , indentF 2 $ build err
                ]

    cmiGetBigMapValueMaybe :: (NicePackedValue k, NiceUnpackedValue v) => BigMapId k v -> k -> ClientM (Maybe v)
    cmiGetBigMapValueMaybe bigMapId k =
      liftIO . runMorleyClientM neMorleyClientEnv $ Client.readBigMapValueMaybe bigMapId k

    cmiGetAllBigMapValuesMaybe :: (NiceUnpackedValue v) => BigMapId k v -> ClientM (Maybe [v])
    cmiGetAllBigMapValuesMaybe bigMapId =
      liftIO . runMorleyClientM neMorleyClientEnv $ Client.readAllBigMapValuesMaybe bigMapId

    getStorageType :: Expression -> ClientM U.Ty
    getStorageType contractExpr = do
      let
        storageTypeExprMb = contractExpr ^?
          _ExpressionSeq
          . each
          . _ExpressionPrim
          . filtered (\prim -> prim ^. mpaPrimL == Prim_storage)
          . mpaArgsL
          . _head

      case storageTypeExprMb of
        Nothing -> clientFailure $ unlinesF
          [ "Contract expression did not contain a 'storage' expression."
          , "Contract expression:"
          , indentF 2 $ build (J.encodeToTextBuilder contractExpr)
          ]
        Just storageTypeExpr ->
          case fromExpression @U.Ty storageTypeExpr of
            Left err -> clientFailure $ unlinesF
              [ "'storage' expression was not a valid type expression."
              , "Storage expression:"
              , indentF 2 $ build (J.encodeToTextBuilder storageTypeExpr)
              , "Decoding error:"
              , indentF 2 $ build err
              ]
            Right storageType -> pure storageType

    cmiGetSomeStorage :: ContractAddress -> ClientM SomeAnnotatedValue
    cmiGetSomeStorage addr = do
      OriginationScript {osCode, osStorage} <-
        liftIO . runMorleyClientM neMorleyClientEnv $ Client.getContractScript addr
      storageType <- getStorageType osCode
      T.withUType storageType \(storageNotes :: T.Notes t) -> do
        withDict (rpcSingIEvi @t) do
          case fromExpression @(T.Value (TAsRPC t)) osStorage of
            Right storageValueRPC ->
              pure $ SomeAnnotatedValue (notesAsRPC storageNotes) storageValueRPC
            Left err ->
              clientFailure $ unlinesF
                [ "Failed to decode storage expression."
                , "Storage expression:"
                , indentF 2 $ build (J.encodeToTextBuilder osStorage)
                , "Decoding error:"
                , indentF 2 $ build err
                ]

    cmiResolveAddress :: AddressOrAlias kind -> ClientM (Client.AddressWithAlias kind)
    cmiResolveAddress =
      liftIO . runMorleyClientM neMorleyClientEnv . Client.resolveAddressWithAlias

    cmiGetPublicKey :: Client.ImplicitAddressWithAlias -> ClientM PublicKey
    cmiGetPublicKey =
      liftIO . runMorleyClientM neMorleyClientEnv . retryOnceOnTimeout . Client.getPublicKey

    getBalanceHelper :: L1AddressKind kind => KindedAddress kind -> ClientM Mutez
    getBalanceHelper = liftIO . runMorleyClientM neMorleyClientEnv . Client.getBalance

    cmiGetDelegate :: L1Address -> ClientM (Maybe KeyHash)
    cmiGetDelegate =
      liftIO . runMorleyClientM neMorleyClientEnv . Client.getDelegate

    cmiGetNow :: ClientM Tezos.Timestamp
    cmiGetNow = timestampFromUTCTime <$> getLastBlockTimestamp neMorleyClientEnv

    cmiGetLevel :: ClientM Natural
    cmiGetLevel = getLastBlockLevel neMorleyClientEnv

    cmiGetApproximateBlockInterval :: ClientM (Time Second)
    cmiGetApproximateBlockInterval = liftIO $ do
      pp <- runMorleyClientM neMorleyClientEnv $ Client.getProtocolParameters
      return . sec $ (unStringEncode $ ppMinimalBlockDelay pp) % 1

    cmiAdvanceTime :: (KnownDivRat unit Second) => Time unit -> ClientM ()
    cmiAdvanceTime delta = do
      let
        -- Round 'delta' to the nearest second, not smaller than 'delta'.
        -- A chain's time resolution is never smaller than a second,
        -- so if 'delta' is 0.1s, we actually need to wait at least 1s.
        deltaSec :: Time Second
        deltaSec = ceilingUnit $ toUnit @Second delta

        deltaSec' :: NominalDiffTime
        deltaSec' = secondsToNominalDiffTime $ toNum @Second deltaSec
      t0 <- getLastBlockTimestamp neMorleyClientEnv
      threadDelay deltaSec
      let
        go :: ClientM ()
        go = do
          now <- getLastBlockTimestamp neMorleyClientEnv
          if (now `diffUTCTime` t0) >= deltaSec'
            then pass
            else threadDelay (sec 1) >> go
      go

    cmiAdvanceToLevel :: (Natural -> Natural) -> ClientM ()
    cmiAdvanceToLevel targetLevelFn = do
      lastLevel <- getLastBlockLevel neMorleyClientEnv
      let targetLevel = max (targetLevelFn lastLevel) lastLevel
      let skippedLevels = targetLevel - lastLevel
      -- In case we need to skip more than one level we'll jump ahead for
      -- 'cmiGetApproximateBlockInterval' for 'skippedLevels - 1' times.
      -- This way we are sure we won't end up in the middle (or towards the end)
      -- of the target level.
      when (skippedLevels > 0) $ do
        when (skippedLevels > 1) $ do
          minBlockInterval <- cmiGetApproximateBlockInterval
          let waitTime = (skippedLevels - 1) * toNum @Second minBlockInterval
          threadDelay . sec $ waitTime % 1
        -- A chain's time resolution is never smaller than a second, so with (less
        -- than) a level to go we can wait for 1s in loop until we reach the target.
        let go :: ClientM ()
            go = do
              curLevel <- cmiGetLevel
              when (targetLevel > curLevel) $ threadDelay (sec 1) >> go
        go

    cmiRunCode
      :: forall cp st vd. (HasRPCRepr st, T.IsoValue (AsRPC st))
      => Sender -> RunCode cp st vd -> ClientM (AsRPC st)
    cmiRunCode (Sender sender) (RunCode rcContract rcParameter rcStorage rcAmount rcLevel rcNow rcBalance rcSource) = do
      liftIO $ runMorleyClientM neMorleyClientEnv do
        -- Pattern match on the contract constructor to reveal
        -- a proof of `NiceStorage st`
        L.Contract{} <- pure rcContract
        T.fromVal <$> Client.runContract Client.RunContractParameters
          { rcpContract = L.toMichelsonContract $ rcContract
          , rcpParameter = rcParameter
          , rcpStorage = rcStorage
          , rcpAmount = rcAmount
          , rcpBalance = rcBalance
          , rcpSource = rcSource
          , rcpLevel = rcLevel
          , rcpNow = rcNow
          , rcpSender = Just $ Client.awaAddress sender
          }
      `catch` \case
        err@(RPC.RunCodeErrors errs)
          | Just clientErr <- runErrorsToClientError errs
          -> throwM =<< exceptionToTransferFailure neMorleyClientEnv [] clientErr
          | otherwise -> throwM err

clientFailure :: Doc -> ClientM a
clientFailure = throwM . CustomTestError . pretty

comment :: Text -> ClientM ()
comment msg = liftIO $ putTextLn msg >> hFlush stdout

getAlias
  :: L1AddressKind kind
  => MorleyClientEnv
  -> KindedAddress kind
  -> ClientM (Alias kind)
getAlias env = liftIO . runMorleyClientM env . TezosClient.getAlias . AddressResolved

getAliasMaybe
  :: L1AddressKind kind
  => MorleyClientEnv
  -> KindedAddress kind
  -> ClientM (Maybe (Alias kind))
getAliasMaybe env = liftIO . runMorleyClientM env . Client.getAliasMaybe . AddressResolved

getLastBlockTimestamp :: MorleyClientEnv -> ClientM UTCTime
getLastBlockTimestamp env = liftIO $
  bhnhTimestamp . bcHeader <$> runMorleyClientM env (Client.getBlockConstants HeadId)

getLastBlockLevel :: MorleyClientEnv -> ClientM Natural
getLastBlockLevel env = do
  bc <- liftIO $ runMorleyClientM env (Client.getBlockConstants HeadId)
  pure . Unsafe.fromIntegral @Int64 @Natural . bhnhLevel $ bcHeader bc

runOperationBatch
  :: MorleyClientEnv
  -> Client.ImplicitAddressWithAlias
  -> [OperationInfo ClevelandInput]
  -> ClientM [OperationInfo ClevelandResult]
runOperationBatch env sender ops = do
  istRef <- ask
  ClientState{csMoneybagAddress=Moneybag moneybag} <- readIORef istRef
  -- Note that tezos key reveal operation cost an additional fee
  -- so that's why we reveal keys in origination and transaction
  -- rather than doing it before scenario execution
  liftIO $ revealKeyUnlessRevealed env sender

  ops' <- forM ops \case
    OpOriginate (SomeOriginateData OriginateData{..}) -> do
      SomeContractAndStorage odContract odStorage <- either throwM pure $
        typeCheckODContractAndStorageIfNeeded odContractAndStorage
      pure $ OpOriginate Client.OriginationData
        { odAliasBehavior = OverwriteDuplicateAlias
        , odMbFee = Nothing
        , ..
        }
    OpTransfer TransferData{..} ->
      pure $ OpTransfer $ Client.TransactionData Client.TD
        { tdReceiver = toL1Address tdTo
        , tdAmount = tdAmount
        , tdEpName = tdEntrypoint
        , tdParam = toVal tdParameter
        , tdMbFee = Nothing
        }
    OpTransferTicket TransferTicketData{..}
      | T.VTicket ttdTicketTicketer ttdTicketContents ttdTicketAmount <- ttdParameter
      , T.Dict <- T.comparableImplies ttdTicketContents
      -> pure $ OpTransferTicket Client.TransferTicketData
        { ttdDestination = L.toAddress $ toL1Address ttdTo
        , ttdMbFee = Nothing
        , ..
        }
    OpReveal key ->
      pure $ OpReveal RevealData
        { rdPublicKey = key
        , rdMbFee = Nothing
        }
    OpDelegation delegate ->
      pure $ OpDelegation DelegationData
        { ddDelegate = delegate
        , ddMbFee = Nothing
        }

  let refill :: Word -> Client.MorleyClientM Word
      refill iter = do
        void $ dryRunOperations sender ops'
        pure iter
        `catch` \errs -> do
          when (iter > 3) $ throwM $ TooManyRefillIterations iter $ Client.awaAddress sender
          realBalance <- Client.getBalance $ Client.awaAddress sender
          let handleRunErrors errs'
                | Just (arg #balance -> balance, arg #required -> required)
                    <- findBalanceTooLow errs' = do
                    logInfo $ sender |+ " balance of " +| realBalance |+ " \n\
                              \is too low, need " +| required |+ ", but got " +| balance |+ ""
                    let reportedDifference = unsafeSubMutez required balance
                        -- required >= balance should always be true if we got 'BalanceTooLow'
                    if iter == 0
                    -- on first iteration, we dry-run the transaction as moneybag (if possible)
                    -- and esitmate the required balance that way;
                    then approximateRequired realBalance
                            `catch` \(_ :: SomeException) -> pure reportedDifference
                    -- on subsequent iterations (which run only if the first
                    -- wasn't enough), we rely on the reported `required` and `balance`,
                    -- NOTE: BalanceTooLow can be thrown either before fees are subtracted, or after.
                    -- In the former case, (required - balance == transfer_amount - real_balance)
                    -- In the latter case, (required - balance == transfer_amount - real_balance - fees)
                    -- Notice that fees are only included if real_balance >= transfer_amount.
                    -- Consequently, if transfer_amount > real_balance AND
                    -- transfer_amount + fees > real_balance + minimalMutez, the amount we transfer here
                    -- will be insufficient. TL;DR, it doesn't work for large transfer_amounts.
                    -- For batched transfers, this gets a bit more complicated, but the same principle
                    -- applies; unless total fees can exceed minimalMutez, or total transfer_amount
                    -- is large, it should work without looping.
                    else pure reportedDifference
                | findCantPayStorageFee errs' = do
                    logInfo $ sender |+ " balance of " +| realBalance |+ "\n\
                              \ is too low to pay storage fee"
                    -- since no required balance is reported, there is no choice
                    approximateRequired realBalance
                      -- if running as moneybag failed for some reason, just throw in some tez
                      -- and hope for the best
                      `catch` \(_ :: SomeException) -> pure minimalMutez
                | otherwise = throwM errs
          amount <- max minimalMutez . addSafetyMutez <$> if
            | Just (Client.UnexpectedRunErrors err) <- fromException errs -> handleRunErrors err
            | Just (RPC.RunCodeErrors err) <- fromException errs -> handleRunErrors err
            | otherwise -> throwM errs
          logInfo $ "Will transfer " +| amount |+ " from " +| moneybag |+ ""
          void $ Client.lTransfer moneybag (Client.awaAddress sender) amount U.DefEpName () Nothing
          refill (iter + 1) -- loop
      addSafetyMutez x = fromMaybe x $ addMutez x safetyMutez
      minimalMutez = 5e5
      safetyMutez = 100
      safetyStorage = 20
      approximateRequired balance = do
        -- dry-run as moneybag and estimate cost+burn+fees
        (appliedResults, fees) <- unzip <$> dryRunOperations moneybag ops'
        ProtocolParameters{..} <- Client.getProtocolParameters
        -- uses quite a bit of unsafe mutez arithmetic, but arguably
        -- if we end up running into overflow while computing the
        -- required balance, then we couldn't run these operations
        -- anyway.
        let totalFees = unsafeSumMutez fees
            unsafeSumMutez = foldr unsafeAddMutez zeroMutez
            zeroMutez = 0
            originationSz = Unsafe.fromIntegral @Int @Natural ppOriginationSize
            (opsSum, originationSize) = bimap unsafeSumMutez sum . unzip
              $ map opcostAndOriginationCount ops
            costPerByte = unTezosMutez ppCostPerByte
            opcostAndOriginationCount = \case
              OpOriginate (SomeOriginateData od) -> (odBalance od, originationSz)
              OpTransfer td -> (tdAmount td, 0)
              OpTransferTicket _ -> (zeroMutez, 0)
              OpReveal _ -> (zeroMutez, 0)
              OpDelegation _ -> (zeroMutez, 0)
            storageDiff AppliedResult{..} = safetyStorage + Unsafe.fromIntegral @TezosInt64 @Natural arPaidStorageDiff
            storageBurnInBytes = originationSize + sum (map storageDiff appliedResults)
            storageBurnInMutez = unsafeMulMutez costPerByte storageBurnInBytes
            required = opsSum `unsafeAddMutez` totalFees `unsafeAddMutez` storageBurnInMutez
        logInfo $ "estimated amount needed is " +| required |+ ", but got " +| balance |+ "\n\
                  \Storage size: " +| storageBurnInBytes |+ "; Operations cost: " +| opsSum |+ "\n\
                  \Fees: " +| totalFees |+ "; Storage burn cost: " +| storageBurnInMutez |+ ""
        pure $ fromMaybe zeroMutez $ subMutez required balance

  refillable <- isAddressRefillable $ Client.awaAddress sender
  results <- liftIO $ runMorleyClientM env $ do
    when refillable $ do
      tookIters <- refill 0
      when (tookIters > 1) $ logWarning $
        "Refill of " +| sender |+ " took " +| tookIters |+ " iterations."
    snd <$> Client.runOperations sender ops'

  forM_ results $ \case
    OpTransfer _ -> pass
    OpTransferTicket _ -> pass
    OpOriginate addr -> do
      alias <- getAlias env addr
      comment $ "Originated smart contract '" +| alias |+
        "' with address " <> pretty addr
    OpReveal () -> pass
    OpDelegation () -> pass

  mapM toClevelandResult results

toClevelandResult :: OperationInfo Result -> ClientM (OperationInfo ClevelandResult)
toClevelandResult = \case
  OpTransfer ops -> OpTransfer <$> mapM intOpEventToContractEvent ops
  OpTransferTicket ops -> OpTransferTicket <$> mapM intOpEventToContractEvent ops
  OpOriginate ops -> pure $ OpOriginate ops
  OpReveal ops -> pure $ OpReveal ops
  OpDelegation ops -> pure $ OpDelegation ops

intOpEventToContractEvent :: WithSource EventOperation -> ClientM ContractEvent
intOpEventToContractEvent WithSource{wsOtherData = EventOperation{..}, ..} = do
  T.AsUType (ceType :: T.Notes t) <- either throwM (pure . toTy) $ fromExpression @U.T eoType
  cePayload <- case eoPayload of
    Nothing -> pure Nothing
    Just payload -> case fromExpression @(T.Value t) . toExpression $ payload of
      Right value -> pure . Just $ SomeAnnotatedValue ceType value
      Left err ->
        clientFailure $ unlinesF
          [ "Failed to decode event payload expression."
          , "Payload expression:"
          , indentF 2 $ build (J.encodeToTextBuilder $ toExpression payload)
          , "Decoding error:"
          , indentF 2 $ build err
          ]
  ceSource :: ContractAddress <- case wsSource of
    Constrained ceSource@ContractAddress{} -> pure ceSource
    Constrained (addr :: KindedAddress kind) ->
      throwM $ CustomTestError $ "Unexpected event source kind: " <> pretty (demote @kind)
        \\ addressKindSanity addr
  pure $ ContractEvent
    { ceTag = maybe "" unMText eoTag
    , ..
    }
  where
    toTy :: U.T -> U.Ty
    toTy t = U.Ty t U.noAnn


dryRunOperations :: Client.ImplicitAddressWithAlias
                  -> [OperationInfo Client.ClientInput]
                  -> Client.MorleyClientM [(AppliedResult, Mutez)]
dryRunOperations s = \case
  [] -> return []
  (x:xs) -> toList . map (second unTezosMutez) <$> Client.dryRunOperationsNonEmpty s (x :| xs)

findBalanceTooLow :: [Client.RunError] -> Maybe ("balance" :! Mutez, "required" :! Mutez)
-- we really shouldn't get several errors of the same type here, so find only the first one
findBalanceTooLow
  (Client.BalanceTooLow balance required:_)
  = Just (balance, required)
findBalanceTooLow (_:xs) = findBalanceTooLow xs
findBalanceTooLow [] = Nothing

findCantPayStorageFee :: [Client.RunError] -> Bool
-- we really shouldn't get several errors of the same type here, so find only the first one
findCantPayStorageFee
  (Client.CantPayStorageFee:_)
  = True
findCantPayStorageFee (_:xs) = findCantPayStorageFee xs
findCantPayStorageFee [] = False

resolveAddressAndAlias :: MorleyClientEnv -> KindedAddress kind -> ClientM AddressAndAlias
resolveAddressAndAlias env addr =
   case addr of
     ImplicitAddress{} -> AddressAndAlias addr <$> getAliasMaybe env addr
     ContractAddress{} -> AddressAndAlias addr <$> getAliasMaybe env addr
     SmartRollupAddress{} -> pure (AddressAndAlias addr Nothing)

exceptionToTransferFailure
  :: MorleyClientEnv
  -> [RPC.OperationResp RPC.WithSource]
  -> RPC.ClientRpcError
  -> ClientM TransferFailure
exceptionToTransferFailure env stack = \case
    RPC.ContractFailed addr expr -> mkTransferFailure addr $ FailedWith (EOTVExpression expr) Nothing
    RPC.BadParameter (MkAddress addr) _ -> mkTransferFailure addr BadParameter
    RPC.EmptyTransaction addr -> mkTransferFailure addr EmptyTransaction
    RPC.ShiftOverflow addr -> mkTransferFailure addr ShiftOverflow
    RPC.GasExhaustion addr -> mkTransferFailure addr GasExhaustion
    internalError -> throwM internalError
    where
      mkTransferFailure :: KindedAddress kind -> TransferFailureReason -> ClientM TransferFailure
      mkTransferFailure addr e = do
        addr' <- resolveAddressAndAlias env addr
        pure $ TransferFailure addr' (rpcToCallSeq stack) e

rpcToCallSeq :: [RPC.OperationResp RPC.WithSource] -> CallSequence
rpcToCallSeq = toCallSeq \case
  RPC.TransactionOpResp RPC.WithSource{wsOtherData=to@RPC.TransactionOperation{..},..} ->
    Just . (wsSource,) . CallSequenceOp to . CallSequence <$> popToCallSeq toDestination
  RPC.TransferTicketOpResp RPC.WithSource{wsOtherData=to@RPC.TransferTicketOperation{..},..} ->
    Just . (wsSource,) . CallSequenceOp to . CallSequence <$> popToCallSeq ttoDestination
  RPC.OriginationOpResp RPC.WithSource{..} ->
    pure $ Just (wsSource, CallSequenceOp wsOtherData mempty)
  RPC.DelegationOpResp RPC.WithSource{..} ->
    pure $ Just (wsSource, CallSequenceOp wsOtherData mempty)
  RPC.RevealOpResp RPC.WithSource{..} ->
    pure $ Just (wsSource, CallSequenceOp wsOtherData mempty)
  RPC.EventOpResp RPC.WithSource{..} ->
    pure $ Just (wsSource, CallSequenceOp wsOtherData mempty)
  RPC.OtherOpResp _ ->
    pure Nothing


exceptionHandler :: MorleyClientEnv -> ClientM a -> ClientM a
exceptionHandler env action = action `catch` \case
  se@(SomeException e)
    | Just err <- cast e -> throwM =<< exceptionToTransferFailure env [] err
    | Just RPC.ClientRpcErrorWithStack{..} <- cast e
    -> throwM =<< exceptionToTransferFailure env (toList crewsStack) crewsError
    | otherwise -> throwM se

resolveSpecificOrDefaultAlias :: SpecificOrDefaultAlias -> ClientM ImplicitAlias
resolveSpecificOrDefaultAlias (SpecificAlias alias) = pure alias
resolveSpecificOrDefaultAlias (DefaultAlias) = do
  stateRef <- ask
  ist@ClientState{csDefaultAliasCounter=DefaultAliasCounter counter} <- readIORef stateRef
  writeIORef stateRef ist{ csDefaultAliasCounter = DefaultAliasCounter $ counter + 1 }
  pure $ mkDefaultAlias counter

setAddressRefillable :: ImplicitAddress -> ClientM ()
setAddressRefillable addr = do
  stRef <- ask
  modifyIORef stRef $ \st@ClientState{..} ->
    st{csRefillableAddresses=Set.insert addr csRefillableAddresses}

isAddressRefillable :: ImplicitAddress -> ClientM Bool
isAddressRefillable addr = do
  stRef <- ask
  Set.member addr . csRefillableAddresses <$> readIORef stRef

----------------------------------------------------------------------------
-- Helpers
----------------------------------------------------------------------------

-- | Runs 'Client.revealKeyUnlessRevealed' with given client environment.
revealKeyUnlessRevealed :: MorleyClientEnv -> Client.ImplicitAddressWithAlias -> IO ()
revealKeyUnlessRevealed env addr = runMorleyClientM env $
  Client.revealKeyUnlessRevealed addr

----------------------------------------------------------------------------
-- Validation
----------------------------------------------------------------------------

-- | Signals an assertion failure during the execution of an action.
newtype TestError
  = CustomTestError Text
  deriving stock Show

instance Exception TestError where
  displayException = pretty
  fromException = fromPossiblyAnnotatedException

instance Buildable TestError where
  build (CustomTestError msg) = build msg