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

-- | Functions useful for implementing instances of type classes from this package.
-- Monads and actual instances are defined in separate modules.

module Morley.Client.App
  ( -- * RunClient
    runRequestAcceptStatusImpl
  , throwClientErrorImpl

  -- * HasTezosRpc
  , getBlockHashImpl
  , getCounterImpl
  , getBlockHeaderImpl
  , getBlockConstantsImpl
  , getScriptSizeAtBlockImpl
  , getBlockOperationsImpl
  , getBlockOperationHashesImpl
  , getProtocolParametersImpl
  , runOperationImpl
  , preApplyOperationsImpl
  , forgeOperationImpl
  , injectOperationImpl
  , getContractScriptImpl
  , getContractStorageAtBlockImpl
  , getContractBigMapImpl
  , getBigMapValueAtBlockImpl
  , getBigMapValuesAtBlockImpl
  , getBalanceImpl
  , getManagerKeyImpl
  , runCodeImpl
  , getChainIdImpl
  , getDelegateImpl
  , waitForOperationImpl
  , getTicketBalanceAtBlockImpl
  , getAllTicketBalancesAtBlockImpl
  , packDataImpl

  -- * Timeouts and retries
  , retryOnTimeout
  , failOnTimeout
  , retryOnceOnTimeout
  , waitBeforeRetry
  , handleInvalidCounterRpc
  ) where

import Unsafe qualified (fromIntegral, (!!))

import Control.Concurrent (threadDelay)
import Data.Aeson qualified as Aeson
import Data.Binary.Builder qualified as Binary
import Data.Text (isInfixOf)
import Fmt (Buildable(..), Doc, build, pretty, (+|), (|+))
import Network.HTTP.Types (Status(..), renderQuery)
import Servant.Client (ClientEnv, runClientM)
import Servant.Client.Core
  (ClientError(..), Request, RequestBody(..), RequestF(..), Response, ResponseF(..), RunClient)
import Servant.Client.Core.RunClient (runRequest)
import Servant.Client.Streaming (withClientM)
import Servant.Types.SourceT (SourceT(unSourceT), StepT(..))
import System.Random (randomRIO)
import UnliftIO (MonadUnliftIO, withRunInIO)
import UnliftIO.Async (wait, waitEither, withAsync)
import UnliftIO.Timeout (timeout)

import Morley.Client.Logging (WithClientLog, logDebug)
import Morley.Client.RPC
import Morley.Client.RPC.API qualified as API
import Morley.Micheline
  (Expression, TezosInt64, TezosNat, toExpression, unStringEncode, unTezosMutez)
import Morley.Michelson.Typed qualified as T
import Morley.Tezos.Address
import Morley.Tezos.Core (ChainId, Mutez, parseChainId)
import Morley.Tezos.Crypto (KeyHash, PublicKey)
import Morley.Util.ByteString (HexJSONByteString)
import Morley.Util.Exception (throwLeft)

----------------
-- RunClient functions
----------------

runRequestAcceptStatusImpl
  :: forall env m. (WithClientLog env m, MonadIO m, MonadThrow m)
  => ClientEnv -> Maybe [Status] -> Request -> m Response
runRequestAcceptStatusImpl env _ req = do
  logRequest req
  response <- either logResponseAndThrowError pure =<<
    liftIO (runClientM (runRequest req) env)
  response <$ logResponse response

  where
    getResponse :: ClientError -> Maybe Response
    getResponse = \case
      FailureResponse _ response -> Just response
      DecodeFailure _ response -> Just response
      UnsupportedContentType _ response -> Just response
      InvalidContentTypeHeader response -> Just response
      ConnectionError{} -> Nothing

    logResponseAndThrowError :: ClientError -> m Response
    logResponseAndThrowError err = do
      whenJust (getResponse err) logResponse
      throwClientErrorImpl err

