{-# LANGUAGE BangPatterns      #-}
{-# LANGUAGE DataKinds         #-}
{-# LANGUAGE DeriveGeneric     #-}
{-# LANGUAGE LambdaCase        #-}
{-# LANGUAGE OverloadedStrings #-}

import           Control.Concurrent
import           Control.Concurrent.Async
import           Control.Exception                          (bracket)
import           Control.Monad
import           Criterion.Main
import           Criterion.Types                            (Config (..))
import qualified Data.ByteString.Lazy                       as BL
import           Data.Word
import           GHC.Generics                               (Generic)
import           Network.GRPC.HighLevel.Server              hiding (serverLoop)
import           Network.GRPC.HighLevel.Server.Unregistered (serverLoop)
import           Network.GRPC.LowLevel
import           Network.GRPC.LowLevel.GRPC                 (threadDelaySecs)
import           Proto3.Suite.Class
import           Proto3.Suite.Types
import           System.Random                              (randomRIO)

data AddRequest = AddRequest {addX   :: Fixed Word32
                              , addY :: Fixed Word32}
  deriving (Show, Eq, Ord, Generic)
instance Message AddRequest

data AddResponse = AddResponse (Fixed Word32)
  deriving (Show, Eq, Ord, Generic)
instance Message AddResponse

addMethod :: MethodName
addMethod = MethodName "/unary"

addClientStreamMethod :: MethodName
addClientStreamMethod = MethodName "/clientstream"

addServerStreamMethod :: MethodName
addServerStreamMethod = MethodName "/serverstream"

addBiDiMethod :: MethodName
addBiDiMethod = MethodName "/bidistream"

addHandler :: Handler 'Normal
addHandler =
  UnaryHandler addMethod $
    \c -> do
      let b = payload c
      return ( AddResponse $ addX b + addY b
             , metadata c
             , StatusOk
             , StatusDetails ""
             )

addClientStreamHandler :: Handler 'ClientStreaming
addClientStreamHandler =
  ClientStreamHandler addClientStreamMethod $
  \_ recv -> do
    answer <- go recv 0
    return (Just answer, mempty, StatusOk, "")
    where go recv !i = do
            req <- recv
            case req of
              Left _ -> return $ AddResponse i
              Right Nothing -> return $ AddResponse i
              Right (Just (AddRequest x y)) -> go recv (i+x+y)

addServerStreamHandler :: Handler 'ServerStreaming
addServerStreamHandler =
  ServerStreamHandler addServerStreamMethod $
  \c send -> do
    let AddRequest (Fixed x) y = payload c
    replicateM_ (fromIntegral x) $ send $ AddResponse y
    return (mempty, StatusOk, "")

addBiDiHandler :: Handler 'BiDiStreaming
addBiDiHandler = BiDiStreamHandler addBiDiMethod (go 0)
  where go :: Fixed Word32 -> ServerRWHandler AddRequest AddResponse
        go !i c recv send = do
          req <- recv
          case req of
            Left _ -> return (mempty, StatusOk, "")
            Right Nothing -> return (mempty, StatusOk, "")
            Right (Just (AddRequest x y)) -> do
              let curr = i + x + y
              void $ send $ AddResponse curr
              go curr c recv send


serverOpts :: ServerOptions
serverOpts =
  defaultOptions{optNormalHandlers         = [addHandler]
                 , optClientStreamHandlers = [addClientStreamHandler]
                 , optServerStreamHandlers = [addServerStreamHandler]
                 , optBiDiStreamHandlers   = [addBiDiHandler]}

main :: IO ()
main = bracket startServer stopServer $ const $ withGRPC $ \grpc ->
  withClient grpc (ClientConfig "localhost" 50051 [] Nothing Nothing) $ \c -> do
    rmAdd <- clientRegisterMethodNormal c addMethod
    rmClientStream <- clientRegisterMethodClientStreaming c addClientStreamMethod
    rmServerStream <- clientRegisterMethodServerStreaming c addServerStreamMethod
    rmBiDiStream <- clientRegisterMethodBiDiStreaming c addBiDiMethod
    defaultMainWith
      defaultConfig{reportFile = Just "benchmarks.html"}
      [ bench "unary request" $ nfIO (addRequest c rmAdd)
      , bench "client stream: 100 messages" $ nfIO (addClientStream c rmClientStream 100)
      , bench "client stream: 1k messages" $ nfIO (addClientStream c rmClientStream 1000)
      , bench "client stream: 10k messages" $ nfIO (addClientStream c rmClientStream 10000)
      , bench "server stream: 100 messages" $ nfIO (addServerStream c rmServerStream 100)
      , bench "server stream: 1k messages" $ nfIO (addServerStream c rmServerStream 1000)
      , bench "server stream: 10k messages" $ nfIO (addServerStream c rmServerStream 10000)
      , bench "bidi stream: 50 messages up, 50 down" $ nfIO (bidiStream c rmBiDiStream 50)
      , bench "bidi stream: 500 message up, 500 down" $ nfIO (bidiStream c rmBiDiStream 500)
      , bench "bidi stream: 5000 messages up, 5000 down" $ nfIO (bidiStream c rmBiDiStream 5000)]

  where startServer = do
          sThrd <- async $ serverLoop serverOpts
          threadDelaySecs 1
          return sThrd

        stopServer sThrd = cancel sThrd >> void (waitCatch sThrd)

        encode = BL.toStrict . toLazyByteString

        addRequest c rmAdd = do
          x <- liftM Fixed $ randomRIO (0,1000)
          y <- liftM Fixed $ randomRIO (0,1000)
          let addEnc = BL.toStrict . toLazyByteString $ AddRequest x y
          clientRequest c rmAdd 5 addEnc mempty >>= \case
            Left e -> fail $ "Got client error on add request: " ++ show e
            Right r -> case fromByteString (rspBody r) of
              Left e -> fail $ "failed to decode add response: " ++ show e
              Right dec
                | dec == AddResponse (x + y) -> return ()
                | otherwise -> fail $ "Got wrong add answer: " ++ show dec ++ "; expected: " ++ show x ++ " + " ++ show y ++ " = " ++ show (x+y)

        addClientStream c rm i = do
          let msg = encode $ AddRequest 1 0
          Right (Just r,_,_,_,_) <- clientWriter c rm 5 mempty $ \send -> do
            replicateM_ i $ send msg
          let decoded = fromByteString r
          when (decoded /= Right (AddResponse (fromIntegral i))) $
            fail $ "clientStream: bad answer: " ++ show decoded ++ "; expected: " ++ show i

        addServerStream c rm i = do
          let msg = encode $ AddRequest (fromIntegral i) 2
          Right (_, _, sd) <- clientReader c rm 5 msg mempty $ \_ recv ->
            replicateM_ i $ do
              Right (Just bs) <- recv
              let Right decoded = fromByteString bs
              when (decoded /= AddResponse 2) $
                fail $ "serverStream: bad response of " ++ show decoded ++ "; expected 2."
          when (sd /= mempty) $ fail $ "bad status details: " ++ show sd

        bidiStream c rm i = do
          Right (_, _, sd) <- clientRW c rm 5 mempty $ \_ recv send done -> do
            forM_ (take i [2,4..]) $ \n -> do
              void $ send $ encode $ AddRequest 1 1
              Right (Just bs) <- recv
              let Right decoded = fromByteString bs
              when (decoded /= AddResponse n) $
                fail $ "bidiStream: got: " ++ show decoded ++ "expected: " ++ show n
            void done
          when (sd /= mempty) $ fail $ "bad StatusDetails: " ++ show sd