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

{-# OPTIONS_GHC -Wno-orphans #-}

-- | Module that defines some basic infrastructure
-- for mocking tezos-node RPC interaction.
module TestM
  ( ContractData (..)
  , ContractState (..)
  , ContractStateBigMap (..)
  , Handlers (..)
  , MockState (..)
  , TestError (..)
  , TestHandlers (..)
  , TestM
  , TestT
  , defaultHandlers
  , defaultMockState
  , runMockTest
  , runMockTestT
  , liftToMockTest

  -- * Lens
  , msContractsL
  ) where

import Colog.Core.Class (HasLog(..))
import Colog.Message (Message)
import Control.Lens (makeLensesFor)
import Control.Monad.Catch.Pure (CatchT(..))
import Data.ByteArray (ScrubbedBytes)
import Data.Time.Clock.POSIX (posixSecondsToUTCTime)

import Fmt
import Morley.Client
import Morley.Client.Logging (ClientLogAction)
import Morley.Client.RPC
import Morley.Client.TezosClient (CalcOriginationFeeData, CalcTransferFeeData, TezosClientConfig)
import Morley.Micheline
import Morley.Michelson.Typed.Scope (UntypedValScope)
import Morley.Tezos.Address
import Morley.Tezos.Core
import Morley.Tezos.Crypto (KeyHash, PublicKey, SecretKey, Signature)
import Morley.Util.ByteString

-- | A test-specific orphan.
instance IsString Alias where
  fromString = mkAlias . fromString

-- | Reader environment to interact with the mock state.
data Handlers m = Handlers
  { -- HasTezosRpc
    hGetBlockHash :: BlockId -> m Text
  , hGetCounter :: BlockId -> Address -> m TezosInt64
  , hGetBlockHeader :: BlockId -> m BlockHeader
  , hGetBlockConstants :: BlockId -> m BlockConstants
  , hGetBlockOperations :: BlockId -> m [[BlockOperation]]
  , hGetProtocolParameters :: BlockId -> m ProtocolParameters
  , hRunOperation :: BlockId -> RunOperation -> m RunOperationResult
  , hPreApplyOperations :: BlockId -> [PreApplyOperation] -> m [RunOperationResult]
  , hForgeOperation :: BlockId -> ForgeOperation -> m HexJSONByteString
  , hInjectOperation :: HexJSONByteString -> m OperationHash
  , hGetContractScript :: BlockId -> Address -> m OriginationScript
  , hGetContractBigMap :: BlockId -> Address -> GetBigMap -> m GetBigMapResult
  , hGetBigMapValue :: BlockId -> Natural -> Text -> m Expression
  , hGetBigMapValues :: BlockId -> Natural -> Maybe Natural -> Maybe Natural -> m Expression
  , hGetBalance :: BlockId -> Address -> m Mutez
  , hRunCode :: BlockId -> RunCode -> m RunCodeResult
  , hGetChainId :: m ChainId
  , hGetManagerKey :: BlockId -> Address -> m (Maybe PublicKey)
  , hGetDelegateAtBlock :: BlockId -> Address -> m (Maybe KeyHash)

  -- HasTezosClient
  , hSignBytes :: AddressOrAlias -> Maybe ScrubbedBytes -> ByteString -> m Signature
  , hGenKey :: AliasOrAliasHint -> m Address
  , hGenFreshKey :: AliasOrAliasHint -> m Address
  , hRevealKey :: Alias -> Maybe ScrubbedBytes -> m ()
  , hWaitForOperation :: OperationHash -> m ()
  , hRememberContract :: Bool -> Address -> AliasOrAliasHint -> m ()
  , hImportKey :: Bool -> AliasOrAliasHint -> SecretKey -> m Alias
  , hResolveAddressMaybe :: AddressOrAlias -> m (Maybe Address)
  , hGetAlias :: AddressOrAlias -> m Alias
  , hGetPublicKey :: AddressOrAlias -> m PublicKey
  , hGetTezosClientConfig :: m TezosClientConfig
  , hCalcTransferFee
    :: AddressOrAlias -> Maybe ScrubbedBytes -> TezosInt64 -> [CalcTransferFeeData] -> m [TezosMutez]
  , hCalcOriginationFee
    :: forall cp st. UntypedValScope st => CalcOriginationFeeData cp st -> m TezosMutez
  , hGetKeyPassword :: Address -> m (Maybe ScrubbedBytes)
  , hRegisterDelegate :: AliasOrAliasHint -> Maybe ScrubbedBytes -> m ()

  -- HasLog
  , hLogAction :: ClientLogAction m
  }

defaultHandlers :: Monad m => Handlers (TestT m)
defaultHandlers = Handlers
  { hGetBlockHash = \_ -> throwM $ UnexpectedRpcCall "getHeadBlock"
  , hGetCounter = \_ _ -> throwM $ UnexpectedRpcCall "getCounter"
  , hGetBlockHeader = \_ -> throwM $ UnexpectedRpcCall "getBlockHeader"
  , hGetBlockConstants = \_ -> throwM $ UnexpectedRpcCall "getBlockConstants"
  , hGetBlockOperations = \_ -> throwM $ UnexpectedRpcCall "getBlockOperations"
  , hGetProtocolParameters = \_ -> throwM $ UnexpectedRpcCall "getProtocolParameters"
  , hRunOperation = \_ _ -> throwM $ UnexpectedRpcCall "runOperation"
  , hPreApplyOperations = \_ _ -> throwM $ UnexpectedRpcCall "preApplyOperations"
  , hForgeOperation = \_ _ -> throwM $ UnexpectedRpcCall "forgeOperation"
  , hInjectOperation = \_ -> throwM $ UnexpectedRpcCall "injectOperation"
  , hGetContractScript = \_ _ -> throwM $ UnexpectedRpcCall "getContractScript"
  , hGetContractBigMap = \_ _ _ -> throwM $ UnexpectedRpcCall "getContractBigMap"
  , hGetBigMapValue = \_ _ _ -> throwM $ UnexpectedRpcCall "getBigMapValue"
  , hGetBigMapValues = \_ _ _ _ -> throwM $ UnexpectedRpcCall "getBigMapValues"
  , hGetBalance = \_ _ -> throwM $ UnexpectedRpcCall "getBalance"
  , hRunCode = \_ _ -> throwM $ UnexpectedRpcCall "runCode"
  , hGetChainId = throwM $ UnexpectedRpcCall "getChainId"
  , hGetManagerKey = \_ _ -> throwM $ UnexpectedRpcCall "getManagerKey"
  , hGetDelegateAtBlock = \_ _ -> throwM $ UnexpectedRpcCall "getDelegateAtBlock"

  , hSignBytes = \_ _ _ -> throwM $ UnexpectedClientCall "signBytes"
  , hGenKey = \_ -> throwM $ UnexpectedClientCall "genKey"
  , hGenFreshKey = \_ -> throwM $ UnexpectedRpcCall "genFreshKey"
  , hRevealKey = \_ _ -> throwM $ UnexpectedClientCall "revealKey"
  , hWaitForOperation = \_ -> throwM $ UnexpectedRpcCall "waitForOperation"
  , hRememberContract = \_ _ _ -> throwM $ UnexpectedClientCall "hRememberContract"
  , hImportKey = \_ _ _ -> throwM $ UnexpectedClientCall "importKey"
  , hResolveAddressMaybe = \_ -> throwM $ UnexpectedRpcCall "resolveAddressMaybe"
  , hGetAlias = \_ -> throwM $ UnexpectedRpcCall "getAlias"
  , hGetPublicKey = \_ -> throwM $ UnexpectedRpcCall "getPublicKey"
  , hGetTezosClientConfig = throwM $ UnexpectedClientCall "getTezosClientConfig"
  , hCalcTransferFee = \_ _ _ _ -> throwM $ UnexpectedClientCall "calcTransferFee"
  , hCalcOriginationFee = \_ -> throwM $ UnexpectedClientCall "calcOriginationFee"
  , hGetKeyPassword = \_ -> throwM $ UnexpectedClientCall "getKeyPassword"
  , hRegisterDelegate = \_ _ -> throwM $ UnexpectedClientCall "registerDelegate"

  , hLogAction = mempty
  }

-- | Type to represent contract state in the @MockState@.
-- This type can represent both implicit accounts and contracts.
data ContractState = ContractState
  { csCounter :: TezosInt64
  , csAlias :: Alias
  , csContractData :: ContractData
  }

data ContractData
  = ContractData OriginationScript (Maybe ContractStateBigMap)
  | ImplicitContractData (Maybe PublicKey)

-- | Type to represent big_map in @ContractState@.
data ContractStateBigMap = ContractStateBigMap
  { csbmKeyType :: Expression
  , csbmValueType :: Expression
  , csbmMap :: Map Text ByteString
  -- ^ Real tezos bigmap also has deserialized keys and values
  , csbmId :: Natural
  -- ^ The big_map's ID
  }

newtype TestHandlers m = TestHandlers {unTestHandlers :: Handlers (TestT m)}

-- | Type to represent chain state in mock tests.
data MockState = MockState
  { msContracts :: Map Address ContractState
  , msHeadBlock :: Text
  -- ^ Hash of the `head` block
  , msFinalHeadBlock :: Text
  -- ^ Hash of the `head~2` block
  , msBlockConstants :: BlockId -> BlockConstants
  , msProtocolParameters :: ProtocolParameters
  }

defaultMockState :: MockState
defaultMockState = MockState
  { msContracts = mempty
  , msHeadBlock = "HEAD"
  , msFinalHeadBlock = "HEAD~2"
  , msBlockConstants = \blkId -> BlockConstants
    { bcProtocol = "PROTOCOL"
    , bcChainId = "CHAIN_ID"
    , bcHeader = BlockHeaderNoHash
      { bhnhTimestamp = posixSecondsToUTCTime 0
      , bhnhLevel = 0
      , bhnhPredecessor = BlockHash "PREV_HASH"
      }
    , bcHash = BlockHash $ pretty blkId
    }
  , msProtocolParameters = ProtocolParameters 257 1040000 60000 15
                                              (TezosMutez [tz|250u|])
  }

type TestT m = StateT MockState (ReaderT (TestHandlers m) (CatchT m))

type TestM = TestT Identity

runMockTestT
  :: forall a m. Monad m =>
     Handlers (TestT m) -> MockState -> TestT m a -> m (Either SomeException a)
runMockTestT handlers mockState action =
  runCatchT $ runReaderT (evalStateT action mockState) (TestHandlers handlers)

runMockTest
  :: forall a. Handlers TestM -> MockState -> TestM a -> Either SomeException a
runMockTest =
  runIdentity ... runMockTestT

getHandler :: Monad m => (Handlers (TestT m) -> fn) -> TestT m fn
getHandler fn = fn . unTestHandlers <$> ask

liftToMockTest :: Monad m => m a -> TestT m a
liftToMockTest = lift . lift . lift

-- | Various mock test errors.
data TestError
  = AlreadyRevealed Address
  | UnexpectedRpcCall Text
  | UnexpectedClientCall Text
  | UnknownContract AddressOrAlias
  | UnexpectedImplicitContract Address
  | ContractDoesntHaveBigMap Address
  | CantRevealContract Address
  | InvalidChainId
  | InvalidProtocol
  | InvalidBranch Text
  | CounterMismatch
  deriving stock Show

instance Exception TestError

instance HasLog (TestHandlers m) Message (TestT m) where
  getLogAction = hLogAction . unTestHandlers
  setLogAction action (TestHandlers handlers) =
    TestHandlers $ handlers { hLogAction = action }

instance Monad m => HasTezosClient (TestT m) where
  signBytes alias mbPassword op = do
    h <- getHandler hSignBytes
    h alias mbPassword op
  genKey alias = do
    h <- getHandler hGenKey
    h alias
  genFreshKey alias = do
    h <- getHandler hGenFreshKey
    h alias
  revealKey alias mbPassword = do
    h <- getHandler hRevealKey
    h alias mbPassword
  waitForOperation op = do
    h <- getHandler hWaitForOperation
    h op
  rememberContract replaceExisting addr alias = do
    h <- getHandler hRememberContract
    h replaceExisting addr alias
  importKey replaceExisting alias key = do
    h <- getHandler hImportKey
    h replaceExisting alias key
  resolveAddressMaybe addr = do
    h <- getHandler hResolveAddressMaybe
    h addr
  getAlias originator = do
    h <- getHandler hGetAlias
    h originator
  getPublicKey alias = do
    h <- getHandler hGetPublicKey
    h alias
  getTezosClientConfig =
    join $ getHandler hGetTezosClientConfig
  calcTransferFee from mbPassword burnCap transferDatas = do
    h <- getHandler hCalcTransferFee
    h from mbPassword burnCap transferDatas
  calcOriginationFee origData = do
    h <- getHandler hCalcOriginationFee
    h origData
  getKeyPassword addr = do
    h <- getHandler hGetKeyPassword
    h addr
  registerDelegate kh pw = do
    h <- getHandler hRegisterDelegate
    h kh pw

instance Monad m => HasTezosRpc (TestT m) where
  getBlockHash block = do
    h <- getHandler hGetBlockHash
    h block
  getCounterAtBlock block addr = do
    h <- getHandler hGetCounter
    h block addr
  getBlockHeader block = do
    h <- getHandler hGetBlockHeader
    h block
  getBlockConstants block = do
    h <- getHandler hGetBlockConstants
    h block
  getBlockOperations block = do
    h <- getHandler hGetBlockOperations
    h block
  getProtocolParametersAtBlock block = do
    h <- getHandler hGetProtocolParameters
    h block
  runOperationAtBlock block op = do
    h <- getHandler hRunOperation
    h block op
  preApplyOperationsAtBlock block ops = do
    h <- getHandler hPreApplyOperations
    h block ops
  forgeOperationAtBlock block op = do
    h <- getHandler hForgeOperation
    h block op
  injectOperation op = do
    h <- getHandler hInjectOperation
    h op
  getContractScriptAtBlock block addr = do
    h <- getHandler hGetContractScript
    h block addr
  getContractStorageAtBlock blockId addr = do
    h <- getHandler hGetContractScript
    osStorage <$> h blockId addr
  getContractBigMapAtBlock block addr getBigMap = do
    h <- getHandler hGetContractBigMap
    h block addr getBigMap
  getBigMapValueAtBlock blockId bigMapId scriptExpr = do
    h <- getHandler hGetBigMapValue
    h blockId bigMapId scriptExpr
  getBigMapValuesAtBlock blockId bigMapId mbOffset mbLength = do
    h <- getHandler hGetBigMapValues
    h blockId bigMapId mbOffset mbLength
  getBalanceAtBlock block addr = do
    h <- getHandler hGetBalance
    h block addr
  runCodeAtBlock block r = do
    h <- getHandler hRunCode
    h block r
  getChainId = join (getHandler hGetChainId)
  getManagerKeyAtBlock block addr = do
    h <- getHandler hGetManagerKey
    h block addr
  getDelegateAtBlock block addr = do
    h <- getHandler hGetDelegateAtBlock
    h block addr

makeLensesFor [("msContracts", "msContractsL")] ''MockState