throwClientErrorImpl :: forall m a . MonadThrow m => ClientError -> m a
throwClientErrorImpl err = case err of
  FailureResponse _ resp
    | 500 <- statusCode (responseStatusCode resp) ->
      handleInternalError (responseBody resp)
  _ -> throwM err
  where
    -- In some cases RPC returns important useful errors as internal ones.
    -- We try to parse the response to a list of 'InternalError'.
    -- If we receive one 'InternalError', we throw it wrapped into
    -- 'ClientInternalError', that's what we observed in most obvious cases.
    -- If we receive more than one, we wrap them into 'UnexpectedInternalErrors'.
    handleInternalError :: LByteString -> m a
    handleInternalError body = case Aeson.decode @[InternalError] body of
      Nothing -> case Aeson.decode @[RunError] body of
        Nothing -> throwM err
        Just runErrs -> throwM $ RunCodeErrors runErrs
      Just [knownErr] -> throwM $ ClientInternalError knownErr
      Just errs -> throwM $ UnexpectedInternalErrors errs

----------------
-- HasTezosRpc functions
----------------

getBlockHashImpl :: (RunClient m, MonadUnliftIO m, MonadCatch m) => BlockId -> m BlockHash
getBlockHashImpl = retryOnceOnTimeout ... fmap BlockHash . API.getBlockHash API.nodeMethods

getCounterImpl :: (RunClient m, MonadUnliftIO m, MonadCatch m) => BlockId -> ImplicitAddress -> m TezosInt64
getCounterImpl = retryOnceOnTimeout ... API.getCounter API.nodeMethods

getBlockHeaderImpl ::
  (RunClient m, MonadUnliftIO m, MonadCatch m) => BlockId -> m BlockHeader
getBlockHeaderImpl = retryOnceOnTimeout ... API.getBlockHeader API.nodeMethods

getBlockConstantsImpl :: (RunClient m, MonadUnliftIO m, MonadCatch m) => BlockId -> m BlockConstants
getBlockConstantsImpl = retryOnceOnTimeout ... API.getBlockConstants API.nodeMethods

getBlockOperationsImpl :: (RunClient m, MonadUnliftIO m, MonadCatch m) => BlockId -> m [[BlockOperation]]
getBlockOperationsImpl =
  retryOnceOnTimeout ... API.getBlockOperations API.nodeMethods

getBlockOperationHashesImpl :: (RunClient m, MonadUnliftIO m, MonadCatch m) => BlockId -> m [[OperationHash]]
getBlockOperationHashesImpl =
  retryOnceOnTimeout ... API.getBlockOperationHashes API.nodeMethods

getProtocolParametersImpl ::
  (RunClient m, MonadUnliftIO m, MonadCatch m) => BlockId -> m ProtocolParameters
getProtocolParametersImpl = retryOnceOnTimeout ... API.getProtocolParameters API.nodeMethods

runOperationImpl ::
  (RunClient m, MonadUnliftIO m, MonadCatch m) => BlockId -> RunOperation -> m RunOperationResult
runOperationImpl = retryOnceOnTimeout ... API.runOperation API.nodeMethods

preApplyOperationsImpl ::
  (RunClient m, MonadUnliftIO m, MonadCatch m) =>
  BlockId -> [PreApplyOperation] -> m [RunOperationResult]
preApplyOperationsImpl =
  retryOnceOnTimeout ... API.preApplyOperations API.nodeMethods

forgeOperationImpl ::
  (RunClient m, MonadUnliftIO m, MonadCatch m) =>
  BlockId -> ForgeOperation -> m HexJSONByteString
forgeOperationImpl = retryOnceOnTimeout ... API.forgeOperation API.nodeMethods

injectOperationImpl ::
  (RunClient m, MonadUnliftIO m, MonadCatch m) =>
  HexJSONByteString -> m OperationHash
injectOperationImpl =
  failOnTimeout ... API.injectOperation API.nodeMethods

getContractScriptImpl ::
  (RunClient m, MonadUnliftIO m, MonadCatch m) =>
  BlockId -> ContractAddress -> m OriginationScript
getContractScriptImpl = retryOnceOnTimeout ... API.getScript API.nodeMethods

