{-# LANGUAGE NamedFieldPuns #-}
{-# LANGUAGE RankNTypes #-}

module Utxorpc.Logged
  ( UtxorpcServiceLogger (..),
    RequestLogger,
    ReplyLogger,
    ServerStreamLogger,
    ServerStreamEndLogger,
    loggedUnary,
    loggedUnaryHandler,
    loggedSStream,
    loggedSStreamHandler,
  )
where

import Control.Monad.IO.Class (MonadIO, liftIO)
import qualified Data.ByteString.Char8 as BS
import Data.UUID (UUID)
import Data.UUID.V4 (nextRandom)
import Network.GRPC.HTTP2.Encoding (GRPCInput, GRPCOutput)
import Network.GRPC.HTTP2.Types (IsRPC (..))
import Network.GRPC.Server (ServiceHandler, UnaryHandler)
import Network.GRPC.Server.Handlers.Trans (ServerStream (..), ServerStreamHandler, serverStream, unary)
import Network.Wai (Request (..))

-- | A record of logging functions that runs in the same monad as
-- the request handlers. Monadic state is passed along throughout the
-- lifecycle of responding to a request. This means that changes to the
-- monadic state in the request logger is seen by the stream logger, stream
-- handler and logger, and reply logger. An `unlift` function to run the monad
-- in IO is provided to @'runUtxorpc'@.
data UtxorpcServiceLogger m = UtxorpcServiceLogger
  { forall (m :: * -> *). UtxorpcServiceLogger m -> RequestLogger m
requestLogger :: RequestLogger m,
    forall (m :: * -> *). UtxorpcServiceLogger m -> RequestLogger m
replyLogger :: ReplyLogger m,
    forall (m :: * -> *).
UtxorpcServiceLogger m -> ServerStreamLogger m
serverStreamLogger :: ServerStreamLogger m,
    forall (m :: * -> *).
UtxorpcServiceLogger m -> ServerStreamEndLogger m
serverStreamEndLogger :: ServerStreamEndLogger m
  }

-- | Log incoming requests.
type RequestLogger m =
  forall i.
  (Show i) =>
  -- | The RPC path
  BS.ByteString ->
  -- | Request metadata
  Request ->
  -- | A UUID generated for this request and passed to stream and reply loggers.
  UUID ->
  -- | The request message
  i ->
  m ()

-- | Log outgoing replies.
type ReplyLogger m =
  forall o.
  (Show o) =>
  -- | The RPC path
  BS.ByteString ->
  -- | Request metadata
  Request ->
  -- Generated for the request that caused this reply
  UUID ->
  -- | The reply message
  o ->
  m ()

-- | Log outgoing server stream messages.
type ServerStreamLogger m =
  forall o.
  (Show o) =>
  -- | The RPC path
  BS.ByteString ->
  -- | Request metadata
  Request ->
  -- | The UUID generated for the request that generated this stream,
  -- and the 0-based index of the message in the stream.
  (UUID, Int) ->
  -- | The stream message
  o ->
  m ()

-- | Log the end of a server stream.
type ServerStreamEndLogger m =
  -- | The RPC path
  BS.ByteString ->
  -- | Request metadata
  Request ->
  -- | The UUID generated for the request that generated this stream,
  -- and the 0-based index of the message in the stream.
  (UUID, Int) ->
  m ()

-- | Creates a ServiceHandler that warp-grpc uses to handle requests
loggedUnary ::
  (MonadIO m, GRPCInput r i, GRPCOutput r o, Show i, Show o) =>
  -- | An `unlift` function for the logger and handler monad
  -- Monad state is carried through from request logger, to handler, to reply logger,
  -- So changes to the monad state in the request logger is seen by the handler and reply logger.
  (forall x. m x -> IO x) ->
  -- | The RPC
  r ->
  -- | Generate a reply from request metadata and a proto Message
  UnaryHandler m i o ->
  -- | A logger that runs in the same monad as the handlers
  Maybe (UtxorpcServiceLogger m) ->
  ServiceHandler
loggedUnary :: forall (m :: * -> *) r i o.
(MonadIO m, GRPCInput r i, GRPCOutput r o, Show i, Show o) =>
(forall x. m x -> IO x)
-> r
-> UnaryHandler m i o
-> Maybe (UtxorpcServiceLogger m)
-> ServiceHandler
loggedUnary forall x. m x -> IO x
unlift r
rpc UnaryHandler m i o
handler Maybe (UtxorpcServiceLogger m)
maybeLogger =
  (forall x. m x -> IO x)
-> r -> UnaryHandler m i o -> ServiceHandler
forall (m :: * -> *) r i o.
(MonadIO m, GRPCInput r i, GRPCOutput r o) =>
(forall x. m x -> IO x)
-> r -> UnaryHandler m i o -> ServiceHandler
unary m x -> IO x
forall x. m x -> IO x
unlift r
rpc (UnaryHandler m i o -> ServiceHandler)
-> UnaryHandler m i o -> ServiceHandler
forall a b. (a -> b) -> a -> b
$ UnaryHandler m i o
-> (UtxorpcServiceLogger m -> UnaryHandler m i o)
-> Maybe (UtxorpcServiceLogger m)
-> UnaryHandler m i o
forall b a. b -> (a -> b) -> Maybe a -> b
maybe UnaryHandler m i o
handler UtxorpcServiceLogger m -> UnaryHandler m i o
loggedHandler Maybe (UtxorpcServiceLogger m)
maybeLogger
  where
    -- Generate UUID here for easier testing of `loggedUnaryHandler`
    loggedHandler :: UtxorpcServiceLogger m -> UnaryHandler m i o
loggedHandler UtxorpcServiceLogger m
logger Request
req i
msg = do
      UUID
uuid <- IO UUID -> m UUID
forall a. IO a -> m a
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO IO UUID
nextRandom
      r
-> UnaryHandler m i o
-> UUID
-> UtxorpcServiceLogger m
-> UnaryHandler m i o
forall (m :: * -> *) i o r.
(MonadIO m, Show i, Show o, IsRPC r) =>
r
-> UnaryHandler m i o
-> UUID
-> UtxorpcServiceLogger m
-> UnaryHandler m i o
loggedUnaryHandler r
rpc UnaryHandler m i o
handler UUID
uuid UtxorpcServiceLogger m
logger Request
req i
msg

loggedUnaryHandler ::
  (MonadIO m, Show i, Show o, IsRPC r) =>
  r ->
  UnaryHandler m i o ->
  UUID ->
  UtxorpcServiceLogger m ->
  UnaryHandler m i o
loggedUnaryHandler :: forall (m :: * -> *) i o r.
(MonadIO m, Show i, Show o, IsRPC r) =>
r
-> UnaryHandler m i o
-> UUID
-> UtxorpcServiceLogger m
-> UnaryHandler m i o
loggedUnaryHandler
  r
rpc
  UnaryHandler m i o
handler
  UUID
uuid
  UtxorpcServiceLogger {RequestLogger m
requestLogger :: forall (m :: * -> *). UtxorpcServiceLogger m -> RequestLogger m
requestLogger :: RequestLogger m
requestLogger, RequestLogger m
replyLogger :: forall (m :: * -> *). UtxorpcServiceLogger m -> RequestLogger m
replyLogger :: RequestLogger m
replyLogger}
  Request
req
  i
msg =
    do
      ByteString -> Request -> UUID -> i -> m ()
RequestLogger m
requestLogger (r -> ByteString
forall t. IsRPC t => t -> ByteString
path r
rpc) Request
req UUID
uuid i
msg
      o
reply <- UnaryHandler m i o
handler Request
req i
msg
      ByteString -> Request -> UUID -> o -> m ()
RequestLogger m
replyLogger ByteString
rpcPath Request
req UUID
uuid o
reply
      o -> m o
forall a. a -> m a
forall (m :: * -> *) a. Monad m => a -> m a
return o
reply
    where
      rpcPath :: ByteString
rpcPath = r -> ByteString
forall t. IsRPC t => t -> ByteString
path r
rpc

-- | Creates a ServiceHandler that warp-grpc uses to handle stream requests
loggedSStream ::
  (MonadIO m, GRPCInput r i, GRPCOutput r o, Show i, Show o) =>
  -- | An unlift function for the logger and handler
  -- Monadic state changes are passed from request logger to stream logger and handlers, and so on.
  -- So changes to the monadic state in the request logger are seen by the handler and other loggers.
  (forall x. m x -> IO x) ->
  -- | The RPC
  r ->
  -- | A function that, given request metadata and a protobuf Message,
  -- generates an initial stream state and a function that folds over the stream state to produce
  -- a stream of messages.
  ServerStreamHandler m i o a ->
  -- | A logger that runs in the same monad as the handler
  Maybe (UtxorpcServiceLogger m) ->
  ServiceHandler
loggedSStream :: forall (m :: * -> *) r i o a.
(MonadIO m, GRPCInput r i, GRPCOutput r o, Show i, Show o) =>
(forall x. m x -> IO x)
-> r
-> ServerStreamHandler m i o a
-> Maybe (UtxorpcServiceLogger m)
-> ServiceHandler
loggedSStream forall x. m x -> IO x
unlift r
rpc ServerStreamHandler m i o a
handler Maybe (UtxorpcServiceLogger m)
Nothing = (forall x. m x -> IO x)
-> r -> ServerStreamHandler m i o a -> ServiceHandler
forall (m :: * -> *) r i o a.
(MonadIO m, GRPCInput r i, GRPCOutput r o) =>
(forall x. m x -> IO x)
-> r -> ServerStreamHandler m i o a -> ServiceHandler
serverStream m x -> IO x
forall x. m x -> IO x
unlift r
rpc ServerStreamHandler m i o a
handler
loggedSStream forall x. m x -> IO x
unlift r
rpc ServerStreamHandler m i o a
handler (Just UtxorpcServiceLogger m
logger) =
  (forall x. m x -> IO x)
-> r -> ServerStreamHandler m i o (a, Int) -> ServiceHandler
forall (m :: * -> *) r i o a.
(MonadIO m, GRPCInput r i, GRPCOutput r o) =>
(forall x. m x -> IO x)
-> r -> ServerStreamHandler m i o a -> ServiceHandler
serverStream m x -> IO x
forall x. m x -> IO x
unlift r
rpc ServerStreamHandler m i o (a, Int)
loggedHandler
  where
    loggedHandler :: ServerStreamHandler m i o (a, Int)
loggedHandler Request
req i
msg = do
      UUID
uuid <- IO UUID -> m UUID
forall a. IO a -> m a
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO IO UUID
nextRandom
      r
-> ServerStreamHandler m i o a
-> UUID
-> UtxorpcServiceLogger m
-> ServerStreamHandler m i o (a, Int)
forall (m :: * -> *) r i o a.
(MonadIO m, IsRPC r, Show i, Show o) =>
r
-> ServerStreamHandler m i o a
-> UUID
-> UtxorpcServiceLogger m
-> ServerStreamHandler m i o (a, Int)
loggedSStreamHandler r
rpc ServerStreamHandler m i o a
handler UUID
uuid UtxorpcServiceLogger m
logger Request
req i
msg

loggedSStreamHandler ::
  (MonadIO m, IsRPC r, Show i, Show o) =>
  r ->
  ServerStreamHandler m i o a ->
  UUID ->
  UtxorpcServiceLogger m ->
  ServerStreamHandler m i o (a, Int)
loggedSStreamHandler :: forall (m :: * -> *) r i o a.
(MonadIO m, IsRPC r, Show i, Show o) =>
r
-> ServerStreamHandler m i o a
-> UUID
-> UtxorpcServiceLogger m
-> ServerStreamHandler m i o (a, Int)
loggedSStreamHandler
  r
rpc
  ServerStreamHandler m i o a
handler
  UUID
uuid
  UtxorpcServiceLogger {RequestLogger m
requestLogger :: forall (m :: * -> *). UtxorpcServiceLogger m -> RequestLogger m
requestLogger :: RequestLogger m
requestLogger, ServerStreamLogger m
serverStreamLogger :: forall (m :: * -> *).
UtxorpcServiceLogger m -> ServerStreamLogger m
serverStreamLogger :: ServerStreamLogger m
serverStreamLogger, ServerStreamEndLogger m
serverStreamEndLogger :: forall (m :: * -> *).
UtxorpcServiceLogger m -> ServerStreamEndLogger m
serverStreamEndLogger :: ServerStreamEndLogger m
serverStreamEndLogger}
  Request
req
  i
msg = do
    -- Log request
    ByteString -> Request -> UUID -> i -> m ()
RequestLogger m
requestLogger ByteString
rpcPath Request
req UUID
uuid i
msg
    -- Get initial stream state and stream output function
    (a
initStreamState, ServerStream {a -> m (Maybe (a, o))
serverStreamNext :: a -> m (Maybe (a, o))
serverStreamNext :: forall (m :: * -> *) o a.
ServerStream m o a -> a -> m (Maybe (a, o))
serverStreamNext}) <- ServerStreamHandler m i o a
handler Request
req i
msg
    -- Wrap stream output function with logging
    let loggedStreamNext :: (a, Int) -> m (Maybe ((a, Int), o))
loggedStreamNext = (a -> m (Maybe (a, o))) -> (a, Int) -> m (Maybe ((a, Int), o))
forall {b} {t} {a}.
Show b =>
(t -> m (Maybe (a, b))) -> (t, Int) -> m (Maybe ((a, Int), b))
mkLoggedStreamNext a -> m (Maybe (a, o))
serverStreamNext
    -- The unwrapped handler returns the initial stream state and stream output function
    -- We add initial log state and return the wrapped stream output function
    ((a, Int), ServerStream m o (a, Int))
-> m ((a, Int), ServerStream m o (a, Int))
forall a. a -> m a
forall (m :: * -> *) a. Monad m => a -> m a
return ((a
initStreamState, Int
0), ((a, Int) -> m (Maybe ((a, Int), o))) -> ServerStream m o (a, Int)
forall (m :: * -> *) o a.
(a -> m (Maybe (a, o))) -> ServerStream m o a
ServerStream (a, Int) -> m (Maybe ((a, Int), o))
loggedStreamNext)
    where
      mkLoggedStreamNext :: (t -> m (Maybe (a, b))) -> (t, Int) -> m (Maybe ((a, Int), b))
mkLoggedStreamNext t -> m (Maybe (a, b))
getNext (t
streamState, Int
index) = do
        -- Get next chunk
        Maybe (a, b)
next <- t -> m (Maybe (a, b))
getNext t
streamState
        case Maybe (a, b)
next of
          Maybe (a, b)
Nothing -> do
            -- Log end of stream
            ServerStreamEndLogger m
serverStreamEndLogger ByteString
rpcPath Request
req (UUID
uuid, Int
index)
            -- Return end of stream
            Maybe ((a, Int), b) -> m (Maybe ((a, Int), b))
forall a. a -> m a
forall (m :: * -> *) a. Monad m => a -> m a
return Maybe ((a, Int), b)
forall a. Maybe a
Nothing
          Just (a
nextStreamState, b
replyMsg) -> do
            -- Log chunk
            ByteString -> Request -> (UUID, Int) -> b -> m ()
ServerStreamLogger m
serverStreamLogger ByteString
rpcPath Request
req (UUID
uuid, Int
index) b
replyMsg
            -- The unwrapped stream output function returns the next stream state and the message to send
            -- We add log state
            Maybe ((a, Int), b) -> m (Maybe ((a, Int), b))
forall a. a -> m a
forall (m :: * -> *) a. Monad m => a -> m a
return (Maybe ((a, Int), b) -> m (Maybe ((a, Int), b)))
-> Maybe ((a, Int), b) -> m (Maybe ((a, Int), b))
forall a b. (a -> b) -> a -> b
$ ((a, Int), b) -> Maybe ((a, Int), b)
forall a. a -> Maybe a
Just ((a
nextStreamState, Int
index Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1), b
replyMsg)

      rpcPath :: ByteString
rpcPath = r -> ByteString
forall t. IsRPC t => t -> ByteString
path r
rpc