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

module Test.Driver.Dialogue.Execution (
    ConnUsage(..)
  , execGlobalSteps
  ) where

import Control.Concurrent
import Control.Concurrent.Async
import Control.Monad
import Control.Monad.Catch
import Control.Monad.State
import Data.List (sortBy)
import Data.Ord (comparing)
import Data.Proxy
import Data.Text qualified as Text
import GHC.Stack
import GHC.TypeLits
import Network.HTTP2.Client qualified as HTTP2.Client

import Network.GRPC.Client qualified as Client
import Network.GRPC.Client.Binary qualified as Client.Binary
import Network.GRPC.Common
import Network.GRPC.Common.Binary
import Network.GRPC.Server qualified as Server
import Network.GRPC.Server.Binary qualified as Server.Binary

import Test.Driver.ClientServer
import Test.Driver.Dialogue.Definition
import Test.Driver.Dialogue.TestClock (TestClock)
import Test.Driver.Dialogue.TestClock qualified as TestClock
import Test.Util

{-------------------------------------------------------------------------------
  Endpoints
-------------------------------------------------------------------------------}

type TestProtocol meth = RawRpc "dialogue" meth

type instance RequestMetadata          (TestProtocol meth) = TestMetadata
type instance ResponseInitialMetadata  (TestProtocol meth) = TestMetadata
type instance ResponseTrailingMetadata (TestProtocol meth) = TestMetadata

type TestRpc1 = TestProtocol "test1"
type TestRpc2 = TestProtocol "test2"
type TestRpc3 = TestProtocol "test3"

withClientProxy ::
     RPC
  -> (forall meth.
          SupportsClientRpc (TestProtocol meth)
       => Proxy meth
       -> a)
  -> a
withClientProxy RPC1 k = k (Proxy @"test1")
withClientProxy RPC2 k = k (Proxy @"test2")
withClientProxy RPC3 k = k (Proxy @"test3")

{-------------------------------------------------------------------------------
  Test failures
-------------------------------------------------------------------------------}

data TestFailure = TestFailure CallStack Failure
  deriving stock (Show)
  deriving anyclass (Exception)

data Failure =
    -- | Thrown by the server when an unexpected new RPC is initiated
    UnexpectedRequest

    -- | Received an unexpected value
  | Unexpected ReceivedUnexpected
  deriving stock (Show)
  deriving anyclass (Exception)

data ReceivedUnexpected = forall a b. (Show a, Show b) => ReceivedUnexpected {
      received   :: a -- The value we received
    , expectInfo :: b -- Some additional info that can help debug the issue
    }

deriving stock instance Show ReceivedUnexpected

expect ::
     (MonadThrow m, Show a, Show info, HasCallStack)
  => info
  -> (a -> Bool)  -- ^ Expected
  -> a            -- ^ Actually received
  -> m ()
expect expectInfo isExpected received
  | isExpected received
  = return ()

  | otherwise
  = throwM $ TestFailure callStack $
      Unexpected $ ReceivedUnexpected{
          received
        , expectInfo
        }

{-------------------------------------------------------------------------------
  Timeouts
-------------------------------------------------------------------------------}

-- | Timeout for waiting for the test clock
timeoutClock :: Int
timeoutClock = 5

-- | Timeout for waiting for the green liht
timeoutGreenLight :: Int
timeoutGreenLight = 5

-- | Timeout for executing all the actions in a client or handler
timeoutLocal :: Int
timeoutLocal = 20

-- | Timeout for waiting for a call to fail
timeoutFailure :: Int
timeoutFailure = 5

-- | Timeout for receiving a stream element
timeoutReceive :: Int
timeoutReceive = 5

{-------------------------------------------------------------------------------
  Health
-------------------------------------------------------------------------------}

