-- SPDX-FileCopyrightText: 2020 Tocqueville Group
--
-- SPDX-License-Identifier: LicenseRef-MIT-TQ

-- | Module with various helpers that are used in morley-client mock tests.
module Test.Util
  ( chainOperationHandlers
  , dumbContractState
  , dumbImplicitContractState
  , dumbManagerKey
  , mapToContractStateBigMap
  , handleGetBigMapValue

    -- * TemplateHaskell test helpers
  , shouldCompileTo
  , shouldCompileIgnoringInstance
    -- * Internals
  , handleRunOperationInternal
  , assertHeadBlockId
  ) where

import Prelude hiding (Type)

import Control.Exception.Safe (throwString)
import Control.Lens (at, (?~))
import Data.Aeson (encode)
import Data.ByteArray (ScrubbedBytes)
import qualified Data.ByteString.Lazy as LBS (toStrict)
import qualified Data.Generics as SYB
import Data.Map as Map (elems, fromList, insert, lookup, toList)
import Data.Singletons (demote)
import Fmt ((+|), (|+))
import Language.Haskell.TH (pprint)
import Language.Haskell.TH.Syntax
  (Dec(..), Name, Q, TyVarBndr(..), Type(..), mkName, nameBase, runQ)
import Network.HTTP.Types.Status (status404)
import Network.HTTP.Types.Version (http20)
import Servant.Client.Core
  (BaseUrl(..), ClientError(..), RequestF(..), ResponseF(..), Scheme(..), defaultRequest)
import Test.Tasty.HUnit (Assertion, (@?=))
import Text.Hex (encodeHex)
import qualified Text.Show (show)

import Lorentz as L (compileLorentz, drop)
import Lorentz.Constraints
import Lorentz.Pack
import Morley.Client.RPC.Types
import Morley.Client.TezosClient.Types
import Morley.Micheline
import Morley.Michelson.Typed
import Morley.Tezos.Address
import Morley.Tezos.Core
import Morley.Tezos.Crypto
import qualified Morley.Tezos.Crypto.Ed25519 as Ed25519
import Morley.Util.ByteString
import TestM

-- | Function to convert given map to big map representation
-- used in mock state.
mapToContractStateBigMap
  :: forall k v. (NicePackedValue k, NicePackedValue v)
  => BigMapId k v -> Map k v -> ContractStateBigMap
mapToContractStateBigMap (BigMapId bigMapId) map' = ContractStateBigMap
  { csbmKeyType = toExpression $ demote @(ToT k)
  , csbmValueType = toExpression $ demote @(ToT v)
  , csbmMap = fromList $
    map (bimap (encodeBase58Check . valueToScriptExpr) lEncodeValue) $
    Map.toList map'
  , csbmId = bigMapId
  }

-- | Initial simple contract mock state.
dumbContractState :: ContractState
dumbContractState = ContractState
  { csCounter = 100500
  , csAlias = "genesis2"
  , csContractData = ContractData
      OriginationScript
        { osCode    = toExpression $ compileLorentz L.drop
        , osStorage = toExpression $ toVal ()
        }
      Nothing
  }

dumbImplicitContractState :: ContractState
dumbImplicitContractState = ContractState
  { csCounter = 100500
  , csAlias = "genesis1"
  , csContractData = ImplicitContractData Nothing
  }

-- | Mock handlers used for transaction sending and contract origination.
chainOperationHandlers :: Monad m => Handlers (TestT m)
chainOperationHandlers = defaultHandlers
  { hGetBlockHash = handleGetBlockHash
  , hGetCounter = handleGetCounter
  , hGetBlockConstants = handleGetBlockConstants
  , hGetProtocolParameters = handleGetProtocolParameters
  , hRunOperation = handleRunOperation
  , hPreApplyOperations = mapM . handlePreApplyOperation
  , hForgeOperation = \blkId arg -> do
      assertHeadBlockId blkId
      pure . HexJSONByteString . LBS.toStrict . encode $ arg
  , hInjectOperation = pure . OperationHash . (<> "_injected") . encodeHex . unHexJSONByteString
  , hGetContractScript = handleGetContractScript
  , hSignBytes =
    \_ _ -> pure . SignatureEd25519 . Ed25519.sign testSecretKey
  , hWaitForOperation = const pass
  , hGetAlias = handleGetAlias
  , hResolveAddressMaybe = handleResolveAddressMaybe
  , hRememberContract = handleRememberContract
  , hCalcTransferFee = \_ _ _ _ -> pure $ [TezosMutez $ toMutez 100500]
  , hCalcOriginationFee = \_ -> pure $ TezosMutez $ toMutez 100500
  , hGetKeyPassword = \_ -> pure Nothing
  , hGenKey = handleGenKey
  , hGetManagerKey = handleGetManagerKey
  , hRevealKey = handleRevealKey
  }
  where
    testSecretKey :: Ed25519.SecretKey
    testSecretKey = Ed25519.detSecretKey "\001\002\003\004"

