{-# LANGUAGE ExistentialQuantification #-}
{-# LANGUAGE LambdaCase                #-}
{-# LANGUAGE TypeFamilies              #-}
module Metro.TP.WebSockets
  ( WebSocket
  , serverConfig
  , clientConfig
  ) where

import           Data.ByteString           (ByteString, empty)
import qualified Data.ByteString.Char8     as BC
import qualified Data.ByteString.Lazy      as BL
import           Metro.Class               (Transport (..))
import           Network.WebSockets        as WS
import qualified Network.WebSockets.Stream as WS

mkStream :: Transport tp => tp -> IO WS.Stream
mkStream :: tp -> IO Stream
mkStream transport :: tp
transport =
  IO (Maybe ByteString) -> (Maybe ByteString -> IO ()) -> IO Stream
WS.makeStream
    (do
        ByteString
bs <- tp -> Int -> IO ByteString
forall transport.
Transport transport =>
transport -> Int -> IO ByteString
recvData tp
transport 8192
        Maybe ByteString -> IO (Maybe ByteString)
forall (m :: * -> *) a. Monad m => a -> m a
return (Maybe ByteString -> IO (Maybe ByteString))
-> Maybe ByteString -> IO (Maybe ByteString)
forall a b. (a -> b) -> a -> b
$ if ByteString -> Bool
BC.null ByteString
bs then Maybe ByteString
forall a. Maybe a
Nothing else ByteString -> Maybe ByteString
forall a. a -> Maybe a
Just ByteString
bs)
    (\case
        Nothing -> () -> IO ()
forall (m :: * -> *) a. Monad m => a -> m a
return ()
        Just bl :: ByteString
bl -> tp -> ByteString -> IO ()
forall transport.
Transport transport =>
transport -> ByteString -> IO ()
sendData tp
transport (ByteString -> IO ()) -> ByteString -> IO ()
forall a b. (a -> b) -> a -> b
$ ByteString -> ByteString
BL.toStrict ByteString
bl)

wsRecvData :: WS.Connection -> Int -> IO ByteString
wsRecvData :: Connection -> Int -> IO ByteString
wsRecvData conn :: Connection
conn _ = do
  DataMessage
msg <- Connection -> IO DataMessage
WS.receiveDataMessage Connection
conn
  case DataMessage
msg of
    WS.Binary bs :: ByteString
bs -> ByteString -> IO ByteString
forall (f :: * -> *) a. Applicative f => a -> f a
pure (ByteString -> IO ByteString) -> ByteString -> IO ByteString
forall a b. (a -> b) -> a -> b
$ ByteString -> ByteString
BL.toStrict ByteString
bs
    _            -> ByteString -> IO ByteString
forall (f :: * -> *) a. Applicative f => a -> f a
pure ByteString
empty

wsSendData :: WS.Connection -> ByteString -> IO ()
wsSendData :: Connection -> ByteString -> IO ()
wsSendData conn :: Connection
conn = Connection -> ByteString -> IO ()
forall a. WebSocketsData a => Connection -> a -> IO ()
WS.sendBinaryData Connection
conn (ByteString -> IO ())
-> (ByteString -> ByteString) -> ByteString -> IO ()
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ByteString -> ByteString
BL.fromStrict

data WebSocket tp = WS WS.Connection tp

instance Transport tp => Transport (WebSocket tp) where
  data TransportConfig (WebSocket tp) =
      WSServer (TransportConfig tp)
    | WSClient (TransportConfig tp) String String
  newTransport :: TransportConfig (WebSocket tp) -> IO (WebSocket tp)
newTransport (WSServer config) = do
    tp
transport <- TransportConfig tp -> IO tp
forall transport.
Transport transport =>
TransportConfig transport -> IO transport
newTransport TransportConfig tp
config
    Stream
stream <- tp -> IO Stream
forall tp. Transport tp => tp -> IO Stream
mkStream tp
transport
    PendingConnection
pendingConn <- Stream -> ConnectionOptions -> IO PendingConnection
WS.makePendingConnectionFromStream Stream
stream ConnectionOptions
WS.defaultConnectionOptions
    (Connection -> tp -> WebSocket tp)
-> tp -> Connection -> WebSocket tp
forall a b c. (a -> b -> c) -> b -> a -> c
flip Connection -> tp -> WebSocket tp
forall tp. Connection -> tp -> WebSocket tp
WS tp
transport (Connection -> WebSocket tp) -> IO Connection -> IO (WebSocket tp)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> PendingConnection -> IO Connection
WS.acceptRequest PendingConnection
pendingConn
  newTransport (WSClient config host port) = do
    tp
transport <- TransportConfig tp -> IO tp
forall transport.
Transport transport =>
TransportConfig transport -> IO transport
newTransport TransportConfig tp
config
    Stream
stream <- tp -> IO Stream
forall tp. Transport tp => tp -> IO Stream
mkStream tp
transport
    (Connection -> tp -> WebSocket tp)
-> tp -> Connection -> WebSocket tp
forall a b c. (a -> b -> c) -> b -> a -> c
flip Connection -> tp -> WebSocket tp
forall tp. Connection -> tp -> WebSocket tp
WS tp
transport (Connection -> WebSocket tp) -> IO Connection -> IO (WebSocket tp)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Stream
-> String
-> String
-> ConnectionOptions
-> Headers
-> IO Connection
WS.newClientConnection Stream
stream String
host String
port ConnectionOptions
WS.defaultConnectionOptions []

  recvData :: WebSocket tp -> Int -> IO ByteString
recvData (WS conn :: Connection
conn _) = Connection -> Int -> IO ByteString
wsRecvData Connection
conn
  sendData :: WebSocket tp -> ByteString -> IO ()
sendData (WS conn :: Connection
conn _) = Connection -> ByteString -> IO ()
wsSendData Connection
conn

  closeTransport :: WebSocket tp -> IO ()
closeTransport (WS _ tp :: tp
tp) = tp -> IO ()
forall transport. Transport transport => transport -> IO ()
closeTransport tp
tp


serverConfig :: Transport tp => TransportConfig tp -> TransportConfig (WebSocket tp)
serverConfig :: TransportConfig tp -> TransportConfig (WebSocket tp)
serverConfig = TransportConfig tp -> TransportConfig (WebSocket tp)
forall tp. TransportConfig tp -> TransportConfig (WebSocket tp)
WSServer

clientConfig :: Transport tp => TransportConfig tp -> String -> String -> TransportConfig (WebSocket tp)
clientConfig :: TransportConfig tp
-> String -> String -> TransportConfig (WebSocket tp)
clientConfig = TransportConfig tp
-> String -> String -> TransportConfig (WebSocket tp)
forall tp.
TransportConfig tp
-> String -> String -> TransportConfig (WebSocket tp)
WSClient