{-# LANGUAGE CPP #-}
{-# OPTIONS_GHC -Wno-orphans #-}

module Test.Stress.Cmdline
  ( -- * Types
    Cmdline(..)
  , Role(..)
  , GlobalOpts(..)

    -- ** Client-specific
  , Connect(..)
  , Exec(..)
  , Call(..)

    -- ** Server-specific
  , Security(..)
  , TlsOpts(..)

    -- ** Auxiliary
  , NotPretty(..)

    -- * Parser
  , getCmdline
  ) where

import Control.Applicative ((<|>))
import Data.Foldable (asum, toList)
import Data.Maybe (fromMaybe)
import GHC.Generics (Generic)
import Network.Socket (HostName, PortNumber)
import Options.Applicative qualified as Opt
import Text.Show.Pretty (PrettyVal(..), parseValue)

import Network.GRPC.Client qualified as Client
import Network.GRPC.Common
import Network.GRPC.Common.Compression (Compression)
import Network.GRPC.Common.Compression qualified as Compr
import Network.GRPC.Server.Run

import Paths_grapesy

{-------------------------------------------------------------------------------
  Definitions
-------------------------------------------------------------------------------}

-- | Command specification
data Cmdline = Cmdline {
      cmdRole       :: Role
    , cmdGlobalOpts :: GlobalOpts
    }
  deriving stock (Show, Generic)
  deriving anyclass (PrettyVal)

-- | Should we run the client, servers, or both?
data Role =
      -- | Run the clients
      Client {
          -- | Connect over TLS?
          clientSecurity   :: Maybe (NotPretty Client.ServerValidation)
        , clientServerPort :: PortNumber
        , clientConnects   :: [Connect]

          -- | Insist on this compression scheme for all messages
        , clientCompression :: Maybe (NotPretty Compression)
        }

      -- | Run the server
    | Server {
          serverConfig :: ServerConfig
        }

      -- | Run the automatic stress test suite
    | Driver {
          driverWorkingDir :: Maybe FilePath
        , driverDuration   :: Int
        , driverGenCharts  :: Bool
        }
  deriving stock (Show, Generic)
  deriving anyclass (PrettyVal)

-- | Connections to execute
data Connect = Connect {
      -- | Execute the connections concurrently or sequentially?
      connectExec :: Exec

      -- | Number of connections
    , connectNum :: Int

      -- | Number of times to repeat the calls on the connection
    , callNum :: Int

      -- | Calls to make on the connections
    , connectCalls :: [Call]
    }
  deriving stock (Show, Generic)
  deriving anyclass (PrettyVal)

-- | Concurrent or sequential execution
data Exec = Concurrent | Sequential
  deriving stock (Show, Generic)
  deriving anyclass (PrettyVal)

-- | Types of RPCs
data Call =
      -- | A single message back and forth
      NonStreaming

      -- | Client sends @N@ messages to the server
    | ClientStreaming Int

      -- | Server sends @N@ messages to the client
    | ServerStreaming Int

      -- | Client and server send @N@ messages to each other
    | BiDiStreaming Int
  deriving stock (Show, Generic)
  deriving anyclass (PrettyVal)

data Security =
      Insecure
    | Secure
  deriving stock (Show, Generic)
  deriving anyclass (PrettyVal)

data TlsOpts = TlsOpts {
      tlsPubCert    :: FilePath
    , tlsChainCerts :: [FilePath]
    , tlsPrivKey    :: FilePath
    }
  deriving stock (Show, Generic)
  deriving anyclass (PrettyVal)

mkConfig :: Maybe TlsOpts -> HostName -> PortNumber -> ServerConfig
mkConfig mtls host port =
    case mtls of
      Just TlsOpts{..} ->
        ServerConfig Nothing $ Just SecureConfig {
            secureHost       = host
          , securePort       = port
          , securePubCert    = tlsPubCert
          , secureChainCerts = tlsChainCerts
          , securePrivKey    = tlsPrivKey
          , secureSslKeyLog  = SslKeyLogNone
          }
      Nothing ->
        (`ServerConfig` Nothing) $ Just InsecureConfig {
            insecureHost = Just host
          , insecurePort = port
          }

data GlobalOpts = GlobalOpts {
      optsTracing :: Bool
    }
  deriving stock (Show, Generic)
  deriving anyclass (PrettyVal)

-------------------------------------------------------------------------------
-- Top-level parsers
-------------------------------------------------------------------------------

getCmdline :: IO Cmdline
getCmdline = do
    defaultPub  <- getDataFileName "grpc-demo.pem"
    defaultPriv <- getDataFileName "grpc-demo.key"

    let info  :: Opt.ParserInfo Cmdline
        info = Opt.info
                 (          parseCmdline defaultPub defaultPriv
                   Opt.<**> Opt.helper
                 )
                 Opt.fullDesc

    Opt.execParser info

parseCmdline ::
     FilePath
  -> FilePath
  -> Opt.Parser Cmdline
parseCmdline defaultPub defaultPriv =
    Cmdline
      <$> (parseRole defaultPub defaultPriv <|> pure (Driver Nothing 60 False))
      <*> parseGlobalOpts

parseRole :: FilePath -> FilePath -> Opt.Parser Role
parseRole defaultPub defaultPriv = Opt.subparser $ mconcat [
      sub "client" "Run the client" $
        parseClientRole defaultPub
    , sub "server" "Run the server" $
        parseServerRole defaultPub defaultPriv
    , sub "driver" "Run the stress test driver" $
        parseDriverRole
    ]

-------------------------------------------------------------------------------
-- Client option parsers
-------------------------------------------------------------------------------

parseClientRole :: FilePath -> Opt.Parser Role
parseClientRole defaultPub =
    Client
      <$> (fmap WrapNotPretty <$> parseClientSecurity defaultPub)
      <*> parseClientPort
      <*> parseClientConnects
      <*> Opt.optional (WrapNotPretty <$> parseCompression)

parseClientSecurity :: FilePath -> Opt.Parser (Maybe Client.ServerValidation)
parseClientSecurity defaultPub =
    Opt.optional (
        Opt.flag' () (mconcat [
            Opt.long "secure"
          , Opt.help "Connect over TLS"
          ])
        *> parseServerValidation
      )
  where
    parseServerValidation :: Opt.Parser Client.ServerValidation
    parseServerValidation =
        aux
          <$> (Opt.switch $ mconcat [
                  Opt.long "no-server-validation"
                , Opt.help "Skip server (certificate) validation"
                ])
          <*> (Opt.switch $ mconcat [
                  Opt.long "cert-store-from-system"
                , Opt.help "Enable the system certificate store"
                ])
          <*> (Opt.option Opt.str $ mconcat [
                  Opt.long "cert-store-from-path"
                , Opt.help "Load certificate store from file or directory (set to empty to disable)"
                , Opt.metavar "PATH"
                , Opt.value defaultPub
                , Opt.showDefault
                ])
      where
        aux :: Bool -> Bool -> FilePath -> Client.ServerValidation
        aux noServerValidation certStoreFromSystem certStoreFromPath =
            if noServerValidation then
              Client.NoServerValidation
            else
              Client.ValidateServer $ mconcat . concat $ [
                  [ Client.certStoreFromSystem
                  | certStoreFromSystem
                  ]

                , [ Client.certStoreFromPath certStoreFromPath
                  | not (null certStoreFromPath)
                  ]
                ]

