{-|
Module      : Gargantext.Core.Notifications.Dispatcher
Description : Dispatcher (handles websocket connections, accepts message from central exchange)
Copyright   : (c) CNRS, 2017-Present
License     : AGPL + CECILL v3
Maintainer  : team@gargantext.org
Stability   : experimental
Portability : POSIX

https://gitlab.iscpif.fr/gargantext/haskell-gargantext/issues/341

Docs:
https://dev.sub.gargantext.org/#/share/Notes/187918

-}

{-# LANGUAGE TemplateHaskell #-}

module Gargantext.Core.Notifications.Dispatcher (
    Dispatcher -- opaque
  , withDispatcher

  -- * Querying a dispatcher
  , dispatcherSubscriptions
  ) where

import Control.Concurrent.Async qualified as Async
import Control.Concurrent.STM.TChan qualified as TChan
import Control.Concurrent.Throttle (throttle)
import Data.Aeson qualified as Aeson
import Data.ByteString.Lazy qualified as BSL
import Data.Text qualified as T
import DeferredFolds.UnfoldlM qualified as UnfoldlM
import Gargantext.Core.Config
    ( GargConfig, LogConfig, gc_logging, gc_notifications_config )
import Gargantext.Core.Config.Types (NotificationsConfig(..))
import Gargantext.Core.Notifications.CentralExchange.Types qualified as CETypes
import Gargantext.Core.Notifications.Dispatcher.Types
import Gargantext.Core.Worker.Types (JobInfo(..))
import Gargantext.Prelude
import Gargantext.System.Logging
    ( HasLogger(logMsg), LogLevel(..), withLogger, logLoc )
import Nanomsg (Pull(..), bind, recv, withSocket)
import Network.WebSockets qualified as WS
import StmContainers.Set qualified as SSet
import System.Timeout (timeout)

{-

Dispatcher is a service, which provides couple of functionalities:
- handles WebSocket connections and manages them
- accepts messages from central exchange
- dispatches these messages to connected users

-}

data Dispatcher =
 Dispatcher { d_subscriptions :: SSet.Set Subscription
            }

dispatcherSubscriptions :: Dispatcher -> SSet.Set Subscription
dispatcherSubscriptions = d_subscriptions

withDispatcher :: GargConfig -> (Dispatcher -> IO a) -> IO a
withDispatcher cfg cb = do
  subscriptions <- SSet.newIO

  Async.withAsync (dispatcherListener cfg subscriptions) $ \_a -> do
    let dispatcher = Dispatcher { d_subscriptions = subscriptions }
    cb dispatcher


-- | This is a nanomsg socket listener. We want to read the messages
-- | as fast as possible and then process them gradually in a separate
-- | thread.
dispatcherListener :: GargConfig -> SSet.Set Subscription -> IO ()
dispatcherListener config subscriptions = do
  withSocket Pull $ \s -> do
    withLogger log_cfg $ \ioLogger -> do
      logMsg ioLogger DEBUG $ "[dispatcherListener] binding to " <> T.unpack _nc_dispatcher_bind
    _ <- bind s $ T.unpack _nc_dispatcher_bind

    tChan <- TChan.newTChanIO

    throttleTChan <- TChan.newTChanIO

    -- NOTE I'm not sure that we need more than 1 worker here, but in
    -- theory, the worker can perform things like user authentication,
    -- DB queries etc so it can be slow sometimes.
    Async.withAsync (throttle (_nc_dispatcher_throttle_ms * 1000) throttleTChan
                     (sendDataMessageThrottled nc log_cfg)) $ \_ -> do
      void $ Async.concurrently (Async.replicateConcurrently 5 $ worker tChan throttleTChan) $ do
        forever $ do
          -- putText "[dispatcher_listener] receiving"
          r <- recv s
          -- C.putStrLn $ "[dispatcher_listener] " <> r
          atomically $ TChan.writeTChan tChan r
  where
    nc@NotificationsConfig { _nc_dispatcher_bind, _nc_dispatcher_throttle_ms } = config ^. gc_notifications_config
    log_cfg = config ^. gc_logging
    worker tChan throttleTChan = withLogger log_cfg $ \ioL -> do
      tId <- myThreadId

      forever $ do
        r <- atomically $ TChan.readTChan tChan
        $(logLoc) ioL DEBUG $ "[" <> show tId <> "] received a message: " <> decodeUtf8 r

        case Aeson.decode (BSL.fromStrict r) of
          Nothing ->
              logMsg ioL DEBUG "[dispatcher_listener] unknown message from central exchange"
          Just ceMessage -> do
            logMsg ioL DEBUG $ "[dispatcher_listener] received " <> show ceMessage
            -- subs <- atomically $ readTVar subscriptions
            filteredSubs <- atomically $ do
              let subs' = UnfoldlM.filter (pure . ceMessageSubPred ceMessage) $ SSet.unfoldlM subscriptions
              UnfoldlM.foldlM' (\acc sub -> pure $ acc <> [sub]) [] subs'
            -- NOTE This isn't safe: we atomically fetch subscriptions,
            -- then send notifications one by one. In the meantime, a
            -- subscription could end or new ones could appear (but is
            -- this really a problem? I new subscription comes up, then
            -- probably they already fetch new tree anyways, and if old
            -- one drops in the meantime, it won't listen to what we
            -- send...)
            -- let filteredSubs = filterCEMessageSubs ceMessage subs
            mapM_ (sendNotification throttleTChan ceMessage) filteredSubs

-- | When processing tasks such as Flow, we can generate quite a few
-- notifications in a short time. We want to limit this with throttle
-- tchan.
sendNotification :: TChan.TChan ((ByteString, Topic), (WS.Connection, WS.DataMessage))
                 -> CETypes.CEMessage
                 -> Subscription
                 -> IO ()
sendNotification throttleTChan ceMessage sub = do
  let ws = s_ws_key_connection sub
  -- 'topic' is where the client subscribed, ceMessage is server's
  -- message to a client
  let topic = s_topic sub
  let mNotification =
        -- | OK so given a websocket subscription and a central
        -- exchange message - decide whether to send this message via
        -- that socket or not
        case (topic, ceMessage) of
          -- (UpdateWorkerProgress jobInfo, CETypes.UpdateWorkerProgress jobInfo' nodeId jobLog) -> do
          (UpdateWorkerProgress jobInfo, CETypes.UpdateWorkerProgress jobInfo' jobLog) -> do
            if jobInfo == jobInfo'
              -- then Just $ NUpdateWorkerProgress jobInfo nodeId (MJobLog jobLog)
              then Just $ NUpdateWorkerProgress jobInfo jobLog  -- (MJobLog jobLog)
              else Nothing
          (UpdateTree nodeId, CETypes.UpdateWorkerProgress jobInfo jobLog) -> do
            if Just nodeId == _ji_mNode_id jobInfo
              then Just $ NUpdateWorkerProgress jobInfo jobLog  -- (MJobLog jobLog)
              else Nothing
          (UpdateTree nodeId, CETypes.UpdateTreeFirstLevel nodeId') ->
            if nodeId == nodeId'
            then Just $ NUpdateTree nodeId
            else Nothing
          (Ping, CETypes.Ping) ->
            Just NPing
          _ -> Nothing

  case mNotification of
    Nothing -> pure ()
    Just notification -> do
      let id' = (wsKey ws, topic)
      atomically $ do
        TChan.writeTChan throttleTChan (id', (wsConn ws, WS.Text (Aeson.encode notification) Nothing))

-- | The "true" message sending to websocket. After it was withheld
-- for a while (for throttling), it is finally sent here
sendDataMessageThrottled :: NotificationsConfig -> LogConfig -> (WS.Connection, WS.DataMessage) -> IO ()
sendDataMessageThrottled (NotificationsConfig { _nc_dispatcher_send_timeout_ms }) log_cfg (conn, msg) = do
  withLogger log_cfg $ \ioL -> do
    logMsg ioL DEBUG $ "[sendDataMessageThrottled] dispatching notification: " <> show msg
    -- | We need a timeout here for the following reason:
    --   when a message is sent and the user disconnects the WS
    --   connection (e.g. refreshes the page), it seems that this message sending hangs.
    --   We don't want to block the thread indefinitely.
    timeoutKickedIn <- timeout (_nc_dispatcher_send_timeout_ms * 1000) $ WS.sendDataMessage conn msg
    case timeoutKickedIn of
      Nothing ->
          $(logLoc) ioL ERROR $ "[sendMessageThrottled] couldn't send msg in timely fashion."
      Just _ -> pure ()


-- | Custom filtering of list of Subscriptions based on
--   'CETypes.CEMessage'.
-- For example, we can add CEMessage.Broadcast to propagate a
-- notification to all connections.
_filterCEMessageSubs :: CETypes.CEMessage -> [Subscription] -> [Subscription]
_filterCEMessageSubs ceMessage subscriptions = filter (ceMessageSubPred ceMessage) subscriptions

-- | Predicate, whether 'Subscription' matches given
-- 'CETypes.CEMessage' (i.e. should given 'Subscription' be informed
-- of this message).
ceMessageSubPred :: CETypes.CEMessage -> Subscription -> Bool
-- ceMessageSubPred (CETypes.UpdateWorkerProgress ji nodeId _jl) (Subscription { s_topic }) =
ceMessageSubPred (CETypes.UpdateWorkerProgress ji _jl) (Subscription { s_topic }) =
     s_topic == UpdateWorkerProgress ji
  || Just s_topic == (UpdateTree <$> _ji_mNode_id ji)
ceMessageSubPred (CETypes.UpdateTreeFirstLevel nodeId) (Subscription { s_topic }) =
  s_topic == UpdateTree nodeId
ceMessageSubPred CETypes.Ping (Subscription { s_topic }) =
  s_topic == Ping
