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

{-# OPTIONS_HADDOCK not-home #-}

-- | The bulk of Cleveland actions.
module Test.Cleveland.Internal.Actions.Misc
  ( module Test.Cleveland.Internal.Actions.Misc
  , noViews
  ) where

import Data.Singletons (demote)
import Fmt (Buildable, Builder, build, indentF, unlinesF, (+|), (|+))
import Time (KnownDivRat, Second, Time, toNum)
import Unsafe qualified (fromIntegral)

import Lorentz (BigMapId, Contract(..), DemoteViewsDescriptor, IsoValue, noViews)
import Lorentz.Bytes
import Lorentz.Constraints
import Morley.AsRPC (HasRPCRepr(..))
import Morley.Client (MorleyClientEnv, OperationInfo(..), RunError(..), UnexpectedErrors(..))
import Morley.Michelson.Runtime (ExecutorError'(..), VotingPowers)
import Morley.Michelson.Runtime.GState (GStateUpdateError(..))
import Morley.Michelson.Runtime.Import qualified as Runtime
import Morley.Michelson.Typed (SomeAnnotatedValue)
import Morley.Michelson.Typed qualified as T
import Morley.Michelson.Typed.AnnotatedValue (castTo, getT, value)
import Morley.Michelson.Untyped qualified as U
import Morley.Tezos.Address (unImplicitAddress)
import Morley.Tezos.Address.Alias (Alias(..))
import Morley.Tezos.Core (ChainId, Mutez, Timestamp)
import Morley.Tezos.Crypto (KeyHash, PublicKey, Signature)
import Morley.Util.SizedList qualified as SL
import Morley.Util.SizedList.Types
import Test.Cleveland.Instances ()
import Test.Cleveland.Internal.Abstract
import Test.Cleveland.Internal.Actions.Assertions
import Test.Cleveland.Internal.Actions.Helpers
import Test.Cleveland.Internal.Actions.MonadOps
import Test.Cleveland.Internal.Actions.Transfer
import Test.Cleveland.Internal.Exceptions (fromPossiblyAnnotatedException)
import Test.Cleveland.Lorentz.Import qualified as LorentzImport
import Test.Cleveland.Lorentz.Types

{-# ANN module ("HLint: ignore Avoid lambda using `infix`" :: Text) #-}

-- $setup
-- >>> :m +Morley.Util.SizedList.Types
-- >>> :{
-- isEquivalentTo :: Show a => a -> a -> Bool
-- isEquivalentTo a b = Debug.show a == Debug.show b
-- infix 0 `isEquivalentTo`
-- :}

-- | Update the current sender on whose behalf transfers and originations are
-- invoked.
withSender :: MonadCleveland caps m => ImplicitAddress -> m a -> m a
withSender addr =
  local (set senderL (Sender addr))

-- | Update the current moneybag that transfers money on the newly created
-- addresses. For the rare occasions when this is necessary.
withMoneybag :: MonadCleveland caps m => ImplicitAddress -> m a -> m a
withMoneybag addr =
  local (set moneybagL (Moneybag addr))

-- | Runs an 'IO' action.
runIO :: (HasCallStack, MonadCleveland caps m) => IO res -> m res
runIO io = do
  withCap getMiscCap \cap -> cmiRunIO cap io

-- | Get the address of the implicit account / contract associated with the given alias.
resolveAddress
  :: (HasCallStack, MonadCleveland caps m)
  => Alias kind -> m (KindedAddress kind)
resolveAddress alias = do
  withCap getMiscCap \cap -> cmiResolveAddress cap alias

-- | Simple combinator that marks address as "refillable".
--
-- If a refillable address lacks funds for the next operation,
-- some funds will automatically be transferred to it.
refillable :: MonadCleveland caps m => m ImplicitAddress -> m ImplicitAddress
refillable action = do
  addr <- action
  withCap getMiscCap \cap -> cmiMarkAddressRefillable cap addr
  pure addr

-- | If the given alias is already associated with an existing address,
-- that address will be reused and returned.
-- Otherwise, generate a new secret key and record it with given alias.
--
-- If the account has too low of a balance, a small amount of XTZ will
-- be transferred to it.
--
-- Notes:
--
-- * By default, the XTZ is transferred from the account associated with the @moneybag@ alias.
--   This can be overridden with the @--cleveland-moneybag-alias@ command line option, the
--   @TASTY_CLEVELAND_MONEYBAG_ALIAS@ env var, or 'withMoneybag'.
newAddress :: (HasCallStack, MonadCleveland caps m) => SpecificOrDefaultAlias -> m ImplicitAddress
newAddress alias = do
  addrs <- newAddresses $ alias :< Nil
  case addrs of
    addr :< Nil -> pure addr

-- | Batched version of `newAddress`
newAddresses
  :: forall n n' caps m.
    (HasCallStack, MonadCleveland caps m, IsoNatPeano n n')
  => SizedList n SpecificOrDefaultAlias
  -> m (SizedList n ImplicitAddress)
newAddresses aliases = do
  addrs <- withCap getMiscCap \cap -> traverse (cmiGenKey cap) aliases
  Moneybag moneybag <- view moneybagL

  -- Addresses may exist from previous scenarios runs and have sufficient
  -- balance for the sake of testing; if so, we can save some time
  balances <- traverse getBalance addrs
  withSender moneybag do
    inBatch do
      sequenceA_ $
        SL.zipWith refillIfLowBalance addrs balances

  pure addrs
  where
    refillIfLowBalance addr balance =
      when (balance < 0.5_e6) do  -- < 0.5 XTZ
        transfer addr (0.9_e6 :: Mutez) -- 0.9 XTZ

-- | Generate a new secret key and record it with given alias. If the
-- alias is already known, the key will be overwritten. The address is
-- guaranteed to be fresh, i. e. no operations on it have been made.
newFreshAddress
  :: (HasCallStack, MonadCleveland caps m)
  => SpecificOrDefaultAlias -> m ImplicitAddress
newFreshAddress alias = do
  withCap getMiscCap \cap -> cmiGenFreshKey cap alias

-- | Get the signature of the preapplied operation.
signBytes :: (HasCallStack, MonadCleveland caps m) => ByteString -> ImplicitAddress -> m Signature
signBytes bytes signer = do
  withCap getMiscCap \cap -> cmiSignBytes cap bytes signer

-- | Create a list of similarly named 'SpecificAlias'es.
--
-- For example,
--
-- >>> enumAliases @2 "operator" `isEquivalentTo` "operator-0" :< "operator-1" :< Nil
-- True
enumAliases
  :: forall n n'.
     (SingIPeano n, IsoNatPeano n n')
  => ImplicitAlias -> SizedList n SpecificOrDefaultAlias
enumAliases (ImplicitAlias pfx) =
  SpecificAlias <$> SL.generate @n (\n -> ImplicitAlias $ pfx <> "-" <> show n)

-- | Type-safer version of 'signBytes'.
signBinary
  :: (HasCallStack, BytesLike bs, MonadCleveland caps m)
  => bs -> ImplicitAddress -> m (TSignature bs)
signBinary bs addr = TSignature <$> signBytes (toBytes bs) addr

-- | Import an untyped contract from file.
importUntypedContract :: (HasCallStack, MonadCleveland caps m) => FilePath -> m U.Contract
importUntypedContract = runIO . Runtime.importUntypedContract

-- | Import a contract from file.
--
-- The compiler must be able to infer the types of parameter, storage and views.
-- In case there are no views or you don't care, you can use 'noViews'.
importContract
  :: ( HasCallStack, NiceParameter param, NiceStorage st
     , NiceViewsDescriptor vd, DemoteViewsDescriptor vd
     , MonadCleveland caps m
     )
  => FilePath -> m (Contract param st vd)
importContract = runIO . LorentzImport.importContract

{- | Run operations in a batch.
Best used with the @ApplicativeDo@ GHC extension.

Example:

@
{-# LANGUAGE ApplicativeDo #-}

contract <- inBatch $ do
  contract <- originate ...
  for_ [1..3] \i ->
    transfer ...
  return contract
@

Batched operations should be applied to chain faster, but note that batches have
their own limits. For instance, at the moment of writing, the gas limit on a
batch is 10x of gas limit applied to a single operation.

A context of a batch is only 'Applicative', not 'Monad'. This means that:

* Return values of one function cannot be passed to another function in the same
  batch, it can only be returned;
* Sometimes the compiler does not recognize that only 'Applicative' context is
  required, in case of any issues with that - follow the error messages.

-}
inBatch :: (HasCallStack, MonadCleveland caps m) => ClevelandOpsBatch a -> m a
inBatch batch = do
  withCap getOpsCap \cap -> runBatched cap batch

-- | Print the given string verbatim as a comment.
-- At the moment, this is a no-op in emulator tests.
comment :: (HasCallStack, MonadCleveland caps m) => Text -> m ()
comment cmt = do
  withCap getMiscCap \cap -> cmiComment cap cmt

-- | Get the balance of the given address.
getBalance
  :: (HasCallStack, MonadCleveland caps m, ToL1Address addr)
  => addr -> m Mutez
getBalance addr = do
  withCap getMiscCap \cap -> cmiGetBalance cap (toL1Address addr)

-- | Get the delegate for the given contract/implicit address.
getDelegate
  :: (HasCallStack, MonadCleveland caps m, ToL1Address addr)
  => addr -> m (Maybe KeyHash)
getDelegate addr = do
  withCap getMiscCap \cap -> cmiGetDelegate cap (toL1Address addr)

-- | Register the given implicit address as a delegate.
registerDelegate
  :: (HasCallStack, MonadCleveland caps m)
  => ImplicitAddress -> m ()
registerDelegate addr = do
  caps <- ask
  r <- lift $ cmiAttempt (getMiscCap caps) $
    runReaderT (setDelegate addr $ Just $ unImplicitAddress addr) caps
  -- NB: we do some exception wrangling such that registerDelegate doesn't error
  -- out if an address is already a delegate. The primary reason for this is the
  -- disconnect between network, which remembers delegation state between
  -- scenarios, and the emulator, which does not. Hence we want registerDelegate
  -- to be idempotent.
  case r of
    Right () -> pass
    Left e
      | Just (UnexpectedRunErrors [DelegateAlreadyActive]) <- fromPossiblyAnnotatedException e -> pass
      | Just (EEFailedToApplyUpdates GStateAlreadySetDelegate{} :: ExecutorError' AddressAndAlias)
        <- fromPossiblyAnnotatedException e
      -> pass
      | otherwise -> lift $ cmiThrow (getMiscCap caps) e

-- | Set/unset delegate
setDelegate :: (HasCallStack, MonadCleveland caps m) => ImplicitAddress -> Maybe KeyHash -> m ()
setDelegate addr kh = void $ withSender addr $
  withOpsCap \cap -> coiRunOperationBatch cap [OpDelegation kh]

-- | Retrieve a contract's storage in its "RPC representation"
-- (i.e., all its big_maps will be replaced by big_map IDs).
--
-- If the storage is of a user-defined type, then 'Test.Cleveland.deriveRPC' /
-- 'Test.Cleveland.deriveManyRPC' should be used to create an RPC representation of the storage
-- type.
--
-- > data MyStorage = MyStorage { field1 :: Natural, field2 :: BigMap Integer MText }
-- > deriveRPC "MyStorage"
getStorage
  :: forall st addr caps m.
    (HasCallStack, MonadCleveland caps m, ToStorageType st addr, IsoValue (AsRPC st))
  => addr
  -> m (AsRPC st)
getStorage contract = do
  someSt <- getSomeStorage contract
  case someSt ^? castTo @(AsRPC st) . value of
    Just st -> pure st
    Nothing -> failure $ unlinesF
      [ "Expected storage to be of type:"
      , indentF 2 $ build $ demote @(T.ToT (AsRPC st))
      , "But its type was:"
      , indentF 2 $ build $ getT someSt
      ]

-- | Retrieve a contract's full storage, including the contents of its big_maps.
--
-- This function can only be used in emulator-only tests.
getFullStorage
  :: forall st addr caps m.
    (HasCallStack, MonadEmulated caps m, ToStorageType st addr)
  => addr
  -> m st
getFullStorage contract = do
  withCap getEmulatedCap \cap -> eiGetStorage cap contract

-- | Similar to 'getStorage', but doesn't require knowing
-- the storage type in advance.
--
-- Use the optics in 'Morley.Michelson.Typed.AnnotatedValue' to
-- read data from the storage.
getSomeStorage
  :: forall addr caps m.
    (HasCallStack, MonadCleveland caps m, ToContractAddress addr)
  => addr
  -> m SomeAnnotatedValue
getSomeStorage contract = do
  withCap getMiscCap \cap -> cmiGetSomeStorage cap (toContractAddress contract)

-- | Retrieve a big_map value, given a big_map ID and a key.
-- Returns 'Nothing' when the big_map ID does not exist, or it exists but
-- does not contain the given key.
getBigMapValueMaybe
  :: forall k v caps m.
   ( HasCallStack, MonadCleveland caps m
   , NiceComparable k, NicePackedValue k, NiceUnpackedValue v
   )
  => BigMapId k v -> k -> m (Maybe v)
getBigMapValueMaybe bmId key = do
  withCap getMiscCap \cap -> cmiGetBigMapValueMaybe cap bmId key

-- | Like 'getBigMapValueMaybe', but fails the tests instead of returning 'Nothing'.
getBigMapValue
  :: forall k v caps m.
   ( HasCallStack, MonadCleveland caps m
   , NiceComparable k, NicePackedValue k, NiceUnpackedValue v
   , Buildable k
   )
  => BigMapId k v -> k -> m v
getBigMapValue bmId k =
  getBigMapValueMaybe bmId k >>= \case
    Just v -> pure v
    Nothing -> failure $ unlinesF @_ @Builder
      [ "Either:"
      , "  1. A big_map with ID '" +| bmId |+ "' does not exist, or"
      , "  2. It exists, but does not contain the key '" +| k |+ "'."
      ]

-- | Retrieve all big_map values, given a big_map ID.
-- Returns 'Nothing' when the big_map ID does not exist.
getAllBigMapValuesMaybe
  :: forall k v caps m.
  ( HasCallStack, MonadCleveland caps m
  , NiceComparable k, NiceUnpackedValue v
  )
  => BigMapId k v -> m (Maybe [v])
getAllBigMapValuesMaybe bmId = do
  withCap getMiscCap \cap -> cmiGetAllBigMapValuesMaybe cap bmId

-- | Like 'getAllBigMapValuesMaybe', but fails the tests instead of returning 'Nothing'.
getAllBigMapValues
  :: forall k v caps m.
  ( HasCallStack, MonadCleveland caps m
  , NiceComparable k, NiceUnpackedValue v
  )
  => BigMapId k v -> m [v]
getAllBigMapValues bmId =
  getAllBigMapValuesMaybe bmId >>= \case
    Just vs -> pure vs
    Nothing -> failure $ "A big map with ID '" +| bmId |+ "' does not exist"

-- | Retrieve a big_map size, given a big_map ID.
-- Returns 'Nothing' when the big_map ID does not exist.
--
-- /O(n)/, because it's implemented with 'Morley.Client.RPC.Getters.getBigMapValues'.
getBigMapSizeMaybe
  :: forall k v caps m.
  ( HasCallStack, MonadCleveland caps m
  , NiceComparable k, NiceUnpackedValue v
  )
  => BigMapId k v -> m (Maybe Natural)
getBigMapSizeMaybe bmId =
    fmap (fmap (Unsafe.fromIntegral @Int @Natural . length)) (getAllBigMapValuesMaybe bmId)

-- | Like 'getBigMapSizeMaybe', but fails the tests instead of returning 'Nothing'.
getBigMapSize
  :: forall k v caps m.
  ( HasCallStack, MonadCleveland caps m
  , NiceComparable k, NiceUnpackedValue v
  )
  => BigMapId k v -> m Natural
getBigMapSize bmId =
  Unsafe.fromIntegral @Int @Natural . length <$> getAllBigMapValues bmId

-- | Get the public key associated with given address.
-- Fail if given address is not an implicit account.
getPublicKey :: (HasCallStack, MonadCleveland caps m) => ImplicitAddress -> m PublicKey
getPublicKey addr = do
  withCap getMiscCap \cap -> cmiGetPublicKey cap addr

-- | Get the chain's @ChainId@.
getChainId :: (HasCallStack, MonadCleveland caps m) => m ChainId
getChainId = do
  withCap getMiscCap \cap -> cmiGetChainId cap

-- | Advance at least the given amount of time, or until a new block is baked,
-- whichever happens last.
--
-- On a real network, this is implemented using @threadDelay@, so it's advisable
-- to use small amounts of time only.
advanceTime
  :: forall unit caps m
  . (HasCallStack, MonadCleveland caps m, KnownDivRat unit Second)
  => Time unit -> m ()
advanceTime time = do
  withCap getMiscCap \cap -> cmiAdvanceTime cap time

-- | Wait till the provided number of levels is past.
advanceLevel
  :: forall caps m
  . (HasCallStack, MonadCleveland caps m)
  => Natural -> m ()
advanceLevel l = do
  withCap getMiscCap \cap -> cmiAdvanceToLevel cap (+ l)

-- | Wait till the provided level is reached.
advanceToLevel
  :: forall caps m
  . (HasCallStack, MonadCleveland caps m)
  => Natural -> m ()
advanceToLevel target = do
  withCap getMiscCap \cap -> cmiAdvanceToLevel cap (const target)

-- | Get the timestamp observed by the last block to be baked.
getNow :: (HasCallStack, MonadCleveland caps m) => m Timestamp
getNow = do
  withCap getMiscCap \cap -> cmiGetNow cap

-- | Get the current level observed by the last block to be baked.
getLevel :: (HasCallStack, MonadCleveland caps m) => m Natural
getLevel = do
  withCap getMiscCap \cap -> cmiGetLevel cap

-- | Get approximate block interval in seconds. Note, that this value
-- is minimal bound and real intervals can be larger, see
-- http://tezos.gitlab.io/active/consensus.html#minimal-block-delay-function
-- for more information about block delays.
getApproximateBlockInterval :: (HasCallStack, MonadCleveland caps m) => m (Time Second)
getApproximateBlockInterval = do
  withCap getMiscCap \cap -> cmiGetApproximateBlockInterval cap

-- | Get minimal block delay in seconds. This is essentially the same as
-- 'getApproximateBlockInterval', but returns a 'Natural' instead of @Time
-- Second@.
--
-- Can be useful when testing code using @MIN_BLOCK_TIME@ instruction.
getMinBlockTime :: (HasCallStack, MonadCleveland caps m) => m Natural
getMinBlockTime = toNum @Second <$> getApproximateBlockInterval

-- | Execute a contract's code without originating it.
-- The chain's state will not be modified.
--
-- Notes:
--
-- * If the contract's code emits operations, they will not be executed.
-- * The sender's account won't be debited.
-- * When running an _originated_ contract, the @BALANCE@ instruction returns the
--   sum of the contract's balance before the transfer operation + the amount of tz being transferred.
--   In other words, this invariant holds: @BALANCE >= AMOUNT@.
--   However, since `runCode` allows overriding the @BALANCE@ instruction,
--   then this invariant no longer holds. It's possible that @BALANCE < AMOUNT@.
runCode
  :: (HasCallStack, MonadCleveland caps m, HasRPCRepr st, IsoValue (AsRPC st))
  => RunCode cp st vd -> m (AsRPC st)
runCode rc = do
  sender <- view senderL
  withCap getMiscCap \cap -> cmiRunCode cap sender rc

-- | Execute multiple testing scenarios independently.
--
-- * Actions performed before 'branchout' will be observed by all branches.
-- * Actions performed in branches will _not_ be observed by any actions performed after 'branchout'.
-- * Actions performed in one branch will _not_ be observed by another branch.
-- * The test succeeds IFF all branches succeed.
-- * If any branch fails, the test ends immediately and the remaining branches
--    won't be executed.
--
-- The following property holds:
--
-- > pre >> branchout [a, b, c] = branchout [pre >> a, pre >> b, pre >> c]
--
-- The list of branches must be non-empty.
branchout :: forall caps m. (MonadEmulated caps m) => [(Text, m ())] -> m ()
branchout branches = do
  caps <- ask
  let branches' :: [(Text, ClevelandBaseMonad caps ())] = second (flip runReaderT caps) <$> branches
  lift $ eiBranchout (getEmulatedCap caps) branches'

-- | Execute one or more actions and roll them back afterwards.
-- Actions performed in 'offshoot' will _not_ be observed by any
-- actions performed after 'offshoot'.
--
-- Similar to 'branchout', but accepts one single branch.
offshoot :: forall caps m. MonadEmulated caps m => Text -> m () -> m ()
offshoot branchName branch = branchout [(branchName, branch)]

{- | Returns the result of the action with the logs it produced. Logs are messages
printed by the Lorentz instruction 'Lorentz.printComment'.

This function can be combined either with lens-based accessors or helper functions to get
more specific information about logs.

Examples:

@
(logsInfo, _) <- getMorleyLogs scenario
logsInfo ^.. each . logsL @== [MorleyLogs ["log"], MorleyLogs ["log2"]]
logsInfo ^.. each . filterLogsByAddrL addr @== [MorleyLogs ["log"]]
@

@
(logsInfo, _) <- getMorleyLogs scenario
collectLogs logsInfo @== MorleyLogs ["log", "log2"]
logsForAddress logsInfo @== [MorleyLogs ["log"]]
@

-}
getMorleyLogs :: forall a caps m. MonadEmulated caps m => m a -> m (LogsInfo, a)
getMorleyLogs action = do
  caps <- ask
  let action' :: ClevelandBaseMonad caps a = runReaderT action caps
  lift $ eiGetMorleyLogs (getEmulatedCap caps) action'

-- | Version of `getMorleyLogs` for actions that don't return a result.
getMorleyLogs_ :: MonadEmulated caps m => m () -> m LogsInfo
getMorleyLogs_ action = fst <$> getMorleyLogs action

-- | Updates voting power accessible via @VOTING_POWER@ and similar
-- instructions.
setVotingPowers :: MonadEmulated caps m => VotingPowers -> m ()
setVotingPowers vp = do
  withCap getEmulatedCap \cap -> eiSetVotingPowers cap vp

-- | A helper constraint synonym to make signatures below a bit shorter
type EqBaseMonad a b = ClevelandBaseMonad a ~ ClevelandBaseMonad b

-- | Perform an action if we are currently in emulation mode.
-- See also 'ifEmulation' note on constraints.
whenEmulation
  :: MonadCleveland caps m
  => (forall caps1 m1. (EqBaseMonad caps caps1, MonadEmulated caps1 m1) => m1 ())
  -> m ()
whenEmulation action = ifEmulation action pass

-- | Perform an action if we are currently in network mode.
-- See also 'ifEmulation' note on constraints.
whenNetwork
  :: MonadCleveland caps m
  => (forall caps1 m1. (EqBaseMonad caps caps1, MonadNetwork caps1 m1) => m1 ())
  -> m ()
whenNetwork action = ifEmulation pass action

{- | Perform one action if we are currently in emulation mode, another otherwise

Functions passed as the first two arguments are universally quantified over
the outer monad, so if additional constraints are required beyond
'MonadEmulated' or 'MonadCleveland', those constraints have to go on the base
monad, e.g.

@
someFunction :: (MonadCleveland caps m, MonadFail (ClevelandBaseMonad caps)) => m ()
someFunction = whenEmulation do
  Just x <- pure (Just 1 :: Maybe Int) -- this would error without MonadFail
  runIO $ print x
@
-}
ifEmulation
  :: forall a caps m
   . MonadCleveland caps m
  => (forall caps1 m1. (EqBaseMonad caps caps1, MonadEmulated caps1 m1) => m1 a)
  -> (forall caps1 m1. (EqBaseMonad caps caps1, MonadNetwork caps1 m1) => m1 a)
  -> m a
ifEmulation onEmu onNet = withCap getMiscCap cmiUnderlyingImpl >>= \case
  Right impl -> withCap (view clevelandCapsL) (runReaderT onNet . NetworkCaps impl)
  Left impl -> withCap (view clevelandCapsL) (runReaderT onEmu . EmulatedCaps impl)

-- | Get a 'MorleyClientEnv' when running a test on network. Useful to run
-- f.ex. @octez-client@ inside a network test.
--
-- This is considered a pretty low-level function, so it's better to avoid it in
-- most cases.
getMorleyClientEnv :: MonadNetwork caps m => m MorleyClientEnv
getMorleyClientEnv = asks (neMorleyClientEnv . getNetworkEnvCap)