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

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

    -- * Internals
  , handleRunOperationInternal
  , assertHeadBlockId
  ) where

import Control.Exception.Safe (throwString)
import Control.Lens (at, (?~))
import Data.Aeson (encode)
import Data.ByteArray (ScrubbedBytes)
import Data.ByteString.Lazy qualified as LBS (toStrict)
import Data.Map as Map (elems, insert, lookup, toList)
import Data.Singletons (demote)
import Fmt ((+|), (|+))
import Network.HTTP.Types.Status (status404)
import Network.HTTP.Types.Version (http20)
import Servant.Client.Core
  (BaseUrl(..), ClientError(..), RequestF(..), ResponseF(..), Scheme(..), defaultRequest)
import Text.Hex (encodeHex)

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

-- | Function to convert given map to big map representation
-- used in fake 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 fake 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
  }

-- | Fake 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 = handleForgeOperation
  , hInjectOperation = pure . OperationHash . (<> "_injected") . encodeHex . unHexJSONByteString
  , hGetContractScript = handleGetContractScript
  , hSignBytes =
    \_ _ -> pure . SignatureEd25519 . Ed25519.sign testSecretKey
  , hWaitForOperation = id
  , hGetAlias = handleGetAlias
  , hResolveAddressMaybe = handleResolveAddressMaybe
  , hRememberContract = handleRememberContract
  , hCalcTransferFee = \_ _ _ _ -> pure $ [TezosMutez [tz|100500u|]]
  , hCalcOriginationFee = \_ -> pure $ TezosMutez [tz|100500u|]
  , 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 BlockHash
handleGetBlockHash blkId = do
  unless (blkId == FinalHeadId) do
    throwString "Expected `getBlockHash` to be called with `head~2`."

  FakeState{..} <- get
  pure fsFinalHeadBlock

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

handleGetBlockConstants
  :: MonadState FakeState m
  => BlockId -> m BlockConstants
handleGetBlockConstants blkId = do
  FakeState{..} <- get
  pure $ fsBlockConstants blkId

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

handleRunOperation :: Monad m => BlockId -> RunOperation -> TestT m RunOperationResult
handleRunOperation blk RunOperation{..} = do
  assertHeadBlockId blk
  FakeState{..} <- get
  -- Ensure that passed chain id matches with one that fake state has
  unless (roChainId == bcChainId (fsBlockConstants blk)) (throwM $ InvalidChainId)
  -- As of release of the ithaca protocol, the "branch" field should be "head~2".
  -- https://web.archive.org/web/20220305165609/https://tezos.gitlab.io/protocols/tenderbake.html
  unless (roiBranch roOperation == fsFinalHeadBlock) do
    throwM $ InvalidBranch $ roiBranch roOperation
  originatedContracts <- handleRunOperationInternal roOperation
  pure $ mkRunOperationResult originatedContracts

handlePreApplyOperation :: Monad m => BlockId -> PreApplyOperation -> TestT m RunOperationResult
handlePreApplyOperation blk PreApplyOperation{..} = do
  assertHeadBlockId blk
  FakeState{..} <- get
  -- Ensure that passed protocol matches with one that mock state has
  unless (paoProtocol == bcProtocol (fsBlockConstants blk)) $
    throwM InvalidProtocol
  -- As of release of the ithaca protocol, the "branch" field should be "head~2".
  -- https://web.archive.org/web/20220305165609/https://tezos.gitlab.io/protocols/tenderbake.html
  unless (paoBranch == fsFinalHeadBlock) do
    throwM $ InvalidBranch paoBranch
  originatedContracts <- concatMapM handleTransactionOrOrigination paoContents
  pure $ mkRunOperationResult originatedContracts

handleForgeOperation :: Monad m => BlockId -> ForgeOperation -> TestT m HexJSONByteString
handleForgeOperation blkId op = do
  assertHeadBlockId blkId
  ms <- get
  -- As of release of the ithaca protocol, the "branch" field should be "head~2".
  -- https://web.archive.org/web/20220305165609/https://tezos.gitlab.io/protocols/tenderbake.html
  unless (foBranch op == fsFinalHeadBlock ms) do
    throwM $ InvalidBranch $ foBranch op
  pure . HexJSONByteString . LBS.toStrict . encode $ op

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

handleTransactionOrOrigination
  :: Monad m => OperationInput -> TestT m [Address]
handleTransactionOrOrigination op = do
  FakeState{..} <- get
  case wcoCustom op of
    -- Ensure that transaction sender exists
    OpTransfer TransactionOperation{..} -> case lookup codSource fsContracts of
      Nothing -> throwM $ UnknownContract $ AddressResolved codSource
      Just ContractState{..} -> do
        -- Ensure that sender counter matches
        unless (csCounter + 1 == codCounter) (throwM CounterMismatch)
        case lookup toDestination fsContracts of
          Nothing -> throwM $ UnknownContract $ AddressResolved toDestination
          Just _ -> pure []
    -- Ensure that originator exists
    OpOriginate _ -> case lookup codSource fsContracts of
      Nothing -> throwM $ UnknownContract $ AddressResolved codSource
      Just ContractState{..} -> do
        -- Ensure that originator counter matches
        unless (csCounter + 1 == codCounter) (throwM CounterMismatch)
        pure [dummyContractAddr]
      where
        dummyContractAddr = unsafe $ parseAddress "KT1LZwEZqbqtLYhwzaidBp6So9LgYDpkpEv7"
    OpReveal _ ->
      -- We do not care about reveals at the moment
      return []
  where
    CommonOperationData{..} = wcoCommon op

-- | In most places, @morley-client@ executes operations against the @head@ block.
assertHeadBlockId :: (HasCallStack, MonadThrow m) => BlockId -> m ()
assertHeadBlockId blockId = unless (blockId == HeadId) $
  throwString "Accessing non-head block is not supported in tests"

handleGetContractScript
  :: ( MonadState FakeState m
     , MonadThrow m
     )
  => BlockId
  -> Address
  -> m OriginationScript
handleGetContractScript blockId addr = do
  assertHeadBlockId blockId
  FakeState{..} <- get
  case lookup addr fsContracts 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 (fsContracts 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

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

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

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

handleGetAlias :: Monad m => AddressOrAlias -> TestT m Alias
handleGetAlias = \case
  AddressAlias alias -> pure alias
  AddressResolved addr -> do
    FakeState{..} <- get
    case lookup addr fsContracts 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 ^. fsContractsL . 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 fake 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
    FakeState{..} <- get
    case find checkAlias $ Map.toList fsContracts 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 . fsContracts)
  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 & fsContractsL . at addr ?~ newContractState
        _ -> error "Inconsitent fake 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 fake tests.
dumbManagerKey :: PublicKey
dumbManagerKey = fromRight (error "impossible") $ parsePublicKey
  "edpkuwTWKgQNnhR5v17H2DYHbfcxYepARyrPGbf1tbMoGQAj8Ljr3V"