mkRunOperationResult :: [Address] -> RunOperationResult
mkRunOperationResult originatedContracts = RunOperationResult
  { rrOperationContents =
    one $ OperationContent $ RunMetadata
      { rmOperationResult = OperationApplied $
        AppliedResult 100500 100500 100500 originatedContracts 0
      , rmInternalOperationResults = []
      }
  }

handleGetBlockHash :: Monad m => BlockId -> TestT m Text
handleGetBlockHash blkId = do
  assertHeadBlockId blkId
  MockState{..} <- get
  pure $ msHeadBlock

handleGetCounter
  :: ( MonadState MockState m
     , MonadThrow m
     )
  => BlockId -> Address -> m TezosInt64
handleGetCounter blk addr = do
  assertHeadBlockId blk
  MockState{..} <- get
  case lookup addr msContracts of
    Nothing -> throwM $ UnknownContract $ AddressResolved addr
    Just ContractState{..} -> pure $ csCounter

handleGetBlockConstants
  :: MonadState MockState m
  => anything -> m BlockConstants
handleGetBlockConstants _ = do
  MockState{..} <- get
  pure $ msBlockConstants

handleGetProtocolParameters
  :: (MonadState MockState m, MonadThrow m)
  => BlockId -> m ProtocolParameters
handleGetProtocolParameters blk = do
  assertHeadBlockId blk
  MockState{..} <- get
  pure $ msProtocolParameters

handleRunOperation :: Monad m => BlockId -> RunOperation -> TestT m RunOperationResult
handleRunOperation blk RunOperation{..} = do
  assertHeadBlockId blk
  MockState{..} <- get
  -- Ensure that passed chain id matches with one that mock state has
  unless (roChainId == bcChainId msBlockConstants) (throwM $ InvalidChainId)
  originatedContracts <- handleRunOperationInternal roOperation
  pure $ mkRunOperationResult originatedContracts

handlePreApplyOperation :: Monad m => BlockId -> PreApplyOperation -> TestT m RunOperationResult
handlePreApplyOperation blk PreApplyOperation{..} = do
  assertHeadBlockId blk
  MockState{..} <- get
  -- Ensure that passed protocol matches with one that mock state has
  unless (paoProtocol == bcProtocol msBlockConstants) (throwM $ InvalidProtocol)
  originatedContracts <- concatMapM handleTransactionOrOrigination paoContents
  pure $ mkRunOperationResult originatedContracts

handleRunOperationInternal :: Monad m => RunOperationInternal -> TestT m [Address]
handleRunOperationInternal RunOperationInternal{..} = do
  concatMapM handleTransactionOrOrigination roiContents

handleTransactionOrOrigination
  :: Monad m => Either TransactionOperation OriginationOperation -> TestT m [Address]
handleTransactionOrOrigination op = do
  MockState{..} <- get
  case op of
    -- Ensure that transaction sender exists
    Left TransactionOperation{..} -> case lookup codSource msContracts of
      Nothing -> throwM $ UnknownContract $ AddressResolved codSource
      Just ContractState{..} -> do
        -- Ensure that sender counter matches
        unless (csCounter + 1 == codCounter) (throwM CounterMismatch)
        case lookup toDestination msContracts of
          Nothing -> throwM $ UnknownContract $ AddressResolved toDestination
          Just _ -> pure []
      where
        CommonOperationData{..} = toCommonData
    -- Ensure that originator exists
    Right OriginationOperation{..} -> case lookup codSource msContracts of
      Nothing -> throwM $ UnknownContract $ AddressResolved codSource
      Just ContractState{..} -> do
        -- Ensure that originator counter matches
        unless (csCounter + 1 == codCounter) (throwM CounterMismatch)
        pure [dummyContractAddr]
      where
        CommonOperationData{..} = ooCommonData
        dummyContractAddr = unsafeParseAddress "KT1LZwEZqbqtLYhwzaidBp6So9LgYDpkpEv7"

