module Test.Stress.Client
  ( client
  ) where

import Control.Concurrent
import Control.Concurrent.Async
import Control.Exception
import Control.Monad
import Data.ByteString.Lazy.Char8 qualified as BS.Char8
import GHC.IO.Exception
import Network.HTTP2.Client (HTTP2Error(..))
import Network.Socket
import Network.TLS (TLSException(..))

import Network.GRPC.Client hiding (Call)
import Network.GRPC.Common
import Network.GRPC.Common.Compression (Compression)
import Network.GRPC.Common.Compression qualified as Compr
import Proto.API.Trivial

import Test.Stress.Cmdline
import Test.Stress.Common

-------------------------------------------------------------------------------
-- Top-level
-------------------------------------------------------------------------------

client ::
     Bool
  -> Maybe ServerValidation
  -> PortNumber
  -> Maybe Compression
  -> [Connect]
  -> IO ()
client v mServerValidation serverPort compr =
    mapM_ $ runConnect v mServerValidation serverPort compr

runConnect ::
     Bool
  -> Maybe ServerValidation
  -> PortNumber
  -> Maybe Compression
  -> Connect
  -> IO ()
runConnect v mServerValidation serverPort compr Connect{..} = do
    say' v serverPort $
      "running calls " ++ show connectCalls ++
        case connectExec of
          Sequential -> " in sequence"
          Concurrent -> " concurrently"
    mapF (runCalls v mServerValidation serverPort compr callNum) $
      zip [1..] $ replicate connectNum connectCalls
  where
    mapF =
      case connectExec of
        Sequential -> mapM_
        Concurrent -> mapConcurrently_

runCalls ::
     Bool
  -> Maybe ServerValidation
  -> PortNumber
  -> Maybe Compression
  -> Int
  -> (Int, [Call])
  -> IO ()
runCalls v mServerValidation serverPort compr callNum (connNum, calls) = do
    say' v serverPort msg
    let connParams = def {
            connCompression = maybe Compr.none Compr.insist compr
          , connHTTP2Settings = def {
                http2TcpAbortiveClose = True
              }
          , connReconnectPolicy =
              exponentialBackoff
                (\d -> do
                  say' v serverPort $ "Reconnecting after " ++ show d ++ "μs"
                  threadDelay d
                )
                1
                (0.1, 0.1)
                maxBound

          }
    allowCertainFailures $
      withConnection connParams server $ \conn ->
        replicateM_ callNum $ mapM_ (runCall v serverPort conn) calls
  where
    addr :: Address
    addr = Address {
          addressHost      = "127.0.0.1"
        , addressPort      = serverPort
        , addressAuthority = Nothing
        }

    server :: Server
    server =
        case mServerValidation of
          Just serverValidation ->
            ServerSecure serverValidation SslKeyLogNone addr
          Nothing ->
            ServerInsecure addr

    allowCertainFailures :: IO () -> IO ()
    allowCertainFailures =
      handle $ \case
        e | Just ServerDisconnected{} <- fromException e ->
          say' v serverPort $ "server disconnected: " ++ show e
          | Just IOError{} <- fromException e ->
          say' v serverPort "failed to connect"
          | Just ConnectionIsTimeout <- fromException e ->
          say' v serverPort "got ConnectionIsTimeout"
          | Just BadThingHappen{} <- fromException e ->
          say' v serverPort "got BadThingHappen"
          | Just HandshakeFailed{} <- fromException e ->
          say' v serverPort "got HandshakeFailed"
          | Just BlockedIndefinitelyOnSTM{} <- fromException e ->
          say' v serverPort "got BlockedIndefinitelyOnSTM"
          | otherwise -> do
          say' v serverPort $ "got exception: " ++ displayException e
          throwIO e

    msg :: String
    msg =
           "opening connection " ++ show connNum ++ " to "
        ++ case mServerValidation of
             Just _ -> "secure "
             Nothing -> "insecure "
        ++ "server at port " ++ show serverPort ++ " with compression "
        ++ maybe "none" (show . Compr.compressionId) compr

runCall :: Bool -> PortNumber -> Connection -> Call -> IO ()
runCall v p conn =
    \case
      NonStreaming ->
        nonStreaming v p conn
      ClientStreaming n ->
        clientStreaming v p conn n
      ServerStreaming n ->
        serverStreaming v p conn n
      BiDiStreaming n ->
        bidiStreaming v p conn n

-------------------------------------------------------------------------------
-- Specific RPCs
-------------------------------------------------------------------------------

-- | One non-streaming, round-trip call
nonStreaming :: Bool -> PortNumber -> Connection -> IO ()
nonStreaming v p conn = do
    say' v p "initiating non-streaming call"
    withRPC conn def (Proxy @(Trivial' "non-streaming")) $ \call -> do
      sendFinalInput call =<< randomMsg
      void $ recvFinalOutput call
      say' v p "received final output for non-streaming call"

-- | Client streaming
--
-- Client sends the server @N@, followed by @N@ messages.
clientStreaming :: Bool -> PortNumber -> Connection -> Int -> IO ()
clientStreaming v p conn n = do
    say' v p "initiating client streaming call"
    withRPC conn def (Proxy @(Trivial' "client-streaming")) $ \call -> do
      say' v p $ "sending " ++ show n ++ " messages"
      sendNextInput call $ BS.Char8.pack (show n)
      msg <- randomMsg
      forM_ [1 .. n-1] $ \_ ->
        void $ sendNextInput call msg
      sendFinalInput call msg
      void $ recvFinalOutput call
      say' v p "received final output for client streaming call"

-- | Server streaming
--
-- Client sends the server @N@, then receives @N@ messages from server.
serverStreaming :: Bool -> PortNumber -> Connection -> Int -> IO ()
serverStreaming v p conn n = do
    say' v p "initiating server streaming call"
    withRPC conn def (Proxy @(Trivial' "server-streaming")) $ \call -> do
      say' v p $ "receiving " ++ show n ++ " messages"
      sendFinalInput call $ BS.Char8.pack (show n)
      forM_ [1 .. n-1] $ \_ -> void $ recvNextOutput call
      void $ recvFinalOutput call
      say' v p "received final output for server streaming call"

-- | Bidirectional streaming
--
-- Client sends the server @N@, then alternates sending and receiving @N*2@
-- total messages.
bidiStreaming :: Bool -> PortNumber -> Connection -> Int -> IO ()
bidiStreaming v p conn n = do
    say' v p "initiating bidi streaming call"
    withRPC conn def (Proxy @(Trivial' "bidi-streaming")) $ \call -> do
      say' v p $ "sending and receiving " ++ show n ++ " messages"
      sendNextInput call $ BS.Char8.pack (show n)
      msg <- randomMsg
      forM_ [1 .. n-1] $ \_ -> do
        sendNextInput call msg
        void $ recvNextOutput call
      sendFinalInput call msg
      void $ recvFinalOutput call
      say' v p "sent and received final messages for bidi streaming call"

-------------------------------------------------------------------------------
-- Utils
-------------------------------------------------------------------------------

say' :: Bool -> PortNumber -> String -> IO ()
say' v p msg =
  say v $ "(client " ++ show p ++ ") " ++ msg