module Jobs
  ( module Types
  , module Jobs.Queue
  , handleJobs
  , stopJobCtl
  ) where

import Import
import Utils.Lens

import Jobs.Types as Types hiding (JobCtl(JobCtlQueue))
import Jobs.Types (JobCtl(JobCtlQueue))
import Jobs.Queue
import Jobs.Crontab

import qualified Data.Conduit.List as C

import qualified Data.Text.Lazy as LT

import Data.Aeson (fromJSON, toJSON)
import qualified Data.Aeson as Aeson
import qualified Data.Aeson.Types as Aeson
import Database.Persist.Sql (fromSqlKey)

import Data.Semigroup (Max(..))

import Utils.Sql

import Control.Monad.Random (evalRand, mkStdGen, getRandomR, uniformMay)

import Cron
import qualified Data.HashMap.Strict as HashMap
import Data.HashMap.Strict (HashMap)
import qualified Data.Set as Set

import qualified Data.List.NonEmpty as NonEmpty

import qualified Data.Map.Strict as Map
import Data.Map.Strict ((!))

import Data.Foldable (foldrM)

import Control.Monad.Trans.Reader (mapReaderT)
import Control.Monad.Trans.Writer (execWriterT)
import Control.Monad.Trans.RWS.Lazy (RWST, mapRWST, evalRWST)
import qualified Control.Monad.State.Class as State
import Control.Monad.Writer.Class (MonadWriter(..))
import Control.Monad.Reader.Class (MonadReader(..))
import Control.Monad.Trans.Resource (MonadResourceBase, runResourceT)
import Control.Monad.Trans.Maybe (MaybeT(..))
import Control.Monad.Trans.Cont (ContT(..), callCC)
import Control.Monad.Random.Lazy (evalRandTIO, mapRandT)
import Control.Monad.Logger

import Data.Time.Zones

import Control.Concurrent.STM (retry)
import Control.Concurrent.STM.Delay


import Jobs.Handler.SendNotification
import Jobs.Handler.SendTestEmail
import Jobs.Handler.QueueNotification
import Jobs.Handler.HelpRequest
import Jobs.Handler.SetLogSettings
import Jobs.Handler.DistributeCorrections
import Jobs.Handler.SendCourseCommunication
import Jobs.Handler.Invitation

import Jobs.HealthReport

  
data JobQueueException = JInvalid QueuedJobId QueuedJob
                       | JLocked QueuedJobId InstanceId UTCTime
                       | JNonexistant QueuedJobId
  deriving (Read, Show, Eq, Generic, Typeable)

instance Exception JobQueueException


handleJobs :: ( MonadResource m
              , MonadLoggerIO m
              )
           => UniWorX -> m ()
-- | Spawn a set of workers that read control commands from `appJobCtl` and address them as they come in
--
-- Uses `unsafeHandler`, as per documentation all HTTP-related fields of state/environment are meaningless placeholders.
-- Handling commands in `HandlerT` provides us with the facilities to render urls, unifies logging, provides a value of the foundation type, ...
handleJobs foundation@UniWorX{..}
  | foundation ^. _appJobWorkers == 0 = return ()
  | otherwise = do
      logger <- askLoggerIO
      let runInIO = flip runLoggingT logger . runResourceT

      jobPoolManager <- allocateLinkedAsync . runInIO $ manageJobPool foundation

      jobCron <- allocateLinkedAsync . runInIO $ manageCrontab foundation

      let jobWorkers = Map.empty
          jobWorkerName = const $ error "Unknown worker"
      jobCrontab <- liftIO $ newTVarIO HashMap.empty
      jobConfirm <- liftIO $ newTVarIO HashMap.empty
      jobShutdown <- liftIO newEmptyTMVarIO
      atomically $ putTMVar appJobState JobState
        { jobContext = JobContext{..}
        , ..
        }

manageJobPool, manageCrontab :: forall m.
                                ( MonadResource m
                                , MonadLogger m
                                )
                             => UniWorX -> m ()