-- We don't pass non-head block anywhere in @morley-client@, this feature
-- exists only for external users.
assertHeadBlockId :: MonadThrow m => BlockId -> m ()
assertHeadBlockId blockId = unless (blockId == HeadId) $
  throwString "Accessing non-head block is not supported in tests"

handleGetContractScript
  :: ( MonadState MockState m
     , MonadThrow m
     )
  => BlockId
  -> Address
  -> m OriginationScript
handleGetContractScript blockId addr = do
  assertHeadBlockId blockId
  MockState{..} <- get
  case lookup addr msContracts of
    Nothing -> throwM $ err404 path
    Just ContractState{..} -> case csContractData of
      ImplicitContractData _ -> throwM $ UnexpectedImplicitContract addr
      ContractData script _ -> pure script
  where
    path = "/chains/main/blocks/head/context/contracts/" <> formatAddress addr <> "/script"

handleGetBigMapValue :: Monad m => BlockId -> Natural -> Text -> TestT m Expression
handleGetBigMapValue blockId bigMapId scriptExpr = do
  assertHeadBlockId blockId
  st <- get

  let allBigMaps :: [ContractStateBigMap] =
        catMaybes $
          Map.elems (msContracts st) <&> \cs -> case (csContractData cs) of
            ContractData _ bigMapMaybe -> bigMapMaybe
            ImplicitContractData _ -> Nothing

  -- Check if a big_map with the given ID exists and, if so, check
  -- whether the giv en key exists.
  case find (\bigMap -> csbmId bigMap == bigMapId) allBigMaps of
    Nothing -> throwM $ err404 path
    Just bigMap ->
      case lookup scriptExpr (csbmMap bigMap ) of
        Nothing -> throwM $ err404 path
        Just serializedValue -> pure $ decodeExpression serializedValue
  where
    path = "/chains/main/blocks/head/context/big_maps/" <> show bigMapId <> "/" <> scriptExpr

-- Here we have an alias with the prefix already added,
-- so we can use 'Alias' instead 'AliasHint'.
getAlias :: AliasOrAliasHint -> Alias
getAlias = \case
  AnAlias x -> x
  AnAliasHint hint -> unsafeCoerceAliasHintToAlias hint

handleRememberContract :: Monad m => Bool -> Address -> AliasOrAliasHint -> TestT m ()
handleRememberContract replaceExisting addr (getAlias -> alias) = do
  let
    cs = dumbContractState { csAlias = alias }
    remember addr' cs' MockState{..} =
      modify $ \s -> s { msContracts = insert addr' cs' msContracts }

  st@MockState{..} <- get
  case lookup addr msContracts of
    Nothing -> remember addr cs st
    _       -> bool pass (remember addr cs st) replaceExisting

handleGenKey :: Monad m => AliasOrAliasHint -> TestT m Address
handleGenKey (getAlias -> alias) = do
  let
    addr = detGenKeyAddress (encodeUtf8 $ unsafeGetAliasText alias)
    newContractState = dumbImplicitContractState { csAlias =  alias }
  modify $ \s ->
    s & msContractsL . at addr ?~ newContractState
  pure addr

handleGetAlias :: Monad m => AddressOrAlias -> TestT m Alias
handleGetAlias = \case
  AddressAlias alias -> pure alias
  AddressResolved addr -> do
    MockState{..} <- get
    case lookup addr msContracts of
      Nothing                -> throwM $ UnknownContract $ AddressResolved addr
      Just ContractState{..} -> pure $ csAlias

handleGetManagerKey :: (Monad m) => BlockId -> Address -> TestT m (Maybe PublicKey)
handleGetManagerKey blk addr = do
  assertHeadBlockId blk
  s <- get
  let mbCs = s ^. msContractsL . at addr
  case mbCs of
    Just ContractState{..} -> case csContractData of
      ImplicitContractData mbManagerKey -> pure mbManagerKey
      ContractData _ _ -> throwString "Only implicit account can have a manager key"
    Nothing -> throwM $ UnknownContract $ AddressResolved addr

-- In scenarios where the system under test checks for 404 errors, we
-- use this function to mock and simulate those errors.
err404 :: Text -> ClientError
err404 path = FailureResponse
  (defaultRequest { requestBody = Nothing, requestPath = (baseUrl , "") })
  response
  where
    baseUrl = BaseUrl
      { baseUrlScheme = Http
      , baseUrlHost = "localhost"
      , baseUrlPort = 8732
      , baseUrlPath = toString path
      }
    response = Response
      { responseStatusCode = status404
      , responseHeaders = mempty
      , responseHttpVersion = http20
      , responseBody = "Contract with given address not found"
      }

