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

{-# OPTIONS_GHC -Wno-orphans #-}

-- | Module that defines some basic infrastructure
-- for faking tezos-node RPC interaction.
module TestM
  ( ContractData (..)
  , ContractState (..)
  , ContractStateBigMap (..)
  , Handlers (..)
  , FakeState (..)
  , TestError (..)
  , TestHandlers (..)
  , SomeAddressOrAlias (..)
  , TestM
  , TestT
  , defaultHandlers
  , defaultFakeState
  , runFakeTest
  , runFakeTestT
  , liftToFakeTest

  -- * Lens
  , fsImplicitsL
  ) 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 (pretty)

import Morley.Client
import Morley.Client.Logging (ClientLogAction)
import Morley.Client.RPC
import Morley.Micheline
import Morley.Tezos.Address
import Morley.Tezos.Address.Alias
import Morley.Tezos.Address.Kinds
import Morley.Tezos.Core
import Morley.Tezos.Crypto (KeyHash, PublicKey, Signature)
import Morley.Util.ByteString

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

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

-- | Reader environment to interact with the fake state.
data Handlers m = Handlers
  { -- HasTezosRpc
    hGetBlockHash :: BlockId -> m BlockHash
  , hGetCounter :: BlockId -> ImplicitAddress -> m TezosInt64
  , hGetBlockHeader :: BlockId -> m BlockHeader
  , hGetBlockConstants :: BlockId -> m BlockConstants
  , hGetBlockOperations :: BlockId -> m [[BlockOperation]]
  , hGetScriptSizeAtBlock :: BlockId -> CalcSize -> m ScriptSize
  , hGetBlockOperationHashes :: BlockId -> m [[OperationHash]]
  , 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 -> ContractAddress -> m OriginationScript
  , hGetContractBigMap :: BlockId -> ContractAddress -> 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 -> ImplicitAddress -> m (Maybe PublicKey)
  , hGetDelegateAtBlock :: BlockId -> ContractAddress -> m (Maybe KeyHash)

  -- HasTezosClient
  , hSignBytes :: ImplicitAddressOrAlias -> Maybe ScrubbedBytes -> ByteString -> m Signature
  , hGenKey :: ImplicitAlias -> m ImplicitAddress
  , hGenFreshKey :: ImplicitAlias -> m ImplicitAddress
  , hRevealKey :: ImplicitAlias -> Maybe ScrubbedBytes -> m ()
  , hWaitForOperation :: m OperationHash -> m OperationHash
  , hRememberContract :: Bool -> ContractAddress -> ContractAlias -> m ()
  , hResolveAddressMaybe :: forall kind. AddressOrAlias kind -> m (Maybe (KindedAddress kind))
  , hGetAlias :: forall kind. AddressOrAlias kind -> m (Alias kind)
  , hGetKeyPassword :: ImplicitAddress -> m (Maybe ScrubbedBytes)
  , hRegisterDelegate :: ImplicitAlias -> 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"
  , hGetScriptSizeAtBlock = \_ _ -> throwM $ UnexpectedRpcCall "getScriptSizeAtBlock"
  , hGetBlockOperations = \_ -> throwM $ UnexpectedRpcCall "getBlockOperations"
  , hGetBlockOperationHashes = \_ -> throwM $ UnexpectedRpcCall "hGetBlockOperationHashes"
  , 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"
  , hResolveAddressMaybe = \_ -> throwM $ UnexpectedRpcCall "resolveAddressMaybe"
  , hGetAlias = \_ -> throwM $ UnexpectedRpcCall "getAlias"
  , hGetKeyPassword = \_ -> throwM $ UnexpectedClientCall "getKeyPassword"
  , hRegisterDelegate = \_ _ -> throwM $ UnexpectedClientCall "registerDelegate"

  , hLogAction = mempty
  }

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

data ContractData k where
  ContractData :: OriginationScript -> Maybe ContractStateBigMap -> ContractData 'AddressKindContract
  ImplicitContractData :: Maybe PublicKey -> ContractData 'AddressKindImplicit

-- | 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 AddressMap k = Map (KindedAddress k) (ContractState k)

-- | Type to represent chain state in mock tests.
data FakeState = FakeState
  { fsContracts :: AddressMap 'AddressKindContract
  , fsImplicits :: AddressMap 'AddressKindImplicit
  , fsHeadBlock :: BlockHash
  -- ^ Hash of the `head` block
  , fsFinalHeadBlock :: BlockHash
  -- ^ Hash of the `head~2` block
  , fsBlockConstants :: BlockId -> BlockConstants
  , fsProtocolParameters :: ProtocolParameters
  }

defaultFakeState :: FakeState
defaultFakeState = FakeState
  { fsContracts = mempty
  , fsImplicits = mempty
  , fsHeadBlock = BlockHash "HEAD"
  , fsFinalHeadBlock = BlockHash "HEAD~2"
  , fsBlockConstants = \blkId -> BlockConstants
    { bcProtocol = "PROTOCOL"
    , bcChainId = "CHAIN_ID"
    , bcHeader = BlockHeaderNoHash
      { bhnhTimestamp = posixSecondsToUTCTime 0
      , bhnhLevel = 0
      , bhnhPredecessor = BlockHash "PREV_HASH"
      }
    , bcHash = BlockHash $ pretty blkId
    }
  , fsProtocolParameters = ProtocolParameters 257 1040000 60000 15
                                              (TezosMutez [tz|250u|])
  }

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

type TestM = TestT Identity

runFakeTestT
  :: forall a m. Monad m =>
     Handlers (TestT m) -> FakeState -> TestT m a -> m (Either SomeException a)
runFakeTestT handlers fakeState action =
  runCatchT $ runReaderT (evalStateT action fakeState) (TestHandlers handlers)

runFakeTest
  :: forall a. Handlers TestM -> FakeState -> TestM a -> Either SomeException a
runFakeTest =
  runIdentity ... runFakeTestT

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

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

-- | Various fake test errors.
data TestError
  = AlreadyRevealed Address
  | UnexpectedRpcCall Text
  | UnexpectedClientCall Text
  | UnknownContract SomeAddressOrAlias
  | ContractDoesntHaveBigMap Address
  | InvalidChainId
  | InvalidProtocol
  | InvalidBranch BlockHash
  | 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
  rememberContract replaceExisting addr alias = do
    h <- getHandler hRememberContract
    h replaceExisting addr alias
  resolveAddressMaybe addr = do
    h <- ask >>= \t -> pure $ hResolveAddressMaybe $ unTestHandlers t
    h addr
  getAlias originator = do
    h <- ask >>= \t -> pure $ hGetAlias $ unTestHandlers t
    h originator
  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
  getScriptSizeAtBlock block script = do
    h <- getHandler hGetScriptSizeAtBlock
    h block script
  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
  getBlockOperationHashes block = do
    h <- getHandler hGetBlockOperationHashes
    h block
  waitForOperation opHash = do
    h <- getHandler hWaitForOperation
    h opHash

makeLensesFor [("fsImplicits", "fsImplicitsL")] ''FakeState