getContractStorageAtBlockImpl ::
  (RunClient m, MonadUnliftIO m, MonadCatch m) =>
  BlockId -> ContractAddress -> m Expression
getContractStorageAtBlockImpl =
  retryOnceOnTimeout ... API.getStorageAtBlock API.nodeMethods

getContractBigMapImpl ::
  (RunClient m, MonadUnliftIO m, MonadCatch m) =>
  BlockId -> ContractAddress -> GetBigMap -> m GetBigMapResult
getContractBigMapImpl = retryOnceOnTimeout ... API.getBigMap API.nodeMethods

getScriptSizeAtBlockImpl ::
  (RunClient m, MonadUnliftIO m, MonadCatch m) => BlockId -> CalcSize -> m ScriptSize
getScriptSizeAtBlockImpl = retryOnceOnTimeout ... API.getScriptSizeAtBlock API.nodeMethods

getBigMapValueAtBlockImpl ::
  (RunClient m, MonadUnliftIO m, MonadCatch m) =>
  BlockId -> Natural -> Text -> m Expression
getBigMapValueAtBlockImpl =
  retryOnceOnTimeout ... API.getBigMapValueAtBlock API.nodeMethods

getBigMapValuesAtBlockImpl ::
  (RunClient m, MonadUnliftIO m, MonadCatch m) =>
  BlockId -> Natural -> Maybe Natural -> Maybe Natural -> m Expression
getBigMapValuesAtBlockImpl =
  retryOnceOnTimeout ... API.getBigMapValuesAtBlock API.nodeMethods

getBalanceImpl ::
  (RunClient m, MonadUnliftIO m, MonadCatch m) => BlockId -> Address -> m Mutez
getBalanceImpl =
  retryOnceOnTimeout ... fmap unTezosMutez ... API.getBalance API.nodeMethods

getTicketBalanceAtBlockImpl ::
  (RunClient m, MonadUnliftIO m, MonadCatch m) => BlockId -> Address -> GetTicketBalance -> m Natural
getTicketBalanceAtBlockImpl =
  retryOnceOnTimeout ... fmap unStringEncode ... API.getTicketBalanceAtBlock API.nodeMethods

getAllTicketBalancesAtBlockImpl
  :: (RunClient m, MonadUnliftIO m, MonadCatch m)
  => BlockId -> ContractAddress -> m [GetAllTicketBalancesResponse]
getAllTicketBalancesAtBlockImpl =
  retryOnceOnTimeout ... API.getAllTicketBalancesAtBlock API.nodeMethods

getDelegateImpl ::
  (RunClient m, MonadUnliftIO m, MonadCatch m) => BlockId -> L1Address -> m (Maybe KeyHash)
getDelegateImpl =
  retryOnceOnTimeout ... API.getDelegate API.nodeMethods

-- | Similar to 'API.getManagerKey', but retries once on timeout.
getManagerKeyImpl ::
  (RunClient m, MonadUnliftIO m, MonadCatch m) => BlockId -> ImplicitAddress -> m (Maybe PublicKey)
getManagerKeyImpl =
  retryOnceOnTimeout ... API.getManagerKey API.nodeMethods

runCodeImpl :: (RunClient m, MonadCatch m) => BlockId -> RunCode -> m RunCodeResult
runCodeImpl = API.runCode API.nodeMethods

getChainIdImpl :: (RunClient m, MonadCatch m) => m ChainId
getChainIdImpl = throwLeft $ parseChainId <$> API.getChainId API.nodeMethods

packDataImpl
  :: (T.ForbidOp t, RunClient m, MonadCatch m, MonadUnliftIO m)
  => BlockId -> T.Value t -> T.Notes t -> m Text
packDataImpl blkId val notes = retryOnceOnTimeout do
  PackDataResult{..} <- API.packData API.nodeMethods blkId PackData
     { pdData = toExpression val
     , pdType = toExpression notes
     , pdGas = Nothing
     }
  pure pdrPacked

