{-|
Module      : Test.Utils.Jobs.Types
Description :
Copyright   : (c) CNRS, 2017-Present
License     : AGPL + CECILL v3
Maintainer  : team@gargantext.org
Stability   : experimental
Portability : POSIX
-}


module Test.Utils.Jobs.Types
  ( TestJobEnv(..)
  , initTestJobEnv
  , initTestWorkerState
  )
where

import Async.Worker.Broker.Types (toA, getMessage)
import Async.Worker.Types qualified as WT
import Control.Concurrent.STM
import Data.Text qualified as T
import Gargantext.Core.Config (hasConfig, gc_database_config, gc_worker)
import Gargantext.Core.Config.Worker (WorkerDefinition(..))
import Gargantext.Core.Worker (performAction)
import Gargantext.Core.Worker.Broker (initBrokerWithDBCreate)
import Gargantext.Core.Worker.Env (WorkerEnv(..))
import Gargantext.Core.Worker.Jobs.Types (Job(..))
import Gargantext.Core.Worker.PGMQTypes (HasWorkerBroker, BrokerMessage, WState)
import Gargantext.Prelude


data TestJobEnv =
  TestJobEnv { inProgress :: Maybe Job
             , done       :: [Job]
             , failed     :: [(Job, SomeException)]
             , killed     :: [Job]
             , timedOut   :: [Job]
             }
initTestJobEnv :: TestJobEnv
initTestJobEnv =
  TestJobEnv { inProgress = Nothing
             , done       = []
             , failed     = []
             , killed     = []
             , timedOut   = [] }

-- | Test worker state. Normally, the message notifications go through
-- the dispatcher system. Here we make a short-cut and just use a
-- TVar to store the processes worker jobs.
-- Job progress, however, is sent via the notifications mechanism,
-- because the worker itself doesn't implement it.
initTestWorkerState :: HasWorkerBroker
                    => TVar TestJobEnv
                    -> WorkerEnv
                    -> WorkerDefinition
                    -> IO WState
initTestWorkerState jobTVar env (WorkerDefinition { .. }) = do
  let gargConfig = env ^. hasConfig
  broker <- initBrokerWithDBCreate (gargConfig ^. gc_database_config) (gargConfig ^. gc_worker)

  pure $ WT.State { broker
                  , queueName = _wdQueue
                  , name = T.unpack _wdName
                  , performAction = performAction env
                  , onMessageReceived = Just $ onJobStarted jobTVar env
                  , onJobFinish = Just $ onJobFinished jobTVar env
                  , onJobTimeout = Just $ onJobTimeout jobTVar env
                  , onJobError = Just $ onJobError jobTVar env
                  , onWorkerKilledSafely = Just $ onWorkerKilled jobTVar env }

onJobStarted :: HasWorkerBroker
             => TVar TestJobEnv
             -> WorkerEnv
             -> WState
             -> BrokerMessage
             -> IO ()
onJobStarted jobTVar _env _state bm = do
  let j = toA $ getMessage bm
  let job = WT.job j
  atomically $ modifyTVar jobTVar $ \testJobEnv -> do
    testJobEnv { inProgress = Just job }

onJobFinished :: HasWorkerBroker
              => TVar TestJobEnv
              -> WorkerEnv
              -> WState
              -> BrokerMessage
              -> IO ()
onJobFinished jobTVar _env _state bm = do
  let j = toA $ getMessage bm
  let job = WT.job j
  atomically $ modifyTVar jobTVar $ \testJobEnv -> do
    testJobEnv { inProgress = Nothing
               , done = done testJobEnv ++ [job] }

onJobTimeout :: HasWorkerBroker
             => TVar TestJobEnv
             -> WorkerEnv
             -> WState
             -> BrokerMessage
             -> IO ()
onJobTimeout jobTVar _env _state bm = do
  let j = toA $ getMessage bm
  let job = WT.job j
  atomically $ modifyTVar jobTVar $ \testJobEnv -> do
    testJobEnv { inProgress = Nothing
               , timedOut = timedOut testJobEnv ++ [job] }

onJobError :: (HasWorkerBroker, HasCallStack)
           => TVar TestJobEnv
           -> WorkerEnv
           -> WState
           -> BrokerMessage
           -> SomeException
           -> IO ()
onJobError jobTVar _env _state bm exc = do
  let j = toA $ getMessage bm
  let job = WT.job j
  atomically $ modifyTVar jobTVar $ \testJobEnv -> do
    testJobEnv { inProgress = Nothing
               , failed = failed testJobEnv ++ [(job, exc)] }

onWorkerKilled :: (HasWorkerBroker, HasCallStack)
               => TVar TestJobEnv
               -> WorkerEnv
               -> WState
               -> Maybe BrokerMessage
               -> IO ()
onWorkerKilled _jobTVar _env _state Nothing = pure ()
onWorkerKilled jobTVar _env _state (Just  bm) = do
  let j = toA $ getMessage bm
  let job = WT.job j
  atomically $ modifyTVar jobTVar $ \testJobEnv -> do
    testJobEnv { inProgress = Nothing
               , killed = killed testJobEnv ++ [job] }
