{-# LANGUAGE GeneralizedNewtypeDeriving #-}

-- This file is part of the Wire Server implementation.
--
-- Copyright (C) 2022 Wire Swiss GmbH <opensource@wire.com>
--
-- This program is free software: you can redistribute it and/or modify it under
-- the terms of the GNU Affero General Public License as published by the Free
-- Software Foundation, either version 3 of the License, or (at your option) any
-- later version.
--
-- This program is distributed in the hope that it will be useful, but WITHOUT
-- ANY WARRANTY; without even the implied warranty of MERCHANTABILITY or FITNESS
-- FOR A PARTICULAR PURPOSE. See the GNU Affero General Public License for more
-- details.
--
-- You should have received a copy of the GNU Affero General Public License along
-- with this program. If not, see <https://www.gnu.org/licenses/>.

module Federator.MockServer
  ( -- * Federator mock server
    MockFederator (..),
    MockException (..),
    withTempMockFederator,
    FederatedRequest (..),
    MockResponse (..),

    -- * Mock utilities
    Mock,
    runMock,
    mockReply,
    mockUnreachableFor,
    mockFail,
    guardRPC,
    guardComponent,
    (~>),
    getRequest,
    getRequestRPC,
    getRequestBody,
    mockWithStatus200,
  )
where

import Control.Exception qualified as Exception
import Control.Exception.Base (throw)
import Control.Monad.Catch hiding (fromException)
import Control.Monad.Trans.Except
import Control.Monad.Trans.Maybe
import Data.Aeson qualified as Aeson
import Data.Default
import Data.Domain
import Data.Text qualified as Text
import Data.Text.Lazy qualified as LText
import Federator.Error
import Federator.Error.ServerError
import Federator.InternalServer
import Federator.Interpreter
import Federator.RPC
import Federator.Validation
import Imports hiding (fromException)
import Network.HTTP.Media qualified as HTTP
import Network.HTTP.Types as HTTP
import Network.HTTP.Types qualified as Wai
import Network.Wai qualified as Wai
import Network.Wai.Handler.Warp qualified as Warp
import Network.Wai.Utilities.Error as Wai hiding (Error)
import Network.Wai.Utilities.MockServer
import Polysemy
import Polysemy.Error hiding (throw)
import Servant.API
import Servant.Server qualified as Servant
import Servant.Server.Generic
import System.Logger qualified as Log
import Wire.API.Federation.API (Component)
import Wire.API.Federation.API.Common
import Wire.API.Federation.Domain
import Wire.API.Federation.Version
import Wire.Sem.Logger.TinyLog

-- | This can be thrown by actions passed to mock federator to simulate
-- failures either in federator itself, or in the services it calls.
data MockException
  = MockErrorResponse HTTP.Status LText
  | MockUnreachableBackendErrorResponse Domain
  deriving (Eq, Show, Typeable)

instance AsWai MockException where
  errorLogLevel _ = Log.Error
  toWai (MockErrorResponse status message) = Wai.mkError status "mock-error" message
  toWai (MockUnreachableBackendErrorResponse d) =
    Wai.mkError HTTP.status503 "mock-error" (unreachableMsg d)

unreachableMsg :: Domain -> LText
unreachableMsg (LText.fromStrict . domainText -> d) =
  "unreachable_backend: " <> d

instance Exception MockException

data FederatedRequest = FederatedRequest
  { frOriginDomain :: Domain,
    frTargetDomain :: Domain,
    frComponent :: Component,
    frRPC :: Text,
    frBody :: LByteString
  }
  deriving (Eq, Show)

mockServer ::
  ( Member (Embed IO) r,
    Member (Error MockException) r,
    Member (Error ServerError) r,
    Member (Error ValidationError) r
  ) =>
  IORef [FederatedRequest] ->
  MockFederator ->
  API (AsServerT (Sem r))
mockServer remoteCalls mock =
  Federator.InternalServer.API
    { status = const $ pure NoContent,
      internalRequest = mockInternalRequest remoteCalls mock
    }

mockInternalRequest ::
  ( Member (Embed IO) r,
    Member (Error MockException) r,
    Member (Error ServerError) r,
    Member (Error ValidationError) r
  ) =>
  IORef [FederatedRequest] ->
  MockFederator ->
  Domain ->
  Component ->
  RPC ->
  Wai.Request ->
  (Wai.Response -> IO Wai.ResponseReceived) ->
  Sem r Wai.ResponseReceived
mockInternalRequest remoteCalls mock targetDomain component (RPC path) req cont = do
  domainTxt <- note NoOriginDomain $ lookup originDomainHeaderName (Wai.requestHeaders req)
  originDomain <- parseDomain domainTxt
  reqBody <- embed $ Wai.strictRequestBody req
  let fedRequest =
        ( FederatedRequest
            { frOriginDomain = originDomain,
              frTargetDomain = targetDomain,
              frComponent = component,
              frRPC = path,
              frBody = reqBody
            }
        )
  mr <-
    if path == "api-version"
      then pure $ MockResponse Wai.status200 "application/json" (Aeson.encode (VersionInfo mock.versions))
      else do
        modifyIORef remoteCalls (<> [fedRequest])
        fromException @MockException
          . handle (throw . handleException)
          $ mock.handler fedRequest
  let headers = ("Content-Type", HTTP.renderHeader mr.contentType) : mock.headers
  embed . cont $ Wai.responseLBS mr.status headers mr.body
  where
    handleException :: SomeException -> MockException
    handleException e = case Exception.fromException e of
      Just mockE -> mockE
      Nothing -> MockErrorResponse HTTP.status500 (LText.pack (displayException e))

data MockFederator = MockFederator
  { headers :: [HTTP.Header],
    handler :: FederatedRequest -> IO MockResponse,
    versions :: [Int]
  }

data MockResponse = MockResponse
  { status :: HTTP.Status,
    contentType :: HTTP.MediaType,
    body :: LByteString
  }
  deriving (Eq, Show)

instance Default MockResponse where
  def =
    MockResponse
      { status = HTTP.status200,
        contentType = "application/json",
        body = Aeson.encode EmptyResponse
      }

instance Default MockFederator where
  def =
    MockFederator
      { headers = [],
        handler = \_ -> pure def,
        versions = map versionInt (toList supportedVersions)
      }

mockWithStatus200 :: (HTTP.MediaType, LByteString) -> MockResponse
mockWithStatus200 (ct, b) =
  def
    { status = status200,
      contentType = ct,
      body = b
    }

-- | Spawn a mock federator on a random port and run an action while it is running.
--
-- A mock federator is a web application that parses requests of the same form
-- as a regular outward service, but runs the provided action instead of
-- forwarding them to a remote federator.
withTempMockFederator ::
  (MonadIO m, MonadMask m) =>
  MockFederator ->
  (Warp.Port -> m a) ->
  m (a, [FederatedRequest])
withTempMockFederator mock action = do
  remoteCalls <- newIORef []
  let interpreter =
        Servant.Handler
          . ExceptT
          . runM
          . discardTinyLogs
          . runError
          . runWaiErrors
            @'[ ValidationError,
                ServerError,
                MockException
              ]
      app = genericServeT interpreter (mockServer remoteCalls mock)
  result <-
    bracket
      (liftIO (startMockServer Nothing app))
      (liftIO . fst)
      (\(_close, port) -> action port)
  calls <- readIORef remoteCalls
  pure (result, calls)

