{-# LANGUAGE OverloadedLabels  #-}
{-# LANGUAGE OverloadedStrings #-}

module Test.Sanity.Compression (tests) where

import Control.Monad
import Data.IORef
import Data.Maybe (isJust)
import Data.Text (Text)
import Test.Tasty
import Test.Tasty.HUnit

import Network.GRPC.Client qualified as Client
import Network.GRPC.Common
import Network.GRPC.Common.Protobuf
import Network.GRPC.Common.StreamElem qualified as StreamElem
import Network.GRPC.Server qualified as Server

import Test.Driver.ClientServer

import Proto.API.Helloworld

{-------------------------------------------------------------------------------
  Top-level
-------------------------------------------------------------------------------}

tests :: TestTree
tests = testGroup "Test.Sanity.Compression" [
      testCase "multipleRPC"  test_multipleRPC
    , testCase "multipleMsgs" test_multipleMsgs
    ]

{-------------------------------------------------------------------------------
  Individual tests
-------------------------------------------------------------------------------}

-- | Test that compression is enabled for the /second/ RPC
test_multipleRPC :: Assertion
test_multipleRPC = do
    counter <- newIORef 0
    testClientServer $ ClientServerTest {
        config = def
      , server = [Server.someRpcHandler $ handleNonStreaming counter]
      , client = simpleTestClient $ \conn ->
            replicateM_ 2 $ do
              Client.withRPC conn def (Proxy @SayHello) $ \call -> do
                Client.sendFinalInput call req
                mResp <- StreamElem.value <$> Client.recvOutputWithMeta call
                case mResp of
                  Nothing -> assertFailure "Expected response"
                  Just (meta, resp) -> do
                    -- /All/ responses from the server should be compressed
                    -- (the request tells the server what the client supports)
                    assertEqual "" True $
                      isJust (inboundCompressedSize meta)
                    assertEqual "" compressibleName $
                      resp ^. #message
    }
  where
    req :: Proto HelloRequest
    req = defMessage & #name .~ compressibleName

-- | Test that multiple messages on /one/ RPC will either all be compressed or
-- all uncompressed.
test_multipleMsgs :: Assertion
test_multipleMsgs = do
    counter <- newIORef 0
    testClientServer $ ClientServerTest {
        config = def
      , server = [Server.someRpcHandler $ handleBidiStreaming counter]
      , client = simpleTestClient $ \conn ->
          replicateM_ 2 $
            Client.withRPC conn def (Proxy @SayHelloBidiStream) $ \call -> do
              replicateM_ 2 $ do
                Client.sendNextInput call req
                mResp <- StreamElem.value <$> Client.recvOutputWithMeta call
                case mResp of
                  Nothing -> assertFailure "Expected response"
                  Just (meta, resp) -> do
                    -- /All/ responses from the server should be compressed
                    -- (the request tells the server what the client supports)
                    assertEqual "" True $
                      isJust (inboundCompressedSize meta)
                    assertEqual "" compressibleName $
                      resp ^. #message
              Client.sendEndOfInput call
      }
  where
    req :: Proto HelloRequest
    req = defMessage & #name .~ compressibleName

{-------------------------------------------------------------------------------
  Server handlers
-------------------------------------------------------------------------------}

handleNonStreaming :: IORef Int -> Server.RpcHandler IO SayHello
handleNonStreaming counter = Server.mkRpcHandler $ \call -> do
    mElem <- Server.recvInputWithMeta call
    case mElem of
      FinalElem (meta, req) NoMetadata -> do
        callNo <- atomicModifyIORef counter $ \i -> (succ i, i)

        -- We expect all messages to be compressed except the first (the client
        -- does not yet know which compression algorithms the server supports)
        let expectCompression :: Bool
            expectCompression = callNo > 0
        assertEqual "" expectCompression $
          isJust (inboundCompressedSize meta)

        Server.sendFinalOutput call (
            defMessage & #message .~ (req ^. #name)
          , NoMetadata
          )
      _otherwise ->
        assertFailure "expected FinalElem"

handleBidiStreaming :: IORef Int -> Server.RpcHandler IO SayHelloBidiStream
handleBidiStreaming counter = Server.mkRpcHandler $ \call -> do
    isFirstCall <- atomicModifyIORef counter $ \i -> (succ i, i == 0)

    let loop :: IO ()
        loop = do
          mElem <- Server.recvInputWithMeta call
          case mElem of
            NoMoreElems NoMetadata ->
              Server.sendTrailers call NoMetadata
            StreamElem (meta, req) -> do
              -- The compression algorithm is established once at the start of
              -- the request; we cannot start compression halfway a conversation
              let expectCompression :: Bool
                  expectCompression = not isFirstCall
              assertEqual "" expectCompression $
                isJust (inboundCompressedSize meta)

              Server.sendNextOutput call $
                defMessage & #message .~ (req ^. #name)
              loop
            FinalElem{} ->
              assertFailure "Unexpected FinalElem"

    loop

{-------------------------------------------------------------------------------
  Auxiliary
-------------------------------------------------------------------------------}

compressibleName :: Text
compressibleName = mconcat (replicate 100 "John")