{-|
Module      : Gargantext.Core.Worker
Description : Asynchronous worker logic
Copyright   : (c) CNRS, 2024
License     : AGPL + CECILL v3
Maintainer  : team@gargantext.org
Stability   : experimental
Portability : POSIX

-}

{-# OPTIONS_GHC -Wno-orphans      #-}  -- orphan HasNodeError IOException

{-# LANGUAGE TemplateHaskell #-}
{-# LANGUAGE UndecidableInstances #-}


module Gargantext.Core.Worker where


import Async.Worker qualified as W
import Async.Worker.Broker.Types (toA, getMessage, messageId)
import Async.Worker.Types qualified as W
import Control.Exception.Safe qualified as CES
import Control.Lens (to)
import Data.Aeson qualified as Aeson
import Data.ByteString.Lazy qualified as BSL
import Data.Text qualified as T
import Database.PostgreSQL.Simple.LargeObjects qualified as PSQL
import Gargantext.API.Admin.Auth (forgotUserPassword)
import Gargantext.API.Admin.Auth.Types (ForgotPasswordAsyncParams(..))
import Gargantext.API.Ngrams.List (postAsyncJSON)
import Gargantext.API.Node.Contact (addContact)
import Gargantext.API.Node.Corpus.Annuaire qualified as Annuaire
import Gargantext.API.Node.Corpus.New (addToCorpusWithTempFile, addToCorpusWithQuery)
import Gargantext.API.Node.DocumentUpload (documentUploadAsync, remoteImportDocuments)
import Gargantext.API.Node.DocumentsFromWriteNodes (documentsFromWriteNodes)
import Gargantext.API.Node.File (addWithFile)
import Gargantext.API.Node.FrameCalcUpload (frameCalcUploadAsync)
import Gargantext.API.Node.New (postNode')
import Gargantext.API.Node.Types (_wtf_file_oid)
import Gargantext.API.Node.Update (updateNode)
import Gargantext.API.Node.Update.Types (UpdateNodeParams(..), Granularity (..))
import Gargantext.API.Server.Named.Ngrams (tableNgramsPostChartsAsync)
import Gargantext.Core.Config (hasConfig, gc_database_config, gc_jobs, gc_worker, gc_logging)
import Gargantext.Core.Config.Types (jc_max_docs_scrapers)
import Gargantext.Core.Config.Worker (WorkerDefinition(..))
import Gargantext.Core.Notifications.CentralExchange qualified as CE
import Gargantext.Core.Notifications.CentralExchange.Types qualified as CET
import Gargantext.Core.Viz.Graph.API (graphRecompute)
import Gargantext.Core.Worker.Broker (initBrokerWithDBCreate)
import Gargantext.Core.Worker.Env
import Gargantext.Core.Worker.Jobs.Types (Job(..), getWorkerMNodeId, ImportRemoteDocumentsPayload(..), ImportRemoteTermsPayload(..))
import Gargantext.Core.Worker.PGMQTypes (BrokerMessage, HasWorkerBroker, WState)
import Gargantext.Core.Worker.Types (JobInfo(..))
import Gargantext.Database.Prelude
import Gargantext.Database.Query.Table.User (getUsersWithEmail)
import Gargantext.Prelude hiding (to)
import Gargantext.System.Logging
import Gargantext.Utils.Jobs.Error (HumanFriendlyErrorText(..))
import Gargantext.Utils.Jobs.Monad (MonadJobStatus(markStarted, markComplete, markFailed))
import System.Posix.Signals (Handler(Catch), installHandler, keyboardSignal)


initWorkerState :: HasWorkerBroker
                => WorkerEnv
                -> WorkerDefinition
                -> IO WState
initWorkerState env (WorkerDefinition { .. }) = do
  let gargConfig = env ^. hasConfig
  broker <- initBrokerWithDBCreate (gargConfig ^. gc_database_config) (gargConfig ^. gc_worker)

  pure $ W.State { broker
                 , queueName = _wdQueue
                 , name = T.unpack _wdName
                 , performAction = performAction env
                 , onMessageReceived = Just $ notifyJobStarted env
                 , onJobFinish = Just $ notifyJobFinished env
                 , onJobTimeout = Just $ notifyJobTimeout env
                 , onJobError = Just $ notifyJobFailed env
                 , onWorkerKilledSafely = Just $ notifyJobKilled env }

notifyJobStarted :: HasWorkerBroker
                 => WorkerEnv
                 -> WState
                 -> BrokerMessage
                 -> IO ()
notifyJobStarted env (W.State { name }) bm = do
  let mId = messageId bm
  let j = toA $ getMessage bm
  let job = W.job j
  withLogger (env ^. w_env_config . gc_logging) $ \ioL ->
    $(logLoc) ioL DEBUG $ T.pack $ "[notifyJobStarted] [" <> name <> " :: " <> show mId <> "] starting job: " <> show j
  let ji = JobInfo { _ji_message_id = messageId bm
                   , _ji_mNode_id = getWorkerMNodeId job }
  let jh = WorkerJobHandle { _w_job_info = ji }
  runWorkerMonad env $ markStarted 1 jh

notifyJobFinished :: HasWorkerBroker
                  => WorkerEnv
                  -> WState
                  -> BrokerMessage
                  -> IO ()
notifyJobFinished env (W.State { name }) bm = do
  let mId = messageId bm
  let j = toA $ getMessage bm
  let job = W.job j
  withLogger (env ^. w_env_config . gc_logging) $ \ioL ->
    logMsg ioL DEBUG $ "[notifyJobFinished] [" <> name <> " :: " <> show mId <> "] finished job: " <> show j
  let ji = JobInfo { _ji_message_id = messageId bm
                   , _ji_mNode_id = getWorkerMNodeId job }
  let jh = WorkerJobHandle { _w_job_info = ji }
  runWorkerMonad env $ markComplete jh

notifyJobTimeout :: HasWorkerBroker
                 => WorkerEnv
                 -> WState
                 -> BrokerMessage
                 -> IO ()
notifyJobTimeout env (W.State { name }) bm = do
  let mId = messageId bm
  let j = toA $ getMessage bm
  let job = W.job j
  withLogger (env ^. w_env_config . gc_logging) $ \ioL ->
    logMsg ioL ERROR $ "[notifyJobTimeout] [" <> name <> " :: " <> show mId <> "] job timed out: " <> show j
  let ji = JobInfo { _ji_message_id = messageId bm
                   , _ji_mNode_id = getWorkerMNodeId job }
  let jh = WorkerJobHandle { _w_job_info = ji }
  runWorkerMonad env $ markFailed (Just $ UnsafeMkHumanFriendlyErrorText "Worker job timed out!") jh

notifyJobFailed :: (HasWorkerBroker, HasCallStack)
                 => WorkerEnv
                 -> WState
                 -> BrokerMessage
                 -> SomeException
                 -> IO ()
notifyJobFailed env (W.State { name }) bm exc = do
  let mId = messageId bm
  let j = toA $ getMessage bm
  let job = W.job j
  withLogger (env ^. w_env_config . gc_logging) $ \ioL ->
    logMsg ioL ERROR $ "[notifyJobFailed] [" <> name <> " :: " <> show mId <> "] failed job: " <> show j <> "   --- ERROR: " <> show exc
  let ji = JobInfo { _ji_message_id = messageId bm
                   , _ji_mNode_id = getWorkerMNodeId job }
  let jh = WorkerJobHandle { _w_job_info = ji }
  runWorkerMonad env $ markFailed (Just $ UnsafeMkHumanFriendlyErrorText $ "Worker job failed: " <> show exc) jh

notifyJobKilled :: (HasWorkerBroker, HasCallStack)
                 => WorkerEnv
                 -> WState
                 -> Maybe BrokerMessage
                 -> IO ()
notifyJobKilled _ _ Nothing = pure ()
notifyJobKilled env (W.State { name }) (Just bm) = do
  let j = toA $ getMessage bm
  let job = W.job j
  withLogger (env ^. w_env_config . gc_logging) $ \ioL ->
    logMsg ioL ERROR $ "[notifyJobKilled] [" <> name <> "] failed job: " <> show j
  let ji = JobInfo { _ji_message_id = messageId bm
                   , _ji_mNode_id = getWorkerMNodeId job }
  let jh = WorkerJobHandle { _w_job_info = ji }
  runWorkerMonad env $ markFailed (Just $ UnsafeMkHumanFriendlyErrorText $ "Worker '" <> T.pack name <> "' was killed") jh


-- | Spawn a worker with PGMQ broker
withPGMQWorker :: HasWorkerBroker
                => WorkerEnv
                -> WorkerDefinition
                -> (Async () -> WState -> IO ())
                -> IO ()
withPGMQWorker env wd cb = do
  state' <- initWorkerState env wd

  withAsync (W.run state') (\a -> cb a state')

withPGMQWorkerSingle :: HasWorkerBroker
                     => WorkerEnv
                     -> WorkerDefinition
                     -> (Async () -> WState -> IO ())
                     -> IO ()
withPGMQWorkerSingle env wd cb = do
  state' <- initWorkerState env wd

  withAsync (W.runSingle state') (\a -> cb a state')


withPGMQWorkerCtrlC :: HasWorkerBroker
                    => WorkerEnv
                    -> WorkerDefinition
                    -> (Async () -> WState -> IO ())
                    -> IO ()
withPGMQWorkerCtrlC env wd cb = do
  withPGMQWorker env wd $ \a state' -> do
    let tid = asyncThreadId a
    _ <- installHandler keyboardSignal (Catch (throwTo tid W.KillWorkerSafely)) Nothing
    cb a state'

withPGMQWorkerSingleCtrlC :: HasWorkerBroker
                          => WorkerEnv
                          -> WorkerDefinition
                          -> (Async () -> WState -> IO ())
                          -> IO ()
withPGMQWorkerSingleCtrlC env wd cb = do
  withPGMQWorkerSingle env wd $ \a state' -> do
    let tid = asyncThreadId a
    _ <- installHandler keyboardSignal (Catch (throwTo tid W.KillWorkerSafely)) Nothing
    cb a state'


-- | How the worker should process jobs
performAction :: HasWorkerBroker
              => WorkerEnv
              -> WState
              -> BrokerMessage
              -> IO ()
performAction env _state bm = do
  let job' = toA $ getMessage bm
  let job = W.job job'
  let ji = JobInfo { _ji_message_id = messageId bm
                   , _ji_mNode_id = getWorkerMNodeId job }
  let jh = WorkerJobHandle { _w_job_info = ji }

  case job of
    Ping -> runWorkerMonad env $ do
      $(logLocM) DEBUG "[performAction] ping"
      liftIO $ CE.notify (env ^. (to _w_env_config)) CET.Ping

    -- | flow action for a single contact
    AddContact { .. } -> runWorkerMonad env $ do
      $(logLocM) DEBUG $ "[performAction] add contact"
      addContact _ac_user _ac_node_id _ac_args jh
      
    -- | Uses temporary file to add documents into corpus
    AddCorpusTempFileAsync { .. } -> runWorkerMonad env $ do
      $(logLocM) DEBUG "[performAction] add to corpus with temporary file"
      CES.finally (addToCorpusWithTempFile _actf_user _actf_cid _actf_args jh)
          (removeLargeObject $ _wtf_file_oid _actf_args)
      
    -- | Perform external API search query and index documents in corpus
    AddCorpusWithQuery { .. } -> runWorkerMonad env $ do
      $(logLocM) DEBUG "[performAction] add corpus with query"
      let limit = Just $ fromIntegral $ env ^. hasConfig . gc_jobs . jc_max_docs_scrapers
      addToCorpusWithQuery _acq_user _acq_cid _acq_args limit jh

    -- | Add to annuaire, from given file (not implemented yet)
    AddToAnnuaireWithForm { .. } -> runWorkerMonad env $ do
      $(logLocM) DEBUG "[performAction] add to annuaire with form"
      Annuaire.addToAnnuaireWithForm _aawf_annuaire_id _aawf_args jh

    -- | Saves file to 'data_filepath' (in TOML), adds this file as a node
    AddWithFile { .. } -> runWorkerMonad env $ do
      $(logLocM) DEBUG "[performAction] add with file"
      addWithFile _awf_authenticatedUser _awf_node_id _awf_args jh

    -- | For given corpus, get write nodes contents and create documents from it
    DocumentsFromWriteNodes { .. } -> runWorkerMonad env $ do
      $(logLocM) DEBUG "[performAction] documents from write nodes"
      documentsFromWriteNodes _dfwn_authenticatedUser _dfwn_node_id _dfwn_args jh

    -- | Forgot password task
    ForgotPasswordAsync { _fpa_args = ForgotPasswordAsyncParams { email } } -> runWorkerMonad env $ do
      $(logLocM) DEBUG $ "[performAction] forgot password: " <> email
      us <- runDBQuery $ getUsersWithEmail (T.toLower email)
      case us of
        [u] -> forgotUserPassword u
        _ -> pure ()

    -- | Add given calc frame into corpus (internall, as a TSV file upload)
    FrameCalcUpload { .. } -> runWorkerMonad env $ do
      $(logLocM) DEBUG "[performAction] frame calc upload"
      frameCalcUploadAsync _fca_authenticatedUser _fca_node_id _fca_args jh

    -- | Process uploaded JSON file
    JSONPost { .. } -> runWorkerMonad env $ do
      $(logLocM) DEBUG $ "[performAction] json post"
      CES.finally (do
                      _jp_ngrams_list' <- readLargeObject (PSQL.Oid $ fromIntegral _jp_ngrams_oid)
                      case Aeson.eitherDecode (BSL.fromStrict _jp_ngrams_list') of
                        Left err -> CES.throwString err
                        Right _jp_ngrams_list -> void $ postAsyncJSON _jp_list_id _jp_ngrams_list jh)
        (removeLargeObject _jp_ngrams_oid)

    -- | Task for updating metrics charts
    NgramsPostCharts { .. } -> runWorkerMonad env $ do
      $(logLocM) DEBUG $ "[performAction] ngrams post charts"
      void $ tableNgramsPostChartsAsync _npc_args jh

    -- | Creates node of given type
    PostNodeAsync { .. } -> runWorkerMonad env $ do
      $(logLocM) DEBUG $ "[performAction] post node async"
      void $ postNode' _pna_authenticatedUser _pna_node_id _pna_args

    -- | Recompute graph (for sigmajs)
    RecomputeGraph { .. } -> runWorkerMonad env $ do
      $(logLocM) DEBUG $ "[performAction] recompute graph"
      void $ graphRecompute _rg_node_id jh

    -- | Updates a node (which triggers graph)
    UpdateNode { .. } -> runWorkerMonad env $ do
      $(logLocM) DEBUG $ "[performAction] update node"
      void $ updateNode _un_node_id _un_args jh

    -- | Upload a document
    UploadDocument { .. } -> runWorkerMonad env $ do
      $(logLocM) DEBUG $ "[performAction] upload document"
      void $ documentUploadAsync _ud_node_id _ud_args jh

    -- | Remotely import documents
    ImportRemoteTerms (ImportRemoteTermsPayload list_id ngrams_list)
      -> runWorkerMonad env $ do
           $(logLocM) DEBUG $ "[performAction] import remote terms"
           void $ postAsyncJSON list_id ngrams_list jh
           -- Trigger an 'UpdateNode' job to update the score(s)
           $(logLocM) DEBUG $ "Updating node scores for corpus node " <> T.pack (show list_id)
           void $ updateNode list_id (UpdateNodeParamsTexts Both) jh
           $(logLocM) DEBUG $ "Done updating node scores for corpus node " <> T.pack (show list_id)

    -- | Remotely import documents
    ImportRemoteDocuments (ImportRemoteDocumentsPayload loggedInUser parentId corpusId docs workSplit)
      -> runWorkerMonad env $ do
           $(logLocM) DEBUG $ "[performAction] import remote documents"
           void $ remoteImportDocuments loggedInUser parentId corpusId workSplit docs
