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

{-# OPTIONS_GHC -Wno-orphans #-}

-- | Module that defines some basic infrastructure for faking @octez-node@ RPC
-- interaction.
module TestM
  ( AccountData (..)
  , AccountState (..)
  , ContractStateBigMap (..)
  , Handlers (..)
  , FakeState (..)
  , TestError (..)
  , TestHandlers (..)
  , TestM
  , TestT
  , defaultHandlers
  , defaultFakeState
  , runFakeTest
  , runFakeTestT
  , liftToFakeTest

  -- * Lens
  , fsImplicitsL
  , fsContractsL
  , asAccountDataL
  ) where

import Colog.Core.Class (HasLog(..))
import Colog.Message (Message)
import Control.Lens (makeLensesFor)
import Control.Monad.Catch.Pure (CatchT(..))
import Data.Aeson (eitherDecodeStrict)
import Data.ByteArray (ScrubbedBytes)
import Data.ByteString (readFile)
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 -> L1Address -> m (Maybe KeyHash)
  , hGetTicketBalanceAtBlock :: BlockId -> Address -> GetTicketBalance ->  m Natural
  , hGetAllTicketBalancesAtBlock :: BlockId -> ContractAddress -> m [GetAllTicketBalancesResponse]

  -- HasTezosClient
  , hSignBytes :: ImplicitAddressWithAlias -> Maybe ScrubbedBytes -> ByteString -> m Signature
  , hGenKey :: ImplicitAlias -> m ImplicitAddressWithAlias
  , hGenFreshKey :: ImplicitAlias -> m ImplicitAddressWithAlias
  , hGetPublicKey :: ImplicitAddressWithAlias -> m PublicKey
  , hWaitForOperation :: m OperationHash -> m OperationHash
  , hRememberContract :: AliasBehavior -> ContractAddress -> ContractAlias -> m ()
  , hGetAliasesAndAddresses :: m [(Text, Text)]
  , hGetKeyPassword :: ImplicitAddressWithAlias -> m (Maybe ScrubbedBytes)

  -- 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"
  , hGetPublicKey = \_ -> throwM $ UnexpectedClientCall "getPublicKey"
  , hWaitForOperation = \_ -> throwM $ UnexpectedRpcCall "waitForOperation"
  , hRememberContract = \_ _ _ -> throwM $ UnexpectedClientCall "rememberContract"
  , hGetAliasesAndAddresses = throwM $ UnexpectedRpcCall "getAliasesAndAddresses"
  , hGetKeyPassword = \_ -> throwM $ UnexpectedClientCall "getKeyPassword"

  , hGetTicketBalanceAtBlock = \_ _ _ -> throwM $ UnexpectedClientCall "getTicketBalanceAtBlock"
  , hGetAllTicketBalancesAtBlock = \_ _ -> throwM $ UnexpectedClientCall "getAllTicketBalancesAtBlock"

  , hLogAction = mempty
  }

-- | Type to represent account state in the @FakeState@.
-- This type can represent both implicit accounts and contracts.
data AccountState k = AccountState
  { asCounter :: TezosInt64
  , asAlias :: Alias k
  , asAccountData :: AccountData k
  }

data AccountData k where
  ContractData ::
    { cdScript :: OriginationScript
    , cdBigMap :: Maybe ContractStateBigMap
    } -> AccountData 'AddressKindContract
  ImplicitData ::
    { idPublicKey :: PublicKey
    , idManagerKey :: Maybe PublicKey
    } -> AccountData 'AddressKindImplicit

-- | Type to represent big_map in @AccountState@.
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) (AccountState 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 = $(do
      -- This file is the output of /chains/main/blocks/head/context/constants
      -- testnet RPC endpoint. Update it with protocol revisions as needed using
      -- scripts/update-morley-client-test-protocol-parameters.sh
      content <- liftIO $ readFile "test/data/constants.json"
      case eitherDecodeStrict @ProtocolParameters content of
        Left err -> fail err
        Right{} -> [|unsafe $ eitherDecodeStrict @ProtocolParameters content|]
          -- ok to use unsafe here as we've just successfully decoded the value.
      )
  }

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
  = UnexpectedRpcCall Text
  | UnexpectedClientCall Text
  | UnknownAccount Address
  | 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
  rememberContract replaceExisting addr alias = do
    h <- getHandler hRememberContract
    h replaceExisting addr alias
  getAliasesAndAddresses = do
    join $ getHandler hGetAliasesAndAddresses
  getKeyPassword addr = do
    h <- getHandler hGetKeyPassword
    h addr
  getPublicKey addrOrAlias = do
    h <- getHandler hGetPublicKey
    h addrOrAlias

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
  getTicketBalanceAtBlock block addr args = do
    h <- getHandler hGetTicketBalanceAtBlock
    h block addr args
  getAllTicketBalancesAtBlock block addr = do
    h <- getHandler hGetAllTicketBalancesAtBlock
    h block addr

makeLensesFor [("fsImplicits", "fsImplicitsL"), ("fsContracts", "fsContractsL")] ''FakeState
makeLensesFor [("asAccountData", "asAccountDataL")] ''AccountState