-- 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/>.
{-# OPTIONS_GHC -Wno-orphans #-}

-- This file is part of the Wire Server implementation.
--
-- Copyright (C) 2025 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 Test.Federator.Client (tests) where

import Control.Exception hiding (handle)
import Control.Monad.Codensity
import Control.Monad.Except
import Data.Aeson qualified as Aeson
import Data.Bifunctor (first)
import Data.ByteString qualified as BS
import Data.ByteString.Builder (Builder, byteString, toLazyByteString)
import Data.ByteString.Lazy qualified as LBS
import Data.Default
import Data.Domain
import Data.Id
import Data.Proxy
import Data.Text.Encoding qualified as Text
import Federator.MockServer
import HTTP2.Client.Manager (defaultHttp2Manager, withHTTP2Request)
import Imports
import Network.HTTP.Types as HTTP
import Network.HTTP2.Client qualified as HTTP2
import Network.Wai qualified as Wai
import Network.Wai.Utilities.Error qualified as Wai
import Network.Wai.Utilities.MockServer
import Network.Wai.Utilities.Server
import Servant.API
import Servant.Client hiding ((//))
import Servant.Client.Core
import Servant.Types.SourceT
import Test.QuickCheck (arbitrary, generate)
import Test.Tasty
import Test.Tasty.HUnit
import Util.Options (Endpoint (Endpoint))
import Wire.API.Federation.API
import Wire.API.Federation.Client
import Wire.API.Federation.Error
import Wire.API.User (UserProfile)

targetDomain :: Domain
targetDomain = Domain "target.example.com"

originDomain :: Domain
originDomain = Domain "origin.example.com"

tests :: TestTree
tests =
  testGroup
    "Federator.Client"
    [ testGroup
        "Servant"
        [ testCase "testClientSuccess" testClientSuccess,
          testCase "testClientStreaming" testClientStreaming,
          testCase "testClientFailure" testClientFailure,
          testCase "testFederatorFailure" testFederatorFailure,
          testCase "testClientException" testClientExceptions,
          testCase "testClientConnectionError" testClientConnectionError
        ],
      testGroup
        "HTTP2 client"
        [ testCase "testResponseHeaders" testResponseHeaders,
          testCase "testStreaming" testStreaming
        ]
    ]

newtype ResponseFailure = ResponseFailure Wai.Error
  deriving (Show)

withMockFederatorClient ::
  MockFederator ->
  FederatorClient c a ->
  IO (Either ResponseFailure a, [FederatedRequest])
withMockFederatorClient mock action = withTempMockFederator mock $ \port -> do
  mgr <- defaultHttp2Manager
  let env =
        FederatorClientEnv
          { ceOriginDomain = originDomain,
            ceTargetDomain = targetDomain,
            ceFederator = Endpoint "127.0.0.1" (fromIntegral port),
            ceHttp2Manager = mgr,
            ceOriginRequestId = RequestId defRequestId
          }
  a <- runFederatorClient env action
  case a of
    Left (FederatorClientError r) -> pure (Left (ResponseFailure r))
    Left err -> assertFailure $ "Unexpected client error: " <> displayException err
    Right x -> pure (Right x)

testClientSuccess :: IO ()
testClientSuccess = do
  handle <- generate arbitrary
  expectedResponse :: UserProfile <- generate arbitrary

  (actualResponse, sentRequests) <-
    withMockFederatorClient
      def {handler = const (pure def {body = Aeson.encode (Just expectedResponse)})}
      $ fedClient @'Brig @"get-user-by-handle" handle

  sentRequests
    @?= [ FederatedRequest
            { frTargetDomain = targetDomain,
              frBody = Aeson.encode handle,
              frOriginDomain = originDomain,
              frRPC = "get-user-by-handle",
              frComponent = Brig
            }
        ]
  first (const ()) actualResponse @?= Right (Just expectedResponse)

type StreamingAPI = StreamGet NewlineFraming PlainText (SourceIO Text)

testClientStreaming :: IO ()
testClientStreaming = withInfiniteMockServer $ \port -> do
  mgr <- defaultHttp2Manager
  let env =
        FederatorClientEnv
          { ceOriginDomain = originDomain,
            ceTargetDomain = targetDomain,
            ceFederator = Endpoint "127.0.0.1" (fromIntegral port),
            ceHttp2Manager = mgr,
            ceOriginRequestId = RequestId defRequestId
          }
      venv = FederatorClientVersionedEnv env Nothing
  let c = clientIn (Proxy @StreamingAPI) (Proxy @(FederatorClient 'Brig))
  runCodensity (runExceptT (runVersionedFederatorClientToCodensity venv c)) $ \case
    Left err -> assertFailure $ "Unexpected error: " <> displayException err
    Right out -> do
      let expected = mconcat (replicate 500 "Hello")
      actual <- takeSourceT (fromIntegral (LBS.length expected)) (fmap Text.encodeUtf8 out)
      actual @?= expected

testClientFailure :: IO ()
testClientFailure = do
  handle <- generate arbitrary

  (actualResponse, _) <-
    withMockFederatorClient
      def {handler = const (throw (MockErrorResponse HTTP.status422 "wrong domain"))}
      $ do
        fedClient @'Brig @"get-user-by-handle" handle

  case actualResponse of
    Right _ -> assertFailure "unexpected success"
    Left (ResponseFailure werr) -> do
      Wai.code werr @?= HTTP.status422
      Wai.message werr @?= "wrong domain"

testFederatorFailure :: IO ()
testFederatorFailure = do
  handle <- generate arbitrary

  (actualResponse, _) <-
    withMockFederatorClient
      def {handler = const (throw (MockErrorResponse HTTP.status403 "invalid path"))}
      $ do
        fedClient @'Brig @"get-user-by-handle" handle

  case actualResponse of
    Right _ -> assertFailure "unexpected success"
    Left (ResponseFailure werr) -> do
      Wai.code werr @?= HTTP.status500
      Wai.label werr @?= "federation-local-error"

testClientExceptions :: IO ()
testClientExceptions = do
  handle <- generate arbitrary

  (response, _) <-
    withMockFederatorClient def {handler = const (evaluate (error "unhandled exception"))} $
      fedClient @'Brig @"get-user-by-handle" handle

  case response of
    Right _ -> assertFailure "unexpected success"
    Left (ResponseFailure werr) -> Wai.code werr @?= HTTP.status500

testClientConnectionError :: IO ()
testClientConnectionError = do
  handle <- generate arbitrary
  mgr <- defaultHttp2Manager
  let env =
        FederatorClientEnv
          { ceOriginDomain = originDomain,
            ceTargetDomain = targetDomain,
            ceFederator = Endpoint "127.0.0.1" 1,
            ceHttp2Manager = mgr,
            ceOriginRequestId = RequestId defRequestId
          }
  result <- runFederatorClient env (fedClient @'Brig @"get-user-by-handle" handle)
  case result of
    Left (FederatorClientHTTP2Error (FederatorClientConnectionError _)) -> pure ()
    Left x -> assertFailure $ "Expected connection error, got: " <> displayException x
    Right _ -> assertFailure "Expected connection with the server to fail"

testResponseHeaders :: IO ()
testResponseHeaders = do
  (r, _) <- withTempMockFederator
    def
      { headers = [("X-Foo", "bar")],
        handler = const $ pure def
      }
    $ \port -> do
      let req =
            HTTP2.requestBuilder
              HTTP.methodPost
              "/rpc/target.example.com/brig/test"
              [("Wire-Origin-Domain", "origin.example.com"), (federationRequestIdHeaderName, "rid")]
              "body"
      mgr <- defaultHttp2Manager
      performHTTP2Request mgr (False, "127.0.0.1", port) req
  case r of
    Left err ->
      assertFailure $
        "Unexpected error while connecting to mock federator: " <> show err
    Right resp -> do
      responseStatusCode resp @?= HTTP.status200
      lookup "X-Foo" (toList (responseHeaders resp)) @?= Just "bar"

testStreaming :: IO ()
testStreaming = withInfiniteMockServer $ \port -> do
  let req = HTTP2.requestBuilder HTTP.methodPost "test" [] mempty
  mgr <- defaultHttp2Manager
  withHTTP2Request mgr (False, "127.0.0.1", port) req $ consumeStreamingResponseWith $ \resp -> do
    let expected = mconcat (replicate 512 "Hello\n")
    actual <- takeSourceT (fromIntegral (LBS.length expected)) (responseBody resp)
    actual @?= expected

withInfiniteMockServer :: (Int -> IO a) -> IO a
withInfiniteMockServer k = bracket (startMockServer Nothing app) fst (k . snd)
  where
    app _ respond = respond $
      Wai.responseStream HTTP.ok200 mempty $ \write flush ->
        let go n = do
              when (n == 0) flush
              write (byteString "Hello\n") *> go (if n == 0 then 100 else n - 1)
         in go (1000 :: Int)

-- SourceT utilities

takeStepT :: Builder -> Int -> StepT IO ByteString -> IO LByteString
takeStepT acc _ Stop = pure (toLazyByteString acc)
takeStepT acc _ (Error _) = pure (toLazyByteString acc)
takeStepT acc s (Skip next) = takeStepT acc s next
takeStepT acc s (Yield chunk next)
  | BS.length chunk >= s =
      pure $ toLazyByteString (acc <> byteString (BS.take s chunk))
  | otherwise = do
      takeStepT (acc <> byteString chunk) (s - BS.length chunk) next
takeStepT acc s (Effect m) = m >>= takeStepT acc s

takeSourceT :: Int -> SourceT IO ByteString -> IO LByteString
takeSourceT s m = unSourceT m (takeStepT mempty s)