-- | Health of the peer (server/client)
--
-- When the client is expecting a response from the server, it needs to know the
-- "health" of the server, that is, is the server still alive, or did it fail
-- with some kind exception? The same is true for the server when it expects a
-- response from the client. Therefore, the client interpretation keeps track of
-- the health of the server, and vice versa.
data PeerHealth =
    PeerAlive

    -- | Peer terminated
    --
    -- The peer might have thrown a deliberate exception, or simply terminated
    -- early without properly closing the connection.
  | PeerTerminated (Maybe DeliberateException)
  deriving stock (Show)

ifPeerAlive :: PeerHealth -> PeerHealth -> PeerHealth
ifPeerAlive PeerAlive             = id
ifPeerAlive (PeerTerminated mErr) = const (PeerTerminated mErr)

{-------------------------------------------------------------------------------
  Client-side interpretation
-------------------------------------------------------------------------------}

clientLocal ::
     HasCallStack
  => TestClock
  -> Client.Call (TestProtocol meth)
  -> LocalSteps
  -> IO ()
clientLocal clock call = \(LocalSteps steps) ->
    evalStateT (go steps) PeerAlive
  where
    go :: [(TestClock.Tick, LocalStep)] -> StateT PeerHealth IO ()
    go []                     = return ()
    go ((tick, step) : steps) = do
        case step of
          ClientAction action -> do
            within timeoutClock step $ TestClock.waitForTick clock tick
            continue <- clientAct tick action `finally` TestClock.advance clock
            when continue $ go steps
          ServerAction action -> do
            TestClock.giveGreenLight clock tick
            reactToServer tick action
            go steps

    -- Client action
    --
    -- Returns 'True' if we should continue executing more actions, or
    -- exit (thereby closing the RPC call)
    clientAct :: TestClock.Tick -> ClientAction -> StateT PeerHealth IO Bool
    clientAct tick action =
        case action of
          Initiate _ ->
            error "clientLocal: unexpected Initiate"
          Send x -> do
            peerHealth <- get
            case peerHealth of
              PeerAlive        -> Client.Binary.sendInput call x
              PeerTerminated _ -> liftIO $ waitForServerDisconnect
            return True
          Terminate mException -> do
            -- See discussion in 'TestClock' for why we need to wait here
            peerHealth <- get
            case peerHealth of
              PeerTerminated _ -> return ()
              PeerAlive        -> within timeoutGreenLight action $
                                    TestClock.waitForGreenLight clock tick
            case mException of
              Just ex -> throwM $ DeliberateException ex
              Nothing -> return False

    reactToServer :: TestClock.Tick -> ServerAction -> StateT PeerHealth IO ()
    reactToServer tick action =
        case action of
          Initiate expectedMetadata -> liftIO $ do
            receivedMetadata <- within timeoutReceive action $
                                  Client.recvResponseInitialMetadata call
            expect (tick, action) (== expectedMetadata) receivedMetadata
          Send (FinalElem a b) -> do
            -- On the client side, when the server sends the final message, we
            -- will receive that final message in one HTTP data frame, and then
            -- the trailers in another. This means that when we get the message,
            -- we do not yet know if this is in fact the last.
            -- (This is different on the server side, because gRPC does not
            -- support trailers on the client side.)
            reactToServer tick $ Send (StreamElem a)
            reactToServer tick $ Send (NoMoreElems b)
          Send expectedElem -> do
            mOut <- try $ within timeoutReceive action $
                      Client.Binary.recvOutput call
            expect (tick, action) (isExpectedElem expectedElem) mOut
          Terminate mErr -> do
            mOut <- try $ within timeoutReceive action $
                      Client.Binary.recvOutput call
            let mErr'       = DeliberateException <$> mErr
                expectation = isGrpcException mErr'
            expect (tick, action) expectation mOut
            modify $ ifPeerAlive $ PeerTerminated mErr'

    -- Wait for the server disconnect to become visible
    --
    -- In principle we could check if we can still /receive/ messages from the
    -- server to see if we can /send/ messages to the server: gRPC does not
    -- allow the server to half-close the connection (only the client). For
    -- consistency, however, we simply wait until sending fails.
    --
    -- See 'waitForClientDisconnect' for additional discussion.
    waitForServerDisconnect :: IO ()
    waitForServerDisconnect =
        within timeoutFailure () $ loop
      where
        loop :: IO ()
        -- We only do this when we know the client has terminated, so the
        -- /type/ of the message we send here as a probe does not matter.
        loop = do
            mFailed <- try $ Client.Binary.sendNextInput call ()
            case mFailed of
              Left (_ :: GrpcException) ->
                return ()
              Right () -> do
                threadDelay 10_000
                loop

    isExpectedElem ::
         StreamElem TestMetadata Int
      -> Either GrpcException (StreamElem TestMetadata Int)
      -> Bool
    isExpectedElem _ (Left _) = False
    isExpectedElem expectedElem (Right streamElem) = expectedElem == streamElem

    isGrpcException ::
         Maybe DeliberateException
      -> Either GrpcException (StreamElem TestMetadata Int)
      -> Bool
    isGrpcException mErr (Left err) = and [
          grpcError        err == GrpcUnknown
        , grpcErrorMessage err == Just (mconcat [
              "Server-side exception: "
            , case mErr of
                Nothing   -> "HandlerTerminated"
                Just err' -> Text.pack $ show err'
            ])
        ]
    isGrpcException _ (Right _) = False

