{-|
Module      : Gargantext.Core.AsyncUpdates.Dispatcher.WebSocket
Description : Dispatcher websocket server
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 TypeOperators #-}
    
module Gargantext.Core.AsyncUpdates.Dispatcher.WebSocket where

import Control.Concurrent.Async qualified as Async
import Control.Lens (view)
import Data.Aeson qualified as Aeson
import Data.UUID.V4 as UUID
import Gargantext.API.Admin.Auth.Types (AuthenticatedUser(_auth_user_id))
import Gargantext.API.Prelude (IsGargServer)
import Gargantext.Core.AsyncUpdates.Dispatcher.Subscriptions
import Gargantext.Core.AsyncUpdates.Dispatcher.Types
import Gargantext.Core.AsyncUpdates.Dispatcher (Dispatcher, dispatcherSubscriptions)
import Gargantext.Core.Config (HasJWTSettings(jwtSettings))
import Gargantext.Prelude
import Gargantext.System.Logging (LogLevel(DEBUG), logMsg, withLogger)
import Network.WebSockets qualified as WS
import Servant
import Servant.API.WebSocket qualified as WS (WebSocketPending)
import Servant.Auth.Server (JWTSettings, verifyJWT)
import Servant.Server.Generic (AsServerT)
import StmContainers.Set as SSet

    
newtype WSAPI mode = WSAPI {
    wsAPIServer :: mode :- "ws" :> Summary "WebSocket endpoint" :> WS.WebSocketPending
  } deriving Generic


wsServer :: ( IsGargServer env err m, HasDispatcher env Dispatcher, HasJWTSettings env ) => WSAPI (AsServerT m)
wsServer = WSAPI { wsAPIServer = streamData }
  where
    streamData :: ( IsGargServer env err m, HasDispatcher env Dispatcher, HasJWTSettings env )
               => WS.PendingConnection -> m ()
    streamData pc = do
      jwtS <- view jwtSettings
      d <- view hasDispatcher
      let subscriptions = dispatcherSubscriptions d
      key <- getWSKey pc
      c <- liftBase $ WS.acceptRequest pc
      let ws = WSKeyConnection (key, c)
      _ <- liftBase $ Async.concurrently (wsLoop jwtS subscriptions ws) (pingLoop ws)
      -- _ <- liftIO $ Async.withAsync (pure ()) (\_ -> wsLoop ws)
      pure ()

  
-- | Send a ping control frame periodically, otherwise the
-- | connection is dropped. NOTE that 'onPing' message is not
-- | supported in the JS API: either the browser supports this or
-- | not:
-- | https://stackoverflow.com/questions/10585355/sending-websocket-ping-pong-frame-from-browser
pingLoop :: WSKeyConnection -> IO ()
pingLoop ws = do
  forever $ do
    -- WS.sendDataMessage (wsConn ws) (WS.Text (Aeson.encode Ping) Nothing)
    WS.sendPing (wsConn ws) ("" :: Text)
    threadDelay $ 10 * 1000000

      
wsLoop :: JWTSettings -> SSet.Set Subscription -> WSKeyConnection -> IO a
wsLoop jwtS subscriptions ws = flip finally disconnect $ do
  withLogger () $ \ioLogger -> do
    logMsg ioLogger DEBUG "[wsLoop] connecting"
    wsLoop' CUPublic ioLogger

  where
    wsLoop' user ioLogger = do
      dm <- WS.receiveDataMessage (wsConn ws)
      
      newUser <- case dm of
        WS.Text dm' _ -> do
          case Aeson.decode dm' of
            Nothing -> do
              logMsg ioLogger DEBUG $ "[wsLoop] unknown message: " <> show dm'
              return user
            Just (WSSubscribe topic) -> do
              -- TODO Fix s_connected_user based on header
              let sub = Subscription { s_connected_user = user
                                     , s_ws_key_connection = ws
                                     , s_topic = topic }
              _ss <- insertSubscription subscriptions sub
              -- putText $ "[wsLoop] subscriptions: " <> show (showSub <$> ss)
              return user
            Just (WSUnsubscribe topic) -> do
              let sub = Subscription { s_connected_user = user
                                     , s_ws_key_connection = ws
                                     , s_topic = topic }
              _ss <- removeSubscription subscriptions sub
              -- putText $ "[wsLoop] subscriptions: " <> show (showSub <$> ss)
              return user
            Just (WSAuthorize token) -> do
              mUser <- liftBase $ verifyJWT jwtS (encodeUtf8 token)

              logMsg ioLogger DEBUG $ "[wsLoop] authorized user: " <> show mUser

              -- TODO Update my subscriptions!

              return $ fromMaybe user (CUUser . _auth_user_id <$> mUser)
            Just WSDeauthorize -> do
              -- TODO Update my subscriptions!
              pure CUPublic
        _ -> do
          logMsg ioLogger DEBUG "[wsLoop] binary ws messages not supported"
          return user
        
      wsLoop' newUser ioLogger

    disconnect = do
      withLogger () $ \ioLogger -> do
        logMsg ioLogger DEBUG "[wsLoop] disconnecting..."
        _ss <- removeSubscriptionsForWSKey subscriptions ws
        -- putText $ "[wsLoop] subscriptions: " <> show (show <$> ss)
        return ()

    
        
getWSKey :: MonadBase IO m => WS.PendingConnection -> m ByteString
getWSKey pc = do
  let reqHead = WS.pendingRequest pc
  -- WebSocket specification says that a pending request should send
  -- some unique, Sec-WebSocket-Key string. We use this to compare
  -- connections (WS.Connection doesn't implement an Eq instance).
  let mKey = head $ filter (\(k, _) -> k == "Sec-WebSocket-Key") $ WS.requestHeaders reqHead
  let key' = snd $ fromMaybe (panicTrace "Sec-WebSocket-Key not found!") mKey
  -- Unfortunately, a single browsers sends the same
  -- Sec-WebSocket-Key so we want to make that even more unique.
  uuid <- liftBase $ UUID.nextRandom
  let key = key' <> "-" <> show uuid
  liftBase $ putText $ "[getWSKey] request headers: " <> (show $ WS.requestHeaders reqHead)

  pure key