waitForOperationImpl :: forall m
                      . (MonadUnliftIO m, HasTezosRpc m)
                     => m OperationHash
                     -> ClientEnv
                     -> m OperationHash
waitForOperationImpl opHash env = do
  finalHead <- getBlockHeader FinalHeadId
  hash <- opHash
  let limit = 10 -- blocks to wait for
      blockout = throwM $ WaitForOperationBlockout limit
      handle other = maybe (wait other >>= maybe blockout pure) pure
  -- NB: we can't really use race here because even on failure on one direction
  -- we still want to check the other.
  withAsync (searchForwards limit hash) \fwd ->
    withAsync (searchBackwards hash (bhLevel finalHead) HeadId) \bwd ->
      waitEither fwd bwd >>= \case
        Left res -> handle bwd res
        Right res -> handle fwd res
  where
    monitorHeads :: Word -> (BlockHeader -> m (MonitorHeadsStep OperationHash)) -> m (Maybe OperationHash)
    monitorHeads limit test = withRunInIO $ \runInIO ->
      withClientM (client runInIO) env (either throwM pure)
      where
        client runInIO = do
          src <- API.monitorHeads
          liftIO $ unSourceT src $ go limit runInIO
        go 0 _ = const $ pure Nothing
        go n runInIO = \case
          Stop -> pure Nothing
          Error str -> throwM $ WaitForOperationStreamingError (fromString str)
          Skip next -> go n runInIO next
          Yield a next -> runInIO (test a) >>= \case
            MonitorHeadsContinue -> go (n - 1) runInIO next
            MonitorHeadsStop res -> pure $ Just res
          Effect eff -> eff >>= go n runInIO
    searchBackwards :: OperationHash -> Int64 -> BlockId -> m (Maybe OperationHash)
    searchBackwards hash stopAtLevel blkHead = do
      checkBlock hash blkHead >>= \case
        MonitorHeadsContinue -> do
          header <- getBlockHeader blkHead
          if bhLevel header == stopAtLevel
          then pure Nothing
          else searchBackwards hash stopAtLevel $ BlockHashId $ bhPredecessor header
        MonitorHeadsStop res -> pure $ Just res
    searchForwards :: Word -> OperationHash -> m (Maybe OperationHash)
    searchForwards limit hash = monitorHeads limit (checkBlock hash . BlockHashId . bhHash)
    checkBlock :: OperationHash -> BlockId -> m (MonitorHeadsStep OperationHash)
    checkBlock hash blkId = do
      opHashes <- concat <$> getBlockOperationHashes blkId
      if hash `notElem` opHashes
      then pure MonitorHeadsContinue
      else do
        ops <- concat <$> getBlockOperations blkId
        case find (\op -> boHash op == unOperationHash hash) ops of
          Just op -> do
            for_ (boContents op)
              \(OperationRespWithMeta{..}) -> case unOperationMetadata =<< orwmMetadata of
                Just res -> case res of
                  OperationApplied{} -> pass
                  OperationFailed errs -> throwM $ UnexpectedRunErrors errs
                Nothing -> pass
            pure $ MonitorHeadsStop hash
          Nothing -> error "impossible"

----------------
-- Logging of requests and responses
----------------

-- | Convert a bytestring to a string assuming this bytestring stores
-- something readable in UTF-8 encoding.
fromBS :: ConvertUtf8 Text bs => bs -> Text
fromBS = decodeUtf8

ppRequestBody :: RequestBody -> Doc
ppRequestBody = build .
  \case
    RequestBodyLBS lbs -> fromBS lbs
    RequestBodyBS bs -> fromBS bs
    RequestBodySource {} -> "<body is not in memory>"

-- | Pretty print a @servant@'s request.
-- Note that we print only some part of t'Request': method, request
-- path and query, request body. We don't print other things that are
-- subjectively less interesting such as HTTP version or media type.
-- But feel free to add them if you want.
ppRequest :: Request -> Doc
ppRequest Request {..} =
  fromBS requestMethod |+ " " +| fromBS (Binary.toLazyByteString requestPath)
  |+ fromBS (renderQuery True $ toList requestQueryString) |+
  maybe mempty (mappend "\n" . ppRequestBody . fst) requestBody