handleResolveAddressMaybe :: Monad m => AddressOrAlias -> TestT m (Maybe Address)
handleResolveAddressMaybe = \case
  AddressResolved addr -> pure (Just addr)
  AddressAlias alias -> do
    MockState{..} <- get
    case find checkAlias $ Map.toList msContracts of
      Just (addr, _) -> pure (Just addr)
      Nothing -> pure Nothing
    where
      checkAlias (_, ContractState { csAlias = alias' }) = alias == alias'

handleRevealKey :: Monad m => Alias -> Maybe ScrubbedBytes -> TestT m ()
handleRevealKey alias _ = do
  contracts <- gets (Map.toList . msContracts)
  let contracts' = filter (\(_, ContractState{..}) -> csAlias == alias) contracts
  case contracts' of
    []  -> throwM $ UnknownContract $ AddressAlias alias
    [(addr, cs@ContractState{..})] ->
      case (addr, csContractData) of
        (ContractAddress _, ContractData _ _) ->
          throwM $ CantRevealContract addr
        (KeyAddress _, ImplicitContractData (Just _)) ->
          throwM $ AlreadyRevealed addr
        (KeyAddress _, ImplicitContractData Nothing) ->
            -- We don't care about the public key itself, but only its presence.
            let newContractState = cs { csContractData = ImplicitContractData $ Just dumbManagerKey }
            in modify $ \s -> s & msContractsL . at addr ?~ newContractState
        _ -> error "Inconsitent mock state. This most likely a bug in tests."
    _   ->
      error $ "Multiple contracts have alias '" +| alias |+
        "'. This is most likely a bug in tests."

-- | Dummy public key used in mock tests.
dumbManagerKey :: PublicKey
dumbManagerKey = fromRight (error "impossible") $ parsePublicKey
  "edpkuwTWKgQNnhR5v17H2DYHbfcxYepARyrPGbf1tbMoGQAj8Ljr3V"

----------------------------------------------------------------------------
-- TemplateHaskell test helpers
----------------------------------------------------------------------------

shouldCompileTo :: HasCallStack => [Dec] -> Q [Dec] -> Assertion
shouldCompileTo actualDecs expectedQ = do
  expectedDecs <- runQ expectedQ
  PrettyDecs (normalizeDecs actualDecs) @?= PrettyDecs (normalizeDecs expectedDecs)

-- | Same as 'shouldCompileTo', but ignores instance declarations of the given class.
shouldCompileIgnoringInstance :: HasCallStack => Name -> [Dec] -> Q [Dec] -> Assertion
shouldCompileIgnoringInstance className actualDecs expectedQ = do
  expectedDecs <- runQ expectedQ
  let actualDecs' = filter (not . isInstance) actualDecs
  PrettyDecs (normalizeDecs actualDecs') @?= PrettyDecs (normalizeDecs expectedDecs)

  where
    isInstance :: Dec -> Bool
    isInstance = \case
      InstanceD _ _ (ConT t `AppT` _) _ | t == className -> True
      _ -> False

-- | Normalize ASTs to make them comparable.
--
-- By default, quoted ASTs and ASTs with names created using 'newName' will have
-- names with unique IDs.
-- For example:
--
-- > decs <- runQ [d|data D = D { f :: Int } |]
-- > putStrLn $ pprint decs
-- >
-- > -- Will generate this AST:
-- > data D_0 = D_1 { f_2 :: Int }
--
-- To be able to check if two ASTs are equivalent, we have to scrub the unique IDs off all names.
--
-- For convenience, to make the output easier to read, we also erase kind annotations when the kind is '*'.
normalizeDecs :: [Dec] -> [Dec]
normalizeDecs decs =
    SYB.everywhere
      (SYB.mkT fixName . SYB.mkT simplifyType . SYB.mkT simplifyTyVar)
      decs
  where
    fixName :: Name -> Name
    fixName = mkName . nameBase

    simplifyType :: Type -> Type
    simplifyType = \case
      SigT t StarT -> t
      t -> t

    simplifyTyVar :: TyVarBndr -> TyVarBndr
    simplifyTyVar = \case
     KindedTV name StarT -> PlainTV name
     tv -> tv

newtype PrettyDecs = PrettyDecs [Dec]
  deriving newtype Eq

instance Show PrettyDecs where
  show (PrettyDecs decs) = pprint decs