--------------------------------------------------------------------------------
-- Mock creation utilities

-- | This is a monad that can be used to create mocked responses. It is a very
-- minimalistic web framework. One can expect a certain request (using
-- 'guardRPC') and reply accordingly (using 'mockReply'). Multiple possible
-- requests and responses can be combined using the 'Alternative' instance.  In
-- simple cases, one can also use the infix '(~>)' combinator, which concisely
-- binds a request to a hardcoded pure response.
newtype Mock a = Mock {unMock :: ReaderT FederatedRequest (MaybeT (ExceptT Text IO)) a}
  deriving newtype (Functor, Applicative, Alternative, Monad, MonadIO)

-- | Convert a mocked response to a function which can be used as input to
-- 'tempMockFederator'.
runMock :: (Text -> IO a) -> Mock a -> FederatedRequest -> IO a
runMock err m req =
  runExceptT (runMaybeT (runReaderT (unMock m) req)) >>= \case
    Right Nothing -> err ("unmocked endpoint called: " <> frRPC req)
    Right (Just x) -> pure x
    Left e -> err e

-- | Retrieve the current request.
getRequest :: Mock FederatedRequest
getRequest = Mock $ ReaderT pure

-- | Retrieve the RPC of the current request.
getRequestRPC :: Mock Text
getRequestRPC = frRPC <$> getRequest

-- | Retrieve and deserialise the body of the current request.
getRequestBody :: (Aeson.FromJSON a) => Mock a
getRequestBody = do
  b <- frBody <$> getRequest
  case Aeson.eitherDecode b of
    Left e -> do
      rpc <- getRequestRPC
      mockFail ("Parse failure in " <> rpc <> ": " <> Text.pack e)
    Right x -> pure x

-- | Expect a given RPC. If the current request does not match, the whole
-- action fails. This can be used in combination with the 'Alternative'
-- instance to provide responses for multiple requests.
guardRPC :: Text -> Mock ()
guardRPC rpc = do
  rpc' <- getRequestRPC
  guard (rpc' == rpc)

guardComponent :: Component -> Mock ()
guardComponent c = do
  c' <- frComponent <$> getRequest
  guard (c == c')

-- | Serialise and return a response.
mockReply :: (Aeson.ToJSON a) => a -> Mock LByteString
mockReply = pure . Aeson.encode

-- | Provide a mock reply simulating an unreachable backend.
mockUnreachableFor :: Set Domain -> Mock LByteString
mockUnreachableFor backends = do
  target <- frTargetDomain <$> getRequest
  guard (target `elem` backends)
  throw (MockUnreachableBackendErrorResponse target)

-- | Abort the mock with an error.
mockFail :: Text -> Mock a
mockFail = Mock . lift . lift . throwE

infixl 5 ~>

-- | Expect a given RPC and simply return a pure response when the current
-- request matches.
(~>) :: (Aeson.ToJSON a) => Text -> a -> Mock LByteString
(~>) rpc x = guardRPC rpc *> mockReply x