logRequest :: WithClientLog env m => Request -> m ()
logRequest req = logDebug $ "RPC request: " +| ppRequest req |+ ""

-- | Pretty print a @servant@'s response.
-- Note that we print only status and body,
-- the rest looks not so interesting in our case.
--
-- If response is not human-readable text in UTF-8 encoding it will
-- print some garbage.  Apparently we don't make such requests for now.
ppResponse :: Response -> Doc
ppResponse Response {..} =
  statusCode responseStatusCode |+ " " +|
  fromBS (statusMessage responseStatusCode) |+ "\n" +|
  fromBS responseBody  |+ ""

logResponse :: WithClientLog env m => Response -> m ()
logResponse resp = logDebug $ "RPC response: " +| ppResponse resp |+ ""

----------------
-- Timeouts and retries
----------------

data TimeoutError = TimeoutError
  deriving stock Show

instance Buildable TimeoutError where
  build TimeoutError =
    "Timeout for action call was reached. Probably, something is wrong with \
    \testing environment."

instance Exception TimeoutError where
  displayException = pretty

-- | Helper function that retries a monadic action in case action hasn't succeed
-- in 'timeoutInterval'. In case retry didn't help, error that indicates
-- timeout is thrown.
retryOnTimeout :: (MonadUnliftIO m, MonadThrow m) => Bool -> m a -> m a
retryOnTimeout wasRetried action = do
  res <- timeout timeoutInterval action
  maybe (if wasRetried then throwM TimeoutError else retryOnTimeout True action)
    pure res

-- | Helper function that consider action failed in case of timeout,
-- because it's unsafe to perform some of the actions twice. E.g. performing two 'injectOperation'
-- action can lead to a situation when operation is injected twice.
failOnTimeout :: (MonadUnliftIO m, MonadThrow m) => m a -> m a
failOnTimeout = retryOnTimeout True

-- | Helper function that retries action once in case of timeout. If retry ended up with timeout
-- as well, action is considered failed. It's safe to retry read-only actions that don't update chain state
-- or @octez-client@ config/environment.
retryOnceOnTimeout :: (MonadUnliftIO m, MonadThrow m) => m a -> m a
retryOnceOnTimeout = retryOnTimeout False

-- | Timeout for 'retryOnTimeout', 'retryOnceOnTimeout' and 'failOnTimeout' helpers in microseconds.
timeoutInterval :: Int
timeoutInterval = 120 * 1e6

-- | Wait for a reasonable amount of time before retrying an action that failed
-- due to invalid counter.
-- The waiting time depends on protocol parameters.
waitBeforeRetry :: (MonadIO m, HasTezosRpc m, WithClientLog env m) => m ()
waitBeforeRetry = do
  i <- liftIO $
    (blockAwaitAmounts Unsafe.!!) <$> randomRIO (0, length blockAwaitAmounts - 1)
  logDebug $ "Invalid counter error occurred, retrying the request after " <> show i <> " blocks"
  ProtocolParameters {..} <- getProtocolParameters
  -- Invalid counter error may occur in case we try to perform multiple operations
  -- from the same address. We should try to wait different amount of times before retry
  -- in case there are multiple actions failed with invalid counter error.
  liftIO $ threadDelay $ i * Unsafe.fromIntegral @TezosNat @Int ppMinimalBlockDelay * 1e6
  where
    blockAwaitAmounts :: [Int]
    blockAwaitAmounts = [1..5]

-- | Retry action if it failed due to invalid counter (already used one).
handleInvalidCounterRpc :: MonadThrow m => m a -> ClientRpcError -> m a
handleInvalidCounterRpc retryAction = \case
  ClientInternalError (CounterInThePast {}) -> retryAction
  ClientInternalError (Failure msg)
    | "Counter" `isInfixOf` msg && "already used for contract" `isInfixOf` msg ->
      retryAction
  anotherErr -> throwM anotherErr