manageCrontab foundation@UniWorX{..} = do
  context <- atomically . fmap jobContext $ readTMVar appJobState
  let awaitTermination = atomically $ do
        shouldTerminate <- readTMVar appJobState >>= fmap not . isEmptyTMVar . jobShutdown
        guard shouldTerminate
  liftIO . race_ awaitTermination . unsafeHandler foundation . void $ do
    atomically . assertM_ (not . Map.null . jobWorkers) $ readTMVar appJobState
    runReaderT ?? foundation $
      writeJobCtlBlock JobCtlDetermineCrontab
    evalRWST (forever execCrontab) context HashMap.empty


manageJobPool foundation@UniWorX{..}
  = flip runContT return . forever . join . atomically $ asum
      [ spawnMissingWorkers
      , reapDeadWorkers
      , terminateGracefully
      ]
  where
    num :: Int
    num = fromIntegral $ foundation ^. _appJobWorkers

    spawnMissingWorkers, reapDeadWorkers, terminateGracefully :: STM (ContT () m ())
    spawnMissingWorkers = do
      oldState <- takeTMVar appJobState
      let missing = num - Map.size (jobWorkers oldState)
      guard $ missing > 0
      return $ do
        $logDebugS "manageJobPool" [st|Spawning #{missing} workers|]
        endo <- execWriterT . replicateM_ missing $ do
          workerId <- newWorkerId
          let logIdent = mkLogIdent workerId
          (bChan, chan) <- atomically $ newBroadcastTChan >>= (\c -> (c, ) <$> dupTChan c)
          let
            streamChan = join . atomically $ do
              shouldTerminate <- readTMVar appJobState >>= fmap not . isEmptyTMVar . jobShutdown
              if
                | shouldTerminate ->
                    return $ return ()
                | otherwise -> do
                    nextVal <- readTChan chan
                    return $ yield nextVal >> streamChan
            runWorker = unsafeHandler foundation . flip runReaderT (jobContext oldState) $ do
              $logInfoS logIdent "Started"
              runConduit $ streamChan .| handleJobs' workerId
              $logInfoS logIdent "Stopped"
          worker <- allocateLinkedAsync runWorker

          tell . Endo $ \cSt -> cSt
            { jobWorkers = Map.insert worker bChan $ jobWorkers cSt
            , jobWorkerName = \a -> bool (jobWorkerName cSt a) workerId $ a == worker
            }
        atomically . putTMVar appJobState $ endo `appEndo` oldState

    reapDeadWorkers = do
      oldState <- takeTMVar appJobState
      deadWorkers <- fmap (Map.fromList . catMaybes) . forM (Map.keys $ jobWorkers oldState) $ \a -> fmap (a,) <$> pollSTM a
      putTMVar appJobState oldState
        { jobWorkers = jobWorkers oldState `Map.withoutKeys` Map.keysSet deadWorkers
        }
      guard . not $ Map.null deadWorkers
      return . forM_ (Map.toList deadWorkers) $ \(jobAsync, result) -> do
        case result of
          Right () -> $logInfoS "JobPoolManager" [st|Job-Executor #{showWorkerId (jobWorkerName oldState jobAsync)} terminated|]
          Left e   -> $logErrorS "JobPoolManager" [st|Job-Executer #{showWorkerId (jobWorkerName oldState jobAsync)} crashed: #{tshow e}|]
        void . lift . allocateLinkedAsync $
          let go = do
                next <- evalRandTIO . mapRandT (liftIO . atomically) . runMaybeT $ do
                  nextVal <- MaybeT . lift . tryReadTChan $ jobWorkers oldState ! jobAsync
                  jobWorkers' <- lift . lift $ jobWorkers <$> readTMVar appJobState
                  receiver <- maybe (lift $ lift retry) return =<< uniformMay jobWorkers'
                  return (nextVal, receiver)
                whenIsJust next $ \(nextVal, receiver) -> do
                  atomically $ writeTChan receiver nextVal
                  go
           in go
        
    terminateGracefully = do
      shouldTerminate <- readTMVar appJobState >>= fmap not . isEmptyTMVar . jobShutdown
      guard shouldTerminate
      return . callCC $ \terminate -> do
        $logInfoS "JobPoolManager" "Shutting down"
        terminate ()

stopJobCtl :: (MonadIO m, MonadBaseControl IO m) => UniWorX -> m ()
-- ^ Stop all worker threads currently running
stopJobCtl UniWorX{appJobState} = do
  didStop <- atomically $ do
    jState <- tryReadTMVar appJobState
    for jState $ \jSt'@JobState{jobShutdown} -> jSt' <$ tryPutTMVar jobShutdown ()
  whenIsJust didStop $ \jSt' -> void . fork . atomically $ do
    workers <- maybe [] (Map.keys . jobWorkers) <$> tryTakeTMVar appJobState
    mapM_ (void . waitCatchSTM) $
      [ jobPoolManager jSt'
      , jobCron jSt'
      ] ++ workers

execCrontab :: RWST JobContext () (HashMap JobCtl (Max UTCTime)) (HandlerT UniWorX IO) ()
-- ^ Keeping a `HashMap` of the latest execution times of `JobCtl`s we have
--   seen, wait for the time of the next job and fire it
execCrontab = do
  mapRWST (liftHandlerT . runDB . setSerializable) $ do
    let
      merge (Entity leId CronLastExec{..})
        | Just job <- Aeson.parseMaybe parseJSON cronLastExecJob
        = State.modify $ HashMap.insertWith (<>) (JobCtlQueue job) (Max cronLastExecTime)
        | otherwise = lift $ delete leId
    runConduit $ transPipe lift (selectSource [] []) .| C.mapM_ merge

  refT <- liftIO getCurrentTime
  settings <- getsYesod appSettings'
  (currentCrontab, (jobCtl, nextMatch)) <- mapRWST (liftIO . atomically) $ do
    crontab <- liftBase . readTVar =<< asks jobCrontab

    State.modify . HashMap.filterWithKey $ \k _ -> HashMap.member k crontab
    prevExec <- State.get
    case earliestJob settings prevExec crontab refT of
      Nothing -> liftBase retry
      Just (_, MatchNone) -> liftBase retry
      Just x -> return (crontab, x)

  let doJob = mapRWST (liftHandlerT . runDBJobs . setSerializable) $ do
          newCrontab <- lift . hoist lift $ determineCrontab'
          if
            | ((==) `on` HashMap.lookup jobCtl) newCrontab currentCrontab
              -> do
                now <- liftIO $ getCurrentTime
                foundation <- getYesod
                let instanceID' = foundation ^. _appInstanceID
                State.modify $ HashMap.alter (Just . ($ Max now) . maybe id (<>)) jobCtl
                case jobCtl of
                  JobCtlQueue job -> do
                    void . lift $ upsertBy
                      (UniqueCronLastExec $ toJSON job)
                      CronLastExec
                        { cronLastExecJob = toJSON job
                        , cronLastExecTime = now
                        , cronLastExecInstance = instanceID'
                        }
                      [ CronLastExecTime =. now ]
                    lift $ queueDBJob job
                  other -> runReaderT ?? foundation $ writeJobCtl other
            | otherwise
              -> mapRWST (liftIO . atomically) $
                   liftBase . void . flip swapTVar newCrontab =<< asks jobCrontab

  case nextMatch of
    MatchAsap -> doJob
    MatchNone -> return ()
    MatchAt nextTime -> do
      JobContext{jobCrontab} <- ask
      nextTime' <- applyJitter jobCtl nextTime
      $logDebugS "Cron" [st|Waiting until #{tshow (utcToLocalTimeTZ appTZ nextTime')} to execute #{tshow jobCtl}|]
      logFunc <- askLoggerIO
      whenM (liftIO . flip runLoggingT logFunc $ waitUntil jobCrontab currentCrontab nextTime')
        doJob
  where
    acc :: NominalDiffTime
    acc = 1e-3

    debouncingAcc :: AppSettings -> JobCtl -> NominalDiffTime
    debouncingAcc AppSettings{appNotificationRateLimit} = \case
      JobCtlQueue (JobQueueNotification _) -> appNotificationRateLimit
      _ -> acc

    applyJitter :: (MonadHandler m, HandlerSite m ~ UniWorX, Hashable seed) => seed -> UTCTime -> m UTCTime
    applyJitter seed t = do
      appInstance <- getsYesod appInstanceID
      let
        halfRange = truncate $ 0.5 / acc
        diff = evalRand ( (* acc) . fromInteger <$> getRandomR (- halfRange, halfRange)) $ mkStdGen (hash appInstance `hashWithSalt` seed)
      return $ addUTCTime diff t

    earliestJob :: AppSettings -> HashMap JobCtl (Max UTCTime) -> Crontab JobCtl -> UTCTime -> Maybe (JobCtl, CronNextMatch UTCTime)
    earliestJob settings lastTimes crontab now = foldr go' Nothing $ HashMap.toList crontab
      where
        go' (jobCtl, cron) mbPrev
          | Just (_, t') <- mbPrev
          , t' < t
          = mbPrev
          | otherwise
          = Just (jobCtl, t)
          where
            t = nextCronMatch appTZ (getMax <$> HashMap.lookup jobCtl lastTimes) (debouncingAcc settings jobCtl) now cron

    waitUntil :: (Eq a, MonadResourceBase m, MonadLogger m) => TVar a -> a -> UTCTime -> m Bool
    waitUntil crontabTV crontab nextTime = runResourceT $ do
      diffT <- diffUTCTime nextTime <$> liftIO getCurrentTime
      let waitTime = fromInteger (truncate $ diffT / acc) * toRational acc
          waitTime'
            | diffT < acc = "Done"
            | otherwise = tshow (realToFrac waitTime :: NominalDiffTime)
      $logDebugS "waitUntil" [st|#{tshow diffT} (#{waitTime'})|]
      if
        | diffT < acc -> return True
        | otherwise   -> do
            delay <- liftIO . newDelay . round $ waitTime * 1e6
            let
              awaitDelayThread = False <$ waitDelay delay
              awaitCrontabChange = do
                crontab' <- readTVar crontabTV
                True <$ guard (crontab /= crontab')
            crontabChanged <- liftIO . atomically $ awaitCrontabChange <|> awaitDelayThread
            bool (waitUntil crontabTV crontab nextTime) (return False) crontabChanged

mkLogIdent :: JobWorkerId -> Text
mkLogIdent wId = "Job-Executor " <> showWorkerId wId

handleJobs' :: JobWorkerId -> Sink JobCtl (ReaderT JobContext Handler) ()
handleJobs' wNum = C.mapM_ $ \jctl -> do
  $logDebugS logIdent $ tshow jctl
  res <- fmap (either Just $ const Nothing) . try $ handleCmd jctl
  sentRes <- mapReaderT (liftIO . atomically) $ do
    resVars <- HashMap.lookup jctl <$> (lift . readTVar =<< asks jobConfirm)
    lift $ foldrM (\resVar -> bool (tryPutTMVar resVar res) $ return True) False (maybe [] NonEmpty.toList resVars)
  case res of
    Just err
      | not sentRes -> $logErrorS logIdent $ tshow err
    _other -> return ()
  where
    logIdent = mkLogIdent wNum
    
    handleQueueException :: MonadLogger m => JobQueueException -> m ()
    handleQueueException (JInvalid jId j) = $logWarnS logIdent $ "Invalid QueuedJob (#" ++ tshow (fromSqlKey jId) ++ "): " ++ tshow j
    handleQueueException (JNonexistant jId) = $logInfoS logIdent $ "Saw nonexistant queue id: " ++ tshow (fromSqlKey jId)
    handleQueueException (JLocked jId lInstance lTime)  = $logDebugS logIdent $ "Saw locked QueuedJob: " ++ tshow (fromSqlKey jId, lInstance, lTime)

    handleCmd JobCtlNoOp = return ()
    handleCmd JobCtlFlush = void . lift . runDB . runConduit $ selectKeys [] [ Asc QueuedJobCreationTime ] .| C.mapM_ (lift . writeJobCtl . JobCtlPerform)
    handleCmd (JobCtlQueue job) = lift $ queueJob' job
    handleCmd (JobCtlPerform jId) = lift . handle handleQueueException . jLocked jId $ \j@QueuedJob{..} -> do
      content <- case fromJSON queuedJobContent of
        Aeson.Success c -> return c
        Aeson.Error t -> do
          $logErrorS logIdent $ "Aeson decoding error: " <> pack t
          throwM $ JInvalid jId j
      
      $logDebugS logIdent . LT.toStrict . decodeUtf8 $ Aeson.encode content

      performJob content

      -- `performJob` is expected to throw an exception if it detects that the job was not done
      runDB $ delete jId
    handleCmd JobCtlDetermineCrontab = do
      newCTab <- liftHandlerT . runDB $ setSerializable determineCrontab'
      -- logDebugS logIdent $ tshow newCTab
      mapReaderT (liftIO . atomically) $ 
        lift . void . flip swapTVar newCTab =<< asks jobCrontab
    handleCmd (JobCtlGenerateHealthReport kind) = do
      hrStorage <- getsYesod appHealthReport
      newReport@(healthReportStatus -> newStatus) <- lift $ generateHealthReport kind

      $logInfoS (tshow kind) $ toPathPiece newStatus
      unless (newStatus == HealthSuccess) $ do
        $logErrorS (tshow kind) $ tshow newReport
  
      liftIO $ do
        now <- getCurrentTime
        let updateReports = Set.insert (now, newReport)
                          . Set.filter (((/=) `on` classifyHealthReport) newReport . snd)
        atomically . modifyTVar' hrStorage $ force . updateReports

jLocked :: QueuedJobId -> (QueuedJob -> Handler a) -> Handler a
jLocked jId act = do
  hasLock <- liftIO $ newTVarIO False

  let
    lock = runDB . setSerializable $ do
      qj@QueuedJob{..} <- maybe (throwM $ JNonexistant jId) return =<< get jId
      instanceID' <- getsYesod $ view instanceID
      threshold <- getsYesod $ view _appJobStaleThreshold
      now <- liftIO getCurrentTime
      hadStale <- maybeT (return False) $ do
        lockTime <- MaybeT $ return queuedJobLockTime
        lockInstance <- MaybeT $ return queuedJobLockInstance
        if
          | lockInstance == instanceID'
          , diffUTCTime now lockTime >= threshold
            -> return True
          | otherwise
            -> throwM $ JLocked jId lockInstance lockTime
      when hadStale .
        $logWarnS "Jobs" $ "Ignored stale lock: " <> tshow qj
      val <- updateGet jId [ QueuedJobLockInstance =. Just instanceID'
                           , QueuedJobLockTime =. Just now
                           ]
      liftIO . atomically $ writeTVar hasLock True
      return val

    unlock = whenM (liftIO . atomically $ readTVar hasLock) $
      runDB . setSerializable $
        update jId [ QueuedJobLockInstance =. Nothing
                   , QueuedJobLockTime =. Nothing
                   ]
      
  bracket lock (const unlock) act


pruneLastExecs :: Crontab JobCtl -> DB ()
pruneLastExecs crontab = runConduit $ selectSource [] [] .| C.mapM_ ensureCrontab
  where
    ensureCrontab (Entity leId CronLastExec{..})
      | Just job <- Aeson.parseMaybe parseJSON cronLastExecJob
      , HashMap.member (JobCtlQueue job) crontab
      = return ()
      | otherwise = delete leId

determineCrontab' :: DB (Crontab JobCtl)
determineCrontab' = (\ct -> ct <$ pruneLastExecs ct) =<< determineCrontab


performJob :: Job -> HandlerT UniWorX IO ()
performJob = $(dispatchTH ''Job)
