jitsi-roomsv2/backend/src/WebSocket.hs

{-# LANGUAGE DerivingVia #-}
{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE UndecidableInstances #-}

module WebSocket (runWebSocketServer) where

import BroadcastUserData (broadcastUserData)
import ClassyPrelude
import Data.UUID (UUID)
import Data.UUID.V4 (nextRandom)
import Network.WebSockets qualified as WS
import Types.AppTypes
  ( Env (..),
    HasConnectedClientState (getConnectedClientState),
  )
import Types.ConnectionState (Client (..), ConnectedClients)
import Types.RoomsState (HasRoomsState (..))

data WSEnv = WSEnv
  { appEnv :: Env,
    connection :: WS.Connection,
    clientId :: UUID
  }

instance HasConnectedClientState WSEnv where
  getConnectedClientState = connectedClientsState . appEnv

instance HasRoomsState WSEnv where
  getRoomsState = roomsState . appEnv

class HasWSEnv a where
  getConn :: a -> WS.Connection
  getClientId :: a -> UUID

instance HasWSEnv WSEnv where
  getConn = connection
  getClientId = clientId

addWSClient ::
  ( HasConnectedClientState env,
    MonadReader env m,
    MonadUnliftIO m
  ) =>
  Client ->
  m ()
addWSClient client = do
  state <- getConnectedClientState <$> ask
  modifyMVar_ state $ \s -> do
    let s' = addClient client s
    return s'

addClient :: Client -> ConnectedClients -> ConnectedClients
addClient client clients = client : clients

disconnectWsClient ::
  ( HasWSEnv env,
    HasConnectedClientState env,
    MonadReader env m,
    MonadUnliftIO m
  ) =>
  m ()
disconnectWsClient = do
  clientId <- getClientId <$> ask
  state <- getConnectedClientState <$> ask
  modifyMVar_ state $ \s ->
    let s' = removeClient clientId s
     in return s'

joinRoom ::
  ( HasWSEnv env,
    HasConnectedClientState env,
    MonadReader env m,
    MonadUnliftIO m
  ) =>
  m ()
joinRoom = do
  clientId <- getClientId <$> ask
  state <- getConnectedClientState <$> ask
  liftIO $ putStrLn "joinedRoom"
  modifyMVar_ state $ \s ->
    let s' = updateClient clientId (\c -> c {joinedRoom = True}) s
     in return s'

removeClient :: UUID -> ConnectedClients -> ConnectedClients
removeClient toRemove = filter ((/= toRemove) . uuid)

updateClient :: UUID -> (Client -> Client) -> ConnectedClients -> ConnectedClients
updateClient toUpdate f allClients = do
  client <- allClients
  if toUpdate == uuid client
    then return $ f client
    else return client

runWebSocketServer ::
  ( MonadIO m,
    MonadReader Env m
  ) =>
  m ()
runWebSocketServer = do
  putStrLn "Websocket up at 127.0.0.1:9160"
  wsApp' <- runWSApp
  liftIO $ WS.runServer "127.0.0.1" 9160 wsApp'

runWSApp ::
  ( MonadIO m,
    MonadReader Env m
  ) =>
  m WS.ServerApp
runWSApp = do
  env <- ask
  return
    ( \pending -> do
        conn <- WS.acceptRequest pending
        uuid <- nextRandom
        let wsEnv = WSEnv {appEnv = env, connection = conn, clientId = uuid}
        WS.withPingThread conn 30 (return ()) $ unWSApp wsApp wsEnv
    )

newClient ::
  ( MonadIO m,
    HasWSEnv env,
    MonadReader env m
  ) =>
  Text ->
  m Client
newClient name = do
  env <- ask
  return $ Client {uuid = getClientId env, name = name, conn = getConn env, joinedRoom = False}

newtype WSApp env a = WSApp {unWSApp :: env -> IO a}
  deriving
    ( Functor,
      Applicative,
      Monad,
      MonadReader env,
      MonadIO,
      MonadUnliftIO
    )
    via ReaderT env IO

wsApp ::
  ( HasWSEnv env,
    HasConnectedClientState env,
    HasRoomsState env,
    MonadReader env m,
    MonadUnliftIO m
  ) =>
  m ()
wsApp = do
  msg <- getMessage
  putStrLn msg
  client <- newClient msg
  addWSClient client
  broadcastUserData
  flip finally disconnectWsClient $ forever $ do
    currentMsg <- getMessage
    joinRoom
    broadcastUserData
    putStrLn currentMsg

getMessage ::
  ( HasWSEnv env,
    MonadIO m,
    MonadReader env m
  ) =>
  m Text
getMessage = do
  conn' <- getConn <$> ask
  liftIO $ WS.receiveData conn'