parseClientPort :: Opt.Parser PortNumber
parseClientPort =
    Opt.option Opt.auto (mconcat [
        Opt.long "port"
      , Opt.help "Connect to the server at PORT"
      , Opt.metavar "PORT"
      ])

parseClientConnects :: Opt.Parser [Connect]
parseClientConnects =
    Opt.some $ Connect
      <$> Opt.flag Sequential Concurrent (mconcat [
              Opt.long "concurrent"
            , Opt.help "Open connections concurrently"
            ])
      <*> Opt.option Opt.auto (mconcat [
              Opt.long "num-connections"
            , Opt.help "Open N connections"
            , Opt.metavar "N"
            , Opt.value 1
            , Opt.showDefault
            ])
      <*> Opt.option Opt.auto (mconcat [
              Opt.long "num-calls"
            , Opt.help "Repeat the calls N times on each connection"
            , Opt.metavar "N"
            , Opt.value 1
            , Opt.showDefault
            ])
      <*> Opt.some parseCall

parseCall :: Opt.Parser Call
parseCall =
        nonStreaming
    <|> clientStreaming
    <|> serverStreaming
    <|> bidiStreaming
  where
    nonStreaming :: Opt.Parser Call
    nonStreaming =
        Opt.flag' NonStreaming $ mconcat [
            Opt.long "non-streaming"
          , Opt.help "Make a single non-streaming call"
          ]

    clientStreaming :: Opt.Parser Call
    clientStreaming =
        ClientStreaming
          <$> Opt.option Opt.auto (mconcat [
                  Opt.long "client-streaming"
                , Opt.help "Stream N messages from the client"
                , Opt.metavar "N"
                ])

    serverStreaming :: Opt.Parser Call
    serverStreaming =
        ServerStreaming
          <$> Opt.option Opt.auto (mconcat [
                  Opt.long "server-streaming"
                , Opt.help "Stream N messages from the server"
                , Opt.metavar "N"
                ])

    bidiStreaming :: Opt.Parser Call
    bidiStreaming =
        BiDiStreaming
          <$> Opt.option Opt.auto (mconcat [
                  Opt.long "bidi-streaming"
                , Opt.help "Stream N messages from both the client and server"
                , Opt.metavar "N"
                ])

