-- 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
  , mapToContractStateBigMap
  , handleGetBigMapValue
  , genesisState
  , revealKeyState
  , addrAndAliasFromGenesisState

    -- * Internals
  , handleRunOperationInternal
  , assertHeadBlockId
  ) where

import Control.Exception.Safe (throwString)
import Control.Lens (at, (?=), (?~))
import Data.Aeson (encode)
import Data.ByteString.Lazy qualified as LBS (toStrict)
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 (toAddress)
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.Client.Types.AliasesAndAddresses
import Morley.Micheline
import Morley.Michelson.Runtime.GState
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 Morley.Util.Constrained
import Morley.Util.Peano
import Morley.Util.SizedList qualified as Sized
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) $ toPairs 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 :: PublicKey -> (ImplicitAddress, AccountState 'AddressKindImplicit)
dumbImplicitState pk = (mkKeyAddress pk, AccountState
  { asCounter = 100500
  , asAlias = "genesis1"
  , asAccountData = ImplicitData pk 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
  , hGetPublicKey = handleGetPublicKey
  }
  where
    testSecretKey :: Ed25519.SecretKey
    testSecretKey = Ed25519.detSecretKey "\001\002\003\004"

mkRunOperationResult :: NonEmpty (OperationResp WithSource, AppliedResult) -> RunOperationResult
mkRunOperationResult results = RunOperationResult
  { rrOperationContents = results <&> \(op, ar) ->
      OperationContent op $ RunMetadata
        { rmOperationResult = OperationApplied ar
        , rmInternalOperationResults = []
        }
  }

mkAppliedResult :: [ContractAddress] -> AppliedResult
mkAppliedResult originatedContracts = AppliedResult 100500 100500 100500 originatedContracts 0

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
  use (fsImplicitsL . at addr) >>= \case
    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
  mkRunOperationResult <$> handleRunOperationInternal roOperation

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
  results <- handleOperationInput paoContents
  pure $ mkRunOperationResult results

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 (NonEmpty (OperationResp WithSource, AppliedResult))
handleRunOperationInternal RunOperationInternal{roiContents} = handleOperationInput roiContents

handleOperationInput
  :: Monad m
  => NonEmpty OperationInput
  -> TestT m (NonEmpty (OperationResp WithSource, AppliedResult))
handleOperationInput (c :| cs) =
  (:|) <$> go c 1 <*> zipWithM go cs [2, 3..]
  where go op = fmap ((toIntOpData op, ) . mkAppliedResult) . handleTransactionOrOrigination op

toIntOpData :: OperationInput -> OperationResp WithSource
toIntOpData (WithCommonOperationData CommonOperationData{..} op) = case op of
  OpTransfer op' -> ws TransactionOpResp op'
  OpTransferTicket op' -> ws TransferTicketOpResp op'
  OpOriginate op' -> ws OriginationOpResp op'
  OpReveal op' -> ws RevealOpResp op'
    -- reveal operation apparently can't be internal, but why not handle this anyway
  OpDelegation op' -> ws DelegationOpResp op'
  where
    ws :: (WithSource a -> r) -> a -> r
    ws f = f . WithSource (toAddress codSource)

handleTransactionOrOrigination
  :: forall m. (Monad m, HasCallStack) => OperationInput -> TezosInt64 -> TestT m [ContractAddress]
handleTransactionOrOrigination op n = do
  FakeState{..} <- get
  AccountState{..} <-
    maybe (throwM $ UnknownAccount $ Constrained codSource) pure $ fsImplicits ^. at codSource
  unless ((asCounter + n) == codCounter) $ throwM CounterMismatch
  let checkContractExists :: Address -> TestT m ()
      checkContractExists dest'@(MkAddress dest) = case dest of
        ContractAddress{} -> when (isNothing $ fsContracts ^. at dest) $
          throwM $ UnknownAccount dest'
        ImplicitAddress{} -> when (isNothing $ fsImplicits ^. at dest) $
          throwM $ UnknownAccount dest'
        SmartRollupAddress{} -> error "smart rollup unsupported"
  case wcoCustom op of
    OpTransfer TransactionOperation{toDestination} -> do
      checkContractExists toDestination
      pure []
    OpTransferTicket TransferTicketOperation{ttoDestination} -> do
      checkContractExists ttoDestination
      pure []
    OpOriginate{} ->
      pure [[ta|KT1LZwEZqbqtLYhwzaidBp6So9LgYDpkpEv7|]]
    OpReveal RevealOperation{..} -> do
      let addr = mkKeyAddress roPublicKey
      modify $ fsImplicitsL . at addr %~ fmap revealKeyState
      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
  use (fsContractsL . at addr) >>= \case
    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 $
          toList (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 csbmMap bigMap ^. at scriptExpr 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' = fsContractsL . at addr' ?= cs'

  use (fsContractsL . at addr) >>= \case
    Nothing -> remember addr cs
    _       -> case replaceExisting of
      KeepDuplicateAlias -> pass
      OverwriteDuplicateAlias -> remember addr cs
      ForbidDuplicateAlias -> throwM $ DuplicateAlias $ unAlias alias

handleGenKey :: Monad m => ImplicitAlias -> TestT m ImplicitAddressWithAlias
handleGenKey alias = do
  let
    pk = toPublic . detSecretKey . encodeUtf8 $ unAlias alias
    (addr, newAccountState) = dumbImplicitState pk
  modify $ fsImplicitsL . at addr ?~ newAccountState{ asAlias = alias }
  pure $ AddressWithAlias addr alias

handleGetAliasesAndAddresses
  :: forall m. Monad m
  => TestT m AliasesAndAddresses
handleGetAliasesAndAddresses = do
  FakeState{fsContracts, fsImplicits} <- get
  pure $ mkAliasesAndAddresses $ convert fsContracts <> convert fsImplicits
  where
    convert
      :: Map (KindedAddress kind) (AccountState kind)
      -> [Constrained NullConstraint AddressWithAlias]
    convert m = toPairs m <&> \(awaAddress, AccountState{asAlias=awaAlias}) ->
      Constrained AddressWithAlias{..}

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{..} -> pure $ idManagerKey asAccountData
    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"
      }

handleGetPublicKey :: Monad m => ImplicitAddressWithAlias -> TestT m PublicKey
handleGetPublicKey aaa@(AddressWithAlias addr alias) = do
  accounts <- gets (toPairs . fsImplicits)
  let accounts' = filter check accounts
      check (addr', AccountState{..}) =  addr' == addr && asAlias == alias
  case accounts' of
    []  -> throwM $ UnknownAccount $ Constrained addr
    [(_, AccountState{..})] -> pure $ idPublicKey asAccountData
    _   ->
      error $ "Multiple accounts found for '" +| aaa |+
        "'. This is most likely a bug in tests."

genesisPk :: forall n. (SingIPeano n, ToPeano 10 > ToPeano n ~ 'True) => PublicKey
genesisPk = toPublic $ Sized.index @n genesisSecrets

genesisState
  :: forall n. (SingIPeano n, ToPeano 10 > ToPeano n ~ 'True)
  => (ImplicitAddress, AccountState 'AddressKindImplicit)
genesisState = dumbImplicitState $ genesisPk @n

addrAndAliasFromGenesisState
  :: forall n. (SingIPeano n, ToPeano 10 > ToPeano n ~ 'True)
  => ImplicitAddressWithAlias
addrAndAliasFromGenesisState =
  let (addr, AccountState{asAlias}) = genesisState @n
  in AddressWithAlias addr asAlias

revealKeyState :: AccountState 'AddressKindImplicit -> AccountState 'AddressKindImplicit
revealKeyState = asAccountDataL %~ \x@ImplicitData{..} -> x{ idManagerKey = Just idPublicKey }