{-|
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.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 Nanomsg (Pull(..), bind, recv, withSocket)
import Network.WebSockets qualified as WS
import StmContainers.Set qualified as SSet
import Gargantext.Core.Config
import Gargantext.System.Logging

{-

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 500_000 throttleTChan (sendDataMessageThrottled 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
    NotificationsConfig { _nc_dispatcher_bind } = 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 :: LogConfig -> (WS.Connection, WS.DataMessage) -> IO ()
sendDataMessageThrottled log_cfg (conn, msg) = do
  withLogger log_cfg $ \ioL ->
    logMsg ioL DEBUG $ "[sendDataMessageThrottled] dispatching notification: " <> show msg
  WS.sendDataMessage conn msg


-- | 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