parseCompression :: Opt.Parser Compression
parseCompression = asum $ map go (toList Compr.allSupportedCompression)
  where
    go :: Compression -> Opt.Parser Compression
    go compr = Opt.flag' compr $ mconcat [
          Opt.long comprId
        , Opt.help $ "Insist on " ++ comprId ++ " compression "
        ]
      where
        comprId :: String
        comprId = show (Compr.compressionId compr)

-------------------------------------------------------------------------------
-- Server option parsers
-------------------------------------------------------------------------------

parseServerRole :: FilePath -> FilePath -> Opt.Parser Role
parseServerRole defaultPub defaultPriv =
    aux
      <$> parseServerPort
      <*> parseServerSecurity defaultPub defaultPriv
  where
    aux :: PortNumber -> Maybe TlsOpts -> Role
    aux port mtls = Server $ mkConfig mtls "127.0.0.1" port

parseServerPort :: Opt.Parser PortNumber
parseServerPort =
    Opt.option Opt.auto (mconcat [
        Opt.long "port"
      , Opt.help "Bind the server to port PORT"
      , Opt.metavar "PORT"
      ])

parseServerSecurity :: FilePath -> FilePath -> Opt.Parser (Maybe TlsOpts)
parseServerSecurity defaultPub defaultPriv =
    Opt.optional $
      Opt.flag' () (mconcat [
          Opt.long "secure"
        , Opt.help "Enable TLS on the servers"
        ])
      *> parseTlsOpts defaultPub defaultPriv

parseTlsOpts :: FilePath -> FilePath -> Opt.Parser TlsOpts
parseTlsOpts defaultPub defaultPriv =
    TlsOpts
      <$> Opt.strOption (mconcat [
              Opt.long "tls-pub"
            , Opt.help "TLS public certificate (X.509 format)"
            , Opt.metavar "CERT_FILE"
            , Opt.showDefault
            , Opt.value defaultPub
            ])
      <*> Opt.many (Opt.strOption $ mconcat [
              Opt.long "tls-cert"
            , Opt.metavar "CERT_FILE"
            , Opt.help "TLS chain certificate (X.509 format)"
            ])
      <*> Opt.strOption (mconcat [
              Opt.long "tls-priv"
            , Opt.metavar "KEY_FILE"
            , Opt.help "TLS private key"
            , Opt.showDefault
            , Opt.value defaultPriv
            ])

-------------------------------------------------------------------------------
-- Driver option parsers
-------------------------------------------------------------------------------

parseDriverRole :: Opt.Parser Role
parseDriverRole =
    Driver
      <$> Opt.optional (Opt.strOption (mconcat [
              Opt.long "working-dir"
            , Opt.help "Write result files to this directory, must already exist"
            , Opt.metavar "DIR"
            ]))
      <*> Opt.option Opt.auto (mconcat [
              Opt.long "duration"
            , Opt.help "Run the stress tests for this many seconds"
            , Opt.metavar "SECONDS"
            , Opt.value 60
            , Opt.showDefault
            ])
      <*> Opt.switch (mconcat [
              Opt.long "gen-charts"
            , Opt.help "Generate heap profile charts for stable components"
            ])

-------------------------------------------------------------------------------
-- Internal auxiliary
-------------------------------------------------------------------------------

parseGlobalOpts :: Opt.Parser GlobalOpts
parseGlobalOpts =
    GlobalOpts
      <$> Opt.switch (mconcat [
              Opt.long "verbose"
            , Opt.short 'v'
            , Opt.help "Trace test execution to stdout"
            ])

sub :: String -> String -> Opt.Parser a -> Opt.Mod Opt.CommandFields a
sub cmd desc parser =
    Opt.command cmd $
      Opt.info (parser Opt.<**> Opt.helper) (Opt.progDesc desc)

-------------------------------------------------------------------------------
-- Auxiliary: pretty-val
-------------------------------------------------------------------------------

newtype NotPretty a = WrapNotPretty { unwrapNotPretty :: a }
  deriving newtype (Show)

instance Show a => PrettyVal (NotPretty a) where
  prettyVal (WrapNotPretty x) =
      fromMaybe
        (error $ "prettyVal: could not parse " ++ show x)
        (parseValue $ show x)

instance PrettyVal PortNumber where
  prettyVal = prettyVal . (fromIntegral :: PortNumber -> Integer)

deriving anyclass instance PrettyVal ServerConfig
deriving anyclass instance PrettyVal InsecureConfig
deriving anyclass instance PrettyVal SecureConfig
deriving anyclass instance PrettyVal SslKeyLog