WebSocket.hs

module Ema.Server.WebSocket where

import Control.Monad.Logger
import Data.LVar (LVar)
import Data.LVar qualified as LVar
import Ema.Asset (
  Asset (AssetGenerated, AssetStatic),
  Format (Html, Other),
 )
import Ema.Route.Class (IsRoute (RouteModel, routePrism))
import Ema.Route.Prism (
  fromPrism_,
 )
import Ema.Server.Common
import Ema.Server.WebSocket.Options (EmaWsHandler (..))
import Ema.Site (EmaStaticSite)
import Network.WebSockets (ConnectionException)
import Network.WebSockets qualified as WS
import Optics.Core (review)
import UnliftIO.Async (race)
import UnliftIO.Exception (try)

wsApp ::
  forall r.
  (Eq r, Show r, IsRoute r, EmaStaticSite r) =>
  (Loc -> LogSource -> LogLevel -> LogStr -> IO ()) ->
  LVar (RouteModel r) ->
  EmaWsHandler r ->
  WS.PendingConnection ->
  IO ()
wsApp logger model emaWsHandler pendingConn = do
  conn :: WS.Connection <- WS.acceptRequest pendingConn
  WS.withPingThread conn 30 pass . flip runLoggingT logger $ do
    let log lvl (s :: Text) =
          logWithoutLoc "ema.ws" lvl s
    log LevelInfo "Connected"
    let wsHandler = unEmaWsHandler emaWsHandler conn
        sendRouteHtmlToClient path s = do
          decodeUrlRoute @r s path & \case
            Left err -> do
              log LevelError $ badRouteEncodingMsg err
              liftIO $ WS.sendTextData conn $ emaErrorHtmlResponse $ badRouteEncodingMsg err
            Right Nothing ->
              liftIO $ WS.sendTextData conn $ emaErrorHtmlResponse decodeRouteNothingMsg
            Right (Just r) -> do
              renderCatchingErrors s r >>= \case
                AssetGenerated Html html ->
                  liftIO $ WS.sendTextData conn $ html <> toLazy wsClientHtml
                -- HACK: We expect the websocket client should check for REDIRECT prefix.
                -- Not bothering with JSON response to avoid having to JSON parse every HTML dump.
                AssetStatic _staticPath ->
                  liftIO $ WS.sendTextData conn $ "REDIRECT " <> toText (review (fromPrism_ $ routePrism s) r)
                AssetGenerated Other _s ->
                  liftIO $ WS.sendTextData conn $ "REDIRECT " <> toText (review (fromPrism_ $ routePrism s) r)
              log LevelDebug $ " ~~> " <> show r
        -- @mWatchingRoute@ is the route currently being watched.
        loop mWatchingRoute = do
          -- Listen *until* either we get a new value, or the client requests
          -- to switch to a new route.
          currentModel <- LVar.get model
          race (LVar.listenNext model) (wsHandler currentModel) >>= \case
            Left newModel -> do
              -- The page the user is currently viewing has changed. Send
              -- the new HTML to them.
              sendRouteHtmlToClient mWatchingRoute newModel
              loop mWatchingRoute
            Right mNextRoute -> do
              -- The user clicked on a route link; send them the HTML for
              -- that route this time, ignoring what we are watching
              -- currently (we expect the user to initiate a watch route
              -- request immediately following this).
              sendRouteHtmlToClient mNextRoute =<< LVar.get model
              loop mNextRoute
    -- Wait for the client to send the first request with the initial route.
    mInitialRoute <- wsHandler =<< LVar.get model
    try (loop mInitialRoute) >>= \case
      Right () -> pass
      Left (connExc :: ConnectionException) -> do
        case connExc of
          WS.CloseRequest _ (decodeUtf8 -> reason) ->
            log LevelInfo $ "Closing websocket connection (reason: " <> reason <> ")"
          _ ->
            log LevelError $ "Websocket error: " <> show connExc