clientGlobal ::
     TestClock
  -> ConnUsage
     -- ^ Use new connection for each RPC call?
     --
     -- Multiple RPC calls on a single connection /ought/ to be independent of
     -- each other. Something going wrong on one should not affect another.
  -> GlobalSteps
  -> TestClient
clientGlobal clock connUsage global connParams testServer delimitTestScope =
    case connUsage of
      ConnPerRPC ->                  go Nothing  [] (getGlobalSteps global)
      SharedConn -> withConn $ \c -> go (Just c) [] (getGlobalSteps global)
  where
    withConn :: (Client.Connection -> IO ()) -> IO ()
    withConn = Client.withConnection connParams testServer

    go :: Maybe Client.Connection -> [Async ()] -> [LocalSteps] -> IO ()
    go _ threads [] = do
        -- Wait for all threads to finish
        --
        -- This also ensures that if any of these threads threw an exception,
        -- that is now rethrown here in the main test. This will also cause us
        -- to leave the scope of all enclosing calls to @withAsync@, thereby
        -- cancelling all other concurrent threads.
        --
        -- (It is therefore important that we catch any /excepted/ exceptions
        -- locally; this is done by the call to @delimitTestScope@.)
        mapM_ wait threads
    go mConn threads (c:cs) =
        withAsync (within timeoutLocal c $ runLocalSteps mConn c) $ \thread ->
          go mConn (thread:threads) cs

    runLocalSteps :: Maybe Client.Connection -> LocalSteps -> IO ()
    runLocalSteps mConn (LocalSteps steps) = delimitTestScope $ do
        case steps of
          (tick, ClientAction (Initiate (metadata, rpc))) : steps' -> do
            TestClock.waitForTick clock tick

            withClientProxy rpc $ startCall mConn metadata steps'
          _otherwise ->
            error $ "clientGlobal: expected Initiate, got " ++ show steps

    startCall :: forall (meth :: Symbol).
         SupportsClientRpc (TestProtocol meth)
      => Maybe Client.Connection
      -> TestMetadata
      -> [(TestClock.Tick, LocalStep)]
      -> Proxy meth -> IO ()
    startCall mConn metadata steps' _ = do
        (case mConn of
            Just conn -> ($ conn)
            Nothing   -> withConn) $ \conn ->
          Client.withRPC conn params (Proxy @(TestProtocol meth)) $ \call -> do
            -- We wait for the /server/ to advance the test clock (so that
            -- we are sure the next step doesn't happen until the
            -- connection is established).
            --
            -- NOTE: We could instead wait for the server to send the
            -- initial metadata; this too would provide evidence that the
            -- connection has been established. However, doing so
            -- precludes a class of correct behaviour: the server might
            -- not respond with that initial metadata until the client has
            -- sent some messages.
            clientLocal clock call (LocalSteps steps')
     where
       -- Timeouts are outside the scope of these tests: it's too finicky
       -- to relate timeouts (in seconds) to specific test execution. We
       -- do test exceptions in general here; the specific exception
       -- arising from a timeout we test elsewhere.
       params = def {
           Client.callRequestMetadata = metadata
         }

{-------------------------------------------------------------------------------
  Server-side interpretation

  The server-side is slightly different, since the infrastructure spawns
  threads on our behalf (one for each incoming RPC).
-------------------------------------------------------------------------------}

serverLocal ::
     TestClock
  -> Server.Call (TestProtocol meth)
  -> LocalSteps -> IO ()
serverLocal clock call = \(LocalSteps steps) -> do
    evalStateT (go steps) PeerAlive
  where
    go :: [(TestClock.Tick, LocalStep)] -> StateT PeerHealth IO ()
    go []                     = return ()
    go ((tick, step) : steps) =
        case step of
          ServerAction action -> do
            within timeoutClock step $ TestClock.waitForTick clock tick
            continue <- serverAct tick action `finally` TestClock.advance clock
            when continue $ go steps
          ClientAction action -> do
            TestClock.giveGreenLight clock tick
            reactToClient tick action
            go steps

    -- Server action
    --
    -- Returns 'True' if we should continue executing the other actions, or
    -- terminate (thereby terminating the handler)
    serverAct :: TestClock.Tick -> ServerAction -> StateT PeerHealth IO Bool
    serverAct tick action =
        case action of
          Initiate metadata -> liftIO $ do
            Server.setResponseInitialMetadata call metadata
            Server.initiateResponse call
            return True
          Send x -> do
            peerHealth <- get
            case peerHealth of
              PeerAlive        -> liftIO $ Server.Binary.sendOutput call x
              PeerTerminated _ -> liftIO $ waitForClientDisconnect
            return True
          Terminate mException -> do
            peerHealth <- get
            case peerHealth of
              PeerTerminated _ -> return ()
              PeerAlive        -> within timeoutGreenLight action $
                                    TestClock.waitForGreenLight clock tick
            case mException of
              Just ex -> throwM $ DeliberateException ex
              Nothing -> return False

    reactToClient :: TestClock.Tick -> ClientAction -> StateT PeerHealth IO ()
    reactToClient tick action =
        case action of
          Initiate _ ->
            error "serverLocal: unexpected ClientInitiateRequest"
          Send expectedElem -> do
            mInp <- liftIO $ try $ within timeoutReceive action $
                      Server.Binary.recvInput call
            expect (tick, action) (isExpectedElem expectedElem) mInp
          Terminate mErr -> do
            mInp <- liftIO $ try $ within timeoutReceive action $
                      Server.Binary.recvInput call
            expect (tick, action) isExpectedDisconnect mInp
            modify $ ifPeerAlive $ PeerTerminated $ DeliberateException <$> mErr

    -- Wait for the client disconnect to become visible
    --
    -- The only way to know that we cannot send messages anymore to a client
    -- that has terminated is by trying. Although the /receiving/ thread may
    -- terminate more-or-less immediately, this does not necessarily indicate
    -- any kind of failure: the client may simply have put the call in
    -- half-closed mode.
    waitForClientDisconnect :: IO ()
    waitForClientDisconnect =
        within timeoutFailure () $ loop
      where
        loop :: IO ()
        -- We only do this when we know the client has terminated, so the
        -- /type/ of the message we send here as a probe does not matter.
        loop = do
            mFailed <- try $ Server.Binary.sendNextOutput call ()
            case mFailed of
              Left (_ :: Server.ClientDisconnected) ->
                return ()
              Right () -> do
                threadDelay 10_000
                loop

    isExpectedElem ::
         StreamElem NoMetadata Int
      -> Either Server.ClientDisconnected (StreamElem NoMetadata Int)
      -> Bool
    isExpectedElem _ (Left _) = False
    isExpectedElem expectedElem (Right streamElem) = expectedElem == streamElem

    isExpectedDisconnect ::
         Either Server.ClientDisconnected (StreamElem NoMetadata Int)
      -> Bool
    isExpectedDisconnect (Left (Server.ClientDisconnected e _))
      | Just HTTP2.Client.ConnectionIsClosed <- fromException e
      = True
      | otherwise
      = False
    isExpectedDisconnect _ = False

serverGlobal ::
     HasCallStack
  => TestClock
  -> MVar GlobalSteps
    -- ^ Unlike in the client case, the grapesy infrastructure spawns a new
    -- thread for each incoming connection. To know which part of the test this
    -- particular handler corresponds to, we take the next 'LocalSteps' from
    -- this @MVar@. Since all requests are started by the client from /one/
    -- thread, the order of these incoming requests is deterministic.
  -> Server.Call (TestProtocol meth)
  -> IO ()
serverGlobal clock globalStepsVar call = do
    steps <- modifyMVar globalStepsVar (getNextSteps . getGlobalSteps)

    -- See discussion in clientGlobal (runLocalSteps)
    TestClock.advance clock

    case getLocalSteps steps of
      (tick, step@(ClientAction (Initiate (metadata, _rpc)))) : steps' -> do
        receivedMetadata <- Server.getRequestMetadata call
        -- It is important that we do this 'expect' outside the scope of the
        -- @modifyMVar@: if we do not, then if the expect fails, we'd leave the
        -- @MVar@ unchanged, and the next request would use the wrong steps.
        expect (tick, step) (== metadata) $ receivedMetadata
        within timeoutLocal steps' $ serverLocal clock call (LocalSteps steps')
      _otherwise ->
         error "serverGlobal: expected ClientInitiateRequest"
  where
    getNextSteps :: [LocalSteps] -> IO (GlobalSteps, LocalSteps)
    getNextSteps [] = do
        throwM $ TestFailure callStack $ UnexpectedRequest
    getNextSteps (LocalSteps steps:global') =
        return (GlobalSteps global', LocalSteps steps)

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

data ConnUsage = SharedConn | ConnPerRPC

execGlobalSteps :: ConnUsage -> GlobalSteps -> IO ClientServerTest
execGlobalSteps connUsage steps = do
    globalStepsVar <- newMVar (order steps)
    clock          <- TestClock.new

    let handler :: forall (meth :: Symbol).
             SupportsServerRpc (TestProtocol meth)
          => Proxy (TestProtocol meth)
          -> Server.SomeRpcHandler IO
        handler _ = Server.someRpcHandler $
          Server.mkRpcHandler @(TestProtocol meth) $ \call ->
            serverGlobal clock globalStepsVar call

    return ClientServerTest {
        config = def {
            isExpectedClientException = \e -> or [
                isDeliberateException e
              , clientTerminatesEarly && isGrpcCancelled e
              ]
          , isExpectedServerException = \e -> or [
                isDeliberateException e
              , serverTerminatesEarly && isHandlerTerminated e
              ]
          }
      , client = clientGlobal clock connUsage steps
      , server = [
            handler (Proxy @TestRpc1)
          , handler (Proxy @TestRpc2)
          , handler (Proxy @TestRpc3)
          ]
      }
  where
    clientTerminatesEarly, serverTerminatesEarly :: Bool
    (clientTerminatesEarly, serverTerminatesEarly) = hasEarlyTermination steps

    -- For 'clientGlobal' the order doesn't matter, because it spawns a thread
    -- for each 'LocalSteps'. The server however doesn't get this option; the
    -- threads /get/ spawnwed for each incoming connection, and must feel off
    -- the appropriate steps. It's therefore important that it will get these
    -- in the order that they come in.
    order :: GlobalSteps -> GlobalSteps
    order (GlobalSteps threads) = GlobalSteps $
        sortBy (comparing firstTick) threads
     where
       firstTick :: LocalSteps -> TestClock.Tick
       firstTick (LocalSteps []) =
           error "execGlobalSteps: unexpected empty LocalSteps"
       firstTick (LocalSteps ((tick, _):_)) =
           tick