-- 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
  , dumbImplicitState
  , 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 (pretty, (+|), (|+))
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.TezosClient (AliasBehavior(..), TezosClientError(DuplicateAlias))
import Morley.Client.Types
import Morley.Micheline
import Morley.Michelson.Typed
import Morley.Tezos.Address
import Morley.Tezos.Address.Alias
import Morley.Tezos.Address.Kinds
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 :: AccountState 'AddressKindContract
dumbContractState = AccountState
  { asCounter = 100500
  , asAlias = "genesis2"
  , asAccountData = ContractData
      OriginationScript
        { osCode    = toExpression $ compileLorentz L.drop
        , osStorage = toExpression $ toVal ()
        }
      Nothing
  }

dumbImplicitState :: AccountState 'AddressKindImplicit
dumbImplicitState = AccountState
  { asCounter = 100500
  , asAlias = "genesis1"
  , asAccountData = ImplicitData 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
  , hGetAliasesAndAddresses = handleGetAliasesAndAddresses
  , hRememberContract = handleRememberContract
  , hGetKeyPassword = \_ -> pure Nothing
  , hGenKey = handleGenKey
  , hGetManagerKey = handleGetManagerKey
  , hRevealKey = handleRevealKey
  }
  where
    testSecretKey :: Ed25519.SecretKey
    testSecretKey = Ed25519.detSecretKey "\001\002\003\004"

mkRunOperationResult :: [ContractAddress] -> 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 -> ImplicitAddress -> m TezosInt64
handleGetCounter blk addr = do
  assertHeadBlockId blk
  FakeState{..} <- get
  case lookup addr fsImplicits of
    Nothing -> throwM $ UnknownAccount $ Constrained addr
    Just AccountState{..} -> pure $ asCounter

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 [ContractAddress]
handleRunOperationInternal RunOperationInternal{..} = do
  concatMapM handleTransactionOrOrigination roiContents

handleTransactionOrOrigination
  :: (Monad m, HasCallStack) => OperationInput -> TestT m [ContractAddress]
handleTransactionOrOrigination op = do
  FakeState{..} <- get
  case wcoCustom op of
    -- Ensure that transaction sender exists
    OpTransfer TransactionOperation{..} -> case lookup codSource fsImplicits of
      Nothing -> throwM $ UnknownAccount $ Constrained codSource
      Just AccountState{..} -> do
        -- Ensure that sender counter matches
        unless (asCounter + 1 == codCounter) (throwM CounterMismatch)
        case toDestination of
          MkAddress dest@ContractAddress{} ->
            case lookup dest fsContracts of
              Nothing -> throwM $ UnknownAccount $ Constrained dest
              Just _ -> pure []
          MkAddress dest@ImplicitAddress{} ->
            case lookup dest fsImplicits of
              Nothing -> throwM $ UnknownAccount $ Constrained dest
              Just _ -> pure []
          MkAddress TxRollupAddress{} -> error "tx rollup unsupported"
    -- Ensure that originator exists
    OpOriginate _ -> case lookup codSource fsImplicits of
      Nothing -> throwM $ UnknownAccount $ Constrained codSource
      Just AccountState{..} -> do
        -- Ensure that originator counter matches
        unless (asCounter + 1 == codCounter) (throwM CounterMismatch)
        pure [dummyContractAddr]
      where
        dummyContractAddr = [ta|KT1LZwEZqbqtLYhwzaidBp6So9LgYDpkpEv7|]
    OpReveal _ ->
      -- We do not care about reveals at the moment
      return []
    OpDelegation _ ->
      -- We do not care about delegations 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
  -> ContractAddress
  -> m OriginationScript
handleGetContractScript blockId addr = do
  assertHeadBlockId blockId
  FakeState{..} <- get
  case lookup addr fsContracts of
    Nothing -> throwM $ err404 path
    Just AccountState{..} -> case asAccountData of
      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 (asAccountData cs) of
            ContractData _ bigMapMaybe -> bigMapMaybe

  -- 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 => AliasBehavior -> ContractAddress -> ContractAlias -> TestT m ()
handleRememberContract DontSaveAlias _ _ = pass
handleRememberContract replaceExisting addr alias = do
  let
    cs = dumbContractState { asAlias = 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
    _       -> case replaceExisting of
      KeepDuplicateAlias -> pass
      OverwriteDuplicateAlias -> remember addr cs st
      ForbidDuplicateAlias -> throwM $ DuplicateAlias $ unAlias alias

handleGenKey :: Monad m => ImplicitAlias -> TestT m ImplicitAddress
handleGenKey alias = do
  let
    addr = detGenKeyAddress (encodeUtf8 $ unAlias alias)
    newAccountState = dumbImplicitState { asAlias =  alias }
  modify $ \s ->
    s & fsImplicitsL . at addr ?~ newAccountState
  pure addr

handleGetAliasesAndAddresses
  :: forall m. Monad m
  => TestT m [(Text, Text)]
handleGetAliasesAndAddresses = do
  FakeState{fsContracts, fsImplicits} <- get
  pure $ convert fsContracts <> convert fsImplicits
  where
    convert :: Map (KindedAddress kind) (AccountState kind) -> [(Text, Text)]
    convert m =
      Map.toList m <&> \(addr, AccountState{asAlias}) -> (pretty asAlias, pretty addr)

handleGetManagerKey :: (Monad m) => BlockId -> ImplicitAddress -> TestT m (Maybe PublicKey)
handleGetManagerKey blk addr = do
  assertHeadBlockId blk
  s <- get
  let mbCs = s ^. fsImplicitsL . at addr
  case mbCs of
    Just AccountState{..} -> case asAccountData of
      ImplicitData mbManagerKey -> pure mbManagerKey
    Nothing -> throwM $ UnknownAccount $ Constrained 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"
      }

handleRevealKey :: Monad m => ImplicitAlias -> Maybe ScrubbedBytes -> TestT m ()
handleRevealKey alias _ = do
  accounts <- gets (Map.toList . fsImplicits)
  let accounts' = filter (\(_, AccountState{..}) -> asAlias == alias) accounts
  case accounts' of
    []  -> throwM $ UnknownAlias alias
    [(addr, cs@AccountState{..})] ->
      case asAccountData of
        ImplicitData (Just _) ->
          throwM $ AlreadyRevealed $ MkAddress addr
        ImplicitData Nothing ->
            -- We don't care about the public key itself, but only its presence.
            let newAccountState = cs { asAccountData = ImplicitData $ Just dumbManagerKey }
            in modify $ \s -> s & fsImplicitsL . at addr ?~ newAccountState
    _   ->
      error $ "Multiple accounts 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"