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

-- | Implementation of generic operations submission.
module Morley.Client.Action.Operation
  ( Result
  , runOperations
  , runOperationsNonEmpty
  -- * helpers
  , dryRunOperationsNonEmpty
  ) where

import Control.Lens (has, (%=), (&~))
import Data.List (zipWith4)
import Data.List.NonEmpty qualified as NE
import Data.Map qualified as Map
import Data.Ratio ((%))
import Data.Singletons (Sing, SingI, demote, sing)
import Data.Text qualified as T
import Fmt (blockListF, blockListF', listF, nameF, pretty, (+|), (|+))

import Morley.Client.Action.Common
import Morley.Client.Logging
import Morley.Client.RPC.Class
import Morley.Client.RPC.Error
import Morley.Client.RPC.Getters
import Morley.Client.RPC.Types
import Morley.Client.TezosClient
import Morley.Client.Types
import Morley.Client.Util (epNameToTezosEp)
import Morley.Micheline (StringEncode(..), TezosInt64, TezosMutez(..), toExpression)
import Morley.Michelson.Typed (Value)
import Morley.Tezos.Address
import Morley.Tezos.Address.Alias
import Morley.Tezos.Core
import Morley.Tezos.Crypto
import Morley.Util.ByteString
import Morley.Util.Constrained
import Morley.Util.Named

-- | Designates output of an operation.
data Result
instance OperationInfoDescriptor Result where
  type TransferInfo Result = [IntOpEvent]
  type TransferTicketInfo Result = [IntOpEvent]
  type OriginationInfo Result = ContractAddress
  type RevealInfo Result = ()
  type DelegationInfo Result = ()

logOperations
  :: forall env m kind.
     ( WithClientLog env m
     , HasTezosClient m
     )
  => AddressWithAlias kind
  -> NonEmpty (OperationInfo ClientInput)
  -> m ()
logOperations sender ops = do
  let opName =
        if | all (has _OpTransfer) ops -> "transactions"
           | all (has _OpOriginate) ops -> "originations"
           | all (has _OpReveal) ops -> "reveals"
           | all (has _OpTransferTicket) ops -> "ticket transfers"
           | otherwise -> "operations"

      buildOp = \case
        (OpTransfer tx, mbAlias) ->
          buildTxDataWithAlias mbAlias tx
        (OpTransferTicket tx, mbAlias) ->
          buildTxTicketDataWithAlias mbAlias tx
        (OpOriginate orig, _) ->
          odName orig |+ " (temporary alias)"
        (OpReveal rv, mbAlias) ->
          "Key " +| rdPublicKey rv |+ maybe "" (\a -> " (" +| a |+ ")") mbAlias
        (OpDelegation delegate, mbAlias) ->
          "Key Hash " +| ddDelegate delegate |+ maybe "" (\a -> " (" +| a |+ ")") mbAlias

  let needsAliasStore = (`any` ops) \case
        OpTransfer{} -> True
        OpTransferTicket{} -> True
        OpOriginate{} -> False
        OpReveal{} -> True
        OpDelegation{} -> True

  anas <- if needsAliasStore
    then Map.fromList . fmap swap <$> getAliasesAndAddresses
    else pure mempty

  let mbAlias :: KindedAddress k -> Maybe Text
      mbAlias addr = Map.lookup (pretty addr) anas

  let aliases = ops <&> \case
        OpTransfer (TransactionData tx) -> withConstrained (tdReceiver tx) mbAlias
        OpTransferTicket TransferTicketData{..} -> withConstrained ttdDestination mbAlias
        OpOriginate _ -> Nothing
        OpReveal r -> mbAlias . mkKeyAddress $ rdPublicKey r
        OpDelegation d -> mbAlias . ImplicitAddress =<< ddDelegate d

  logInfo $ T.strip $ -- strip trailing newline
    "Running " +| opName +| " by " +| sender |+ ":\n" +|
    blockListF' "-" buildOp (ops `NE.zip` aliases)

-- | Perform sequence of operations.
--
-- Returns operation hash (or @Nothing@ in case empty list was provided) and result of
-- each operation (nothing for transactions and an address for originated contracts
runOperations
  :: forall m env.
     ( HasTezosRpc m
     , HasTezosClient m
     , WithClientLog env m
     )
  => ImplicitAddressWithAlias
  -> [OperationInfo ClientInput]
  -> m (Maybe OperationHash, [OperationInfo Result])
runOperations sender operations = case operations of
  [] -> return (Nothing, [])
  op : ops -> do
    (opHash, res) <- runOperationsNonEmpty sender $ op :| ops
    return $ (Just opHash, toList res)

-- | How many times to retry if an operation fails after injection
injectionRetryCount :: Natural
injectionRetryCount = 2

-- | Perform non-empty sequence of operations.
--
-- Returns operation hash and result of each operation
-- (nothing for transactions and an address for originated contracts).
runOperationsNonEmpty
  :: forall m env.
     ( HasTezosRpc m
     , HasTezosClient m
     , WithClientLog env m
     )
  => ImplicitAddressWithAlias
  -> NonEmpty (OperationInfo ClientInput)
  -> m (OperationHash, NonEmpty (OperationInfo Result))
runOperationsNonEmpty sender operations =
  runOperationsNonEmptyHelper @'RealRun injectionRetryCount sender operations

-- | Flag that is used to determine @runOperationsNonEmptyHelper@ behaviour.
data RunMode = DryRun | RealRun

isRealRun :: forall (runMode :: RunMode). (SingI runMode) => Bool
isRealRun = case sing @runMode of
  SRealRun -> True
  SDryRun  -> False

-- | Type family which is used to determine the output type of the
-- @runOperationsNonEmptyHelper@.
type family RunResult (a :: RunMode) where
  RunResult 'DryRun = NonEmpty (AppliedResult, TezosMutez)
  RunResult 'RealRun = (OperationHash, NonEmpty (OperationInfo Result))

data SingRunResult :: RunMode -> Type where
  SDryRun :: SingRunResult 'DryRun
  SRealRun :: SingRunResult 'RealRun

type instance Sing = SingRunResult

instance SingI 'DryRun where
  sing = SDryRun

instance SingI 'RealRun where
  sing = SRealRun

-- | Perform dry-run for sequence of operations.
--
-- Returned @AppliedResult@ contains information about estimated limits,
-- storage changes, etc. Additionally, estimated fees are returned.
dryRunOperationsNonEmpty
  :: forall m env.
     ( HasTezosRpc m
     , HasTezosClient m
     , WithClientLog env m
     )
  => ImplicitAddressWithAlias
  -> NonEmpty (OperationInfo ClientInput)
  -> m (NonEmpty (AppliedResult, TezosMutez))
dryRunOperationsNonEmpty sender operations =
  runOperationsNonEmptyHelper @'DryRun 0 sender operations

-- | Perform non-empty sequence of operations and either dry-run
-- and return estimated limits and fees or perform operation injection.
-- Behaviour is defined via @RunMode@ flag argument.
runOperationsNonEmptyHelper
  :: forall (runMode :: RunMode) m env.
     ( HasTezosRpc m
     , HasTezosClient m
     , WithClientLog env m
     , SingI runMode
     )
  => Natural
  -> ImplicitAddressWithAlias
  -> NonEmpty (OperationInfo ClientInput)
  -> m (RunResult runMode)
runOperationsNonEmptyHelper retryCount sender@(AddressWithAlias senderAddress _) operations' = do
  operations <- fromMaybe operations' <$> runMaybeT do
    guard $ isRealRun @runMode && mayNeedSenderRevealing (toList operations')
    Nothing <- lift $ getManagerKey senderAddress
    pk <- lift $ getPublicKey sender
    pure $ OpReveal (RevealData pk Nothing) :| toList operations'
  logOperations sender operations
  -- If we intend to fail on duplicate aliases, we want to fail now, rather than
  -- after contract origination.
  forM_ operations $ \case
    OpOriginate OriginationData{odAliasBehavior = ForbidDuplicateAlias, ..} -> do
      resolved <- resolveAddressMaybe (AddressAlias odName)
      whenJust resolved $ const $ throwM $ DuplicateAlias (unAlias odName)
    _ -> pass
  mbPassword <- getKeyPassword sender
  pp <- getProtocolParameters
  OperationConstants{..} <- preProcessOperation senderAddress

  let convertOps i = WithCommonOperationData commonData . \case
        OpTransfer (TransactionData TD {tdReceiver=Constrained tdReceiver, ..}) ->
          OpTransfer TransactionOperation
          { toDestination = MkAddress tdReceiver
          , toAmount = TezosMutez tdAmount
          , toParameters = toParametersInternals tdEpName tdParam
          }
        OpTransferTicket (TransferTicketData {..})
          | _ :: Value t <- ttdTicketContents ->
          OpTransferTicket TransferTicketOperation
          { ttoTicketContents = toExpression ttdTicketContents
          , ttoTicketTy       = toExpression (demote @t)
          , ttoTicketTicketer = ttdTicketTicketer
          , ttoTicketAmount   = StringEncode ttdTicketAmount
          , ttoDestination    = ttdDestination
          , ttoEntrypoint     = epNameToTezosEp ttdEntrypoint
          }
        OpOriginate OriginationData{..} ->
          OpOriginate OriginationOperation
          { ooBalance = TezosMutez odBalance
          , ooDelegate = odDelegate
          , ooScript = mkOriginationScript odContract odStorage
          }
        OpReveal RevealData{..} ->
          OpReveal RevealOperation
          { roPublicKey = rdPublicKey
          }
        OpDelegation DelegationData{..} ->
          OpDelegation DelegationOperation
          { doDelegate = ddDelegate
          }
        where
          commonData = mkCommonOperationData pp
            ! #sender senderAddress
            ! #counter (ocCounter + i)
            ! #num_operations (fromIntegral (length operations))

  let opsToRun = NE.zipWith convertOps (1 :| [(2 :: TezosInt64)..]) operations
      mbFees = operations <&> \case
        OpTransfer (TransactionData TD {..}) -> tdMbFee
        OpTransferTicket TransferTicketData{..} -> ttdMbFee
        OpOriginate OriginationData{..} -> odMbFee
        OpReveal RevealData{..} -> rdMbFee
        OpDelegation DelegationData {..} -> ddMbFee

  -- Perform run_operation with dumb signature in order
  -- to estimate gas cost, storage size and paid storage diff
  let runOp = RunOperation
        { roOperation = RunOperationInternal
          { roiBranch = ocLastBlockHash
          , roiContents = opsToRun
          , roiSignature = stubSignature
          }
        , roChainId = bcChainId ocBlockConstants
        }
  (results, _) <- getAppliedResults (Left runOp)

  let -- Learn how to forge given operations
      forgeOp :: NonEmpty OperationInput -> m ByteString
      forgeOp ops =
        fmap unHexJSONByteString . forgeOperation $ ForgeOperation
          { foBranch = ocLastBlockHash
          , foContents = ops
          }

  let -- Attach a signature to forged operation + return the signature itself
      signForgedOp :: ByteString -> m (Signature, ByteString)
      signForgedOp op = do
        signature' <- signBytes sender mbPassword (addOperationPrefix op)
        return (signature', prepareOpForInjection op signature')

  -- Fill in fees
  let
    updateOp :: OperationInput -> Maybe Mutez -> AppliedResult -> Bool -> m (OperationInput, Maybe (Signature, ByteString))
    updateOp opToRun@(WithCommonOperationData _ internalOp) mbFee ar isFirst = do
      let gasSafetyGuard = 100
          -- @gasSafetyGuard@ is added to origination operations and transfers to non-implicit
          -- accounts, see https://gitlab.com/tezos/tezos/-/blob/v13.0/src/proto_013_PtJakart/lib_client/injection.ml#L750
          additionalGas = case internalOp of
            OpTransferTicket{} -> gasSafetyGuard
            OpOriginate _ -> gasSafetyGuard
            OpTransfer (TransactionOperation {..}) -> case toDestination of
              MkAddress ImplicitAddress{} -> 0
              MkAddress ContractAddress{} -> gasSafetyGuard
              MkAddress SmartRollupAddress{} -> gasSafetyGuard
            OpReveal _ -> 0
            OpDelegation _ -> 0
      let storageLimit = computeStorageLimit [ar] pp + 20
          -- similarly to @octez-client@, we add 20 for safety
      let gasLimit = ceiling (arConsumedMilliGas ar % 1000) + additionalGas
          updateCommonDataForFee fee =
            updateCommonData gasLimit storageLimit (TezosMutez fee)

      (_fee, op, mReadySignedOp) <- convergingFee
        @OperationInput
        @(Maybe (Signature, ByteString))  -- ready operation and its signature
        (\fee ->
          return $ opToRun &~
            wcoCommonDataL %= updateCommonDataForFee fee
          )
        (\op -> do
          forgedOp <- forgeOp $ one op
          -- In the Tezos implementation the first transaction
          -- in the series pays for signature.
          -- Signature of hash should be constant in size,
          -- so we can pass any signature, not necessarily the final one
          (fullForgedOpLength, mExtra) <-
            if isFirst
              then do
                res@(_signature, signedOp) <- signForgedOp forgedOp
                return (length signedOp, Just res)
              else
                -- Forge output automatically includes additional 32-bytes header
                -- which should be ommited for all operations in batch except the first one.
                pure (length forgedOp - 32, Nothing)
          return
            ( maybe (computeFee ocFeeConstants fullForgedOpLength gasLimit) id mbFee
            , mExtra
            )
          )

      return (op, mReadySignedOp)

  let
    zipWith4NE
      :: (a -> b -> c -> d -> e) -> NonEmpty a -> NonEmpty b -> NonEmpty c -> NonEmpty d
      -> NonEmpty e
    zipWith4NE f (a :| as) (b :| bs) (c :| cs) (d :| ds) =
      (f a b c d) :| zipWith4 f as bs cs ds
  -- These two lists must have the same length here.
  -- @opsToRun@ is constructed directly from @params@.
  -- The length of @results@ is checked in @getAppliedResults@.
  (updOps, readySignedOps) <- fmap NE.unzip . sequenceA $
    zipWith4NE updateOp opsToRun mbFees results (True :| repeat False)

  -- Forge operation with given limits and get its hexadecimal representation
  (signature', signedOp) <- case readySignedOps of
    -- Save one forge + sign call pair in case of one operation
    Just readyOp :| [] -> pure readyOp
    -- In case of batch we have to reforge the full operation
    _ -> forgeOp updOps >>= signForgedOp

  -- Operation still can fail due to insufficient gas or storage limit, so it's required
  -- to preapply it before injecting
  let preApplyOp = PreApplyOperation
        { paoProtocol = bcProtocol ocBlockConstants
        , paoBranch = ocLastBlockHash
        , paoContents = updOps
        , paoSignature = signature'
        }
  (ars2, iopsData) <- getAppliedResults (Right preApplyOp)
  case sing @runMode of
    SDryRun -> do
      let fees = map (codFee . wcoCommon) updOps
      return $ NE.zip ars2 fees
    SRealRun -> do
      operationHash <- waitForOperation $ injectOperation (HexJSONByteString signedOp)
      let contractAddrs = arOriginatedContracts <$> ars2
      opsRes <- forM (NE.zip operations contractAddrs) $ \case
        (OpTransfer _, addrs) -> do
          unless (null addrs) $
            logInfo . T.strip $
              "The following contracts were originated during transactions: " +|
              listF addrs |+ ""
          return $ OpTransfer $ mapMaybe iopsDataToEmitOp iopsData
        (OpTransferTicket _, addrs) -> do
          unless (null addrs) $
            logInfo . T.strip $
              "The following contracts were originated during transactions: " +|
              listF addrs |+ ""
          return $ OpTransferTicket $ mapMaybe iopsDataToEmitOp iopsData
        (OpOriginate _, []) ->
          throwM RpcOriginatedNoContracts
        (OpOriginate OriginationData{..}, [addr]) -> do
          logDebug $ "Saving " +| addr |+ " for " +| odName |+ "\n"
          rememberContract odAliasBehavior addr odName
          logInfo $ "Originated contract: " <> pretty odName
          return $ OpOriginate addr
        (OpOriginate _, addrs@(_ : _ : _)) ->
          throwM $ RpcOriginatedMoreContracts addrs
        (OpReveal _, _) ->
          return $ OpReveal ()
        (OpDelegation _, _) ->
          return $ OpDelegation ()
      forM_ ars2 logStatistics
      return (operationHash, opsRes)
      `catch` \case
        (UnexpectedRunErrors errs) | retryCount > 0 -> do
          logError $ pretty $ nameF "When injecting operations, there were unexpected errors" $
            blockListF errs
          runOperationsNonEmptyHelper @runMode (retryCount - 1) sender operations
        e -> throwM e
  where
    iopsDataToEmitOp :: InternalOperationData -> Maybe IntOpEvent
    iopsDataToEmitOp = \case
      IODEvent evt -> Just evt
      _ -> Nothing

    mayNeedSenderRevealing :: [OperationInfo i] -> Bool
    mayNeedSenderRevealing = any \case
      OpTransfer{} -> True
      OpTransferTicket{} -> True
      OpOriginate{} -> True
      OpReveal{} -> False
      OpDelegation{} -> True

    logStatistics :: AppliedResult -> m ()
    logStatistics ar = do
      let showTezosInt64 = show . unStringEncode
      logInfo $ "Consumed milli-gas: " <> showTezosInt64 (arConsumedMilliGas ar)
      logInfo $ "Storage size: " <> showTezosInt64 (arStorageSize ar)
      logInfo $ "Paid storage size diff: " <> showTezosInt64 (arPaidStorageDiff ar)