{-# LANGUAGE GeneralizedNewtypeDeriving #-}

module RunSchedule
  (
    -- types
    Scheduler(..)
    -- functions
  , runWorkScheduler
  , newSchedulerState
  , runWork
  ) where

import           RunAuxiliary
import           RunEnvironment
import           RunNode
import           RunProcess
import           RunWork
import           RunWorkSet

import           Data.Maybe (fromJust, catMaybes, isNothing)
import           Data.List (intercalate)
import qualified Data.Set as Set (member)
import           Data.Traversable (mapM)
import           Data.Time.Clock (getCurrentTime)

import qualified Data.Map.Strict as Map

import           Text.Printf (printf)

import qualified System.Process as Process
import           System.Posix.Signals (Handler, Handler(CatchOnce), installHandler,
                                      sigINT, sigTERM, raiseSignal)
import           System.Random

import           Control.Monad (forM_, filterM)
import           Control.Monad.Loops (whileM_)
import           Control.Monad.Reader
import           Control.Monad.State

import           Control.Concurrent
import           Control.Concurrent.Async


------------------------------------------------------------

-- auxiliary type used during determination of the next work package
-- data of this type contains: (group id, threads, node)
type GroupThreadsNode = (GroupId,NumThreads,AnyNode)


------------------------------------------------------------

data JobData = JobData {
    jNode    :: AnyNode         -- node used for the job
  , jThreads :: NumThreads      -- number of threads used for the job
  , jWork    :: AnyWork         -- work unit computed by the job
  , jHandle  :: Async JobResult -- the result of the job is just the job id,
                                -- and exit status of the job
  }


data SchedulerState = SchedulerState {
    sdWorkSet      :: WorkSet                -- remaining work set
  , sdNodeThreads  :: Map.Map String NumThreads
                                             -- currently available threads for each compute node,
                                             -- accessed by the name of the node
  , sdJobs         :: Map.Map JobId JobData  -- currently running jobs
  , sdJobId        :: JobId                  -- counter for job id
  , sdProcessStore :: MVar (Map.Map JobId Process.ProcessHandle)
                                             -- store all currently active process handles
  , sdRandomGen    :: StdGen
  }


newtype Scheduler a = Scheduler {
    runScheduler :: ReaderT RunEnv (StateT SchedulerState IO) a
  } deriving (Functor, Applicative, Monad, MonadIO,
              MonadState SchedulerState, MonadReader RunEnv)


------------------------------------------------------------


-- config :: RunEnv
-- state :: SchedulerState
-- runStateT (runReaderT (runScheduler k) config) state
runWorkScheduler :: Scheduler a -> RunEnv -> SchedulerState -> IO a
runWorkScheduler k r s = fst <$> runStateT (runReaderT (runScheduler k) r) s


-- construct a new scheduler state
-- it already set the work set data, but node and thread data are just nulled
newSchedulerState :: WorkSet -> IO SchedulerState
newSchedulerState wm = do
  -- to runenv
  -- create a lock for atomic IO and a process handle map
--  lockIO <- newMVar ()
  pstoreVar <- newMVar Map.empty

  -- get a new random number generator
  rgen <- newStdGen

  return SchedulerState {
      sdWorkSet = wm
    , sdNodeThreads = Map.empty
    , sdJobs = Map.empty
    , sdJobId = 0
    , sdProcessStore = pstoreVar
    , sdRandomGen = rgen
    }


-- initialise the map storing available threads for each node
initNodeResources :: Scheduler ()
initNodeResources = do
  ntm <- Map.map threadsTotal <$> reader envNodes
  modify (\sd -> sd { sdNodeThreads = ntm })


------------------------------------------------------------

-- we are finished if there is no work left and no running jobs left
isFinished :: SchedulerState -> Bool
isFinished sd = null (sdWorkSet sd) && null (sdJobs sd)


getRandom :: SchedulerState -> (Int,SchedulerState)
getRandom sd = let (i,rgen') = random $ sdRandomGen sd
               in (i, sd { sdRandomGen = rgen' })


getRandomR :: (Int,Int) -> SchedulerState -> (Int,SchedulerState)
getRandomR r sd = let (i,rgen') = randomR r $ sdRandomGen sd
                  in (i, sd { sdRandomGen = rgen' })


-- depending on order settings it either returns 0 (for first element)
-- or a random index
getIndex :: Scheduler Int
getIndex = do
  isSorted <- reader envExecSorted
  if isSorted
    then return 0
    else state getRandom


getJobId :: Scheduler Int
getJobId = do
  jid <- (+1) <$> gets sdJobId
  modify (\sd -> sd { sdJobId = jid })
  return jid


addJob :: JobId -> JobData -> Scheduler ()
addJob jid jd = do
  jobs <- gets sdJobs
  let threads = jThreads jd
      nd = jNode jd

  -- reduce available threads for node
  modify $ subThreads nd threads
  modify (\sd -> sd { sdJobs = Map.insert jid jd jobs })


-- release the job resources
removeJob :: JobId -> Scheduler ()
removeJob jid = do
  jobs <- gets sdJobs
  let jd = fromJust $ Map.lookup jid jobs
      threads = jThreads jd
      nd = jNode jd

  -- increase available threads for node
  modify $ addThreads nd threads
  modify (\sd -> sd { sdJobs = Map.delete jid jobs })


------------------------------------------------------------

schedulerJobIds :: Scheduler [JobId]
schedulerJobIds = gets $ Map.keys . sdJobs


schedulerJobsData :: Scheduler [JobData]
schedulerJobsData = gets $ Map.elems . sdJobs


schedulerJobsAsyncs :: Scheduler [Async JobResult]
schedulerJobsAsyncs = map jHandle <$> schedulerJobsData


getJobById :: JobId -> Scheduler JobData
getJobById jid = do
  job_maybe <- gets $ Map.lookup jid . sdJobs
  case job_maybe of
    Just job -> return job
    _        -> error "getJobById: invalid job id"

-- return the number of threads used by a job
schedulerJobsThreads :: JobId -> Scheduler NumThreads
schedulerJobsThreads jid = jThreads <$> getJobById jid


schedulerJobsNode :: JobId -> Scheduler AnyNode
schedulerJobsNode jid = jNode <$> getJobById jid


schedulerJobsWork :: JobId -> Scheduler AnyWork
schedulerJobsWork jid = jWork <$> getJobById jid


------------------------------------------------------------

checkSufficientThreads :: GroupThreadsNode -> Scheduler Bool
checkSufficientThreads (_,threads,nd) = do
  let name = nodeName nd
      lookupThreads = fromJust . Map.lookup name
  tavail <- lookupThreads <$> gets sdNodeThreads
  return $ threads <= tavail


-- add available threads for a given node
addThreads :: AnyNode -> NumThreads -> SchedulerState -> SchedulerState
addThreads nd threads sd = sd { sdNodeThreads = nt' }
  where nt' = Map.adjust (threads+) (nodeName nd) $ sdNodeThreads sd


-- subtract available threads for a given node
subThreads :: AnyNode -> NumThreads -> SchedulerState -> SchedulerState
subThreads nd threads = addThreads nd (-threads)


supplementThreads :: (GroupId, AnyNode) -> Scheduler GroupThreadsNode
supplementThreads (gid,nd) = do
  -- take a random number of threads between tmin and tupper
  -- at this point we cannot take available threads into account (for tupper),
  -- as otherwise we end up using the tmin number of threads after awhile
  -- (if a task with lots of threads is finished we get a small random number and can
  -- schedule two or more works with a few threads, now if a work unit with few threads
  -- is finished, tupper is small and we use a few number of threads for the new task,
  -- and we will never be able to schedule with lots of threads)
  let
    ttotal = threadsTotal nd
    (tmin, tmax) = getThreadsForGroup gid nd
    -- use at least tmin, to make sure that the range is proper
    tupper = max tmin $ min tmax ttotal
    r = (tmin, tupper)
  
  threads <- state $ getRandomR r
  return (gid, threads, nd)


-- for each group find all compute nodes which have enough resources to execute
-- a work unit from the given group, return a list of all valid combinations
allGroupThreadsNode :: Scheduler [GroupThreadsNode]
allGroupThreadsNode = do
  -- groups which are still present in the workset map
  gs <- getWorkSetGroups <$> gets sdWorkSet
  -- list of node names available to the scheduler
  ns <- Map.keys <$> gets sdNodeThreads
  nm <- reader envNodes
  -- replace name with complete node
  let ns' = map (\n -> fromJust $ Map.lookup n nm) ns
  -- get a list of all (gid,threads,nd) tuples, and filter for those combinations,
  -- which are valid (node has sufficiently many threads available for the selected number of threads)
  gns_all <- mapM supplementThreads [(gid,nd) | gid <- gs, nd <- ns']
  filterM checkSufficientThreads gns_all


-- select the next group and node (and the number of threads to use for it),
-- the work unit of the selected group is selected later on
selectGroupThreadsNode :: Scheduler (Maybe GroupThreadsNode)
selectGroupThreadsNode = do
  gtns <- allGroupThreadsNode
  i <- getIndex
  return (fst $ removeList i gtns)


removeWork :: GroupId -> Int -> SchedulerState -> (AnyWork,SchedulerState) 
removeWork gid i sd = (w, sd { sdWorkSet = wm' })
  where
    wm = sdWorkSet sd
    gs = head $ filter (Set.member gid) $ Map.keys wm
    (Just w, wm') = removeMapLists gs i wm


-- helper function which deals with the regular case and selects a concrete work
-- once a node and a group have been choosen
takeWork :: GroupThreadsNode -> Scheduler (GroupId,NumThreads,AnyNode,AnyWork)
takeWork (gid,threads,nd) = do
  -- from the preparations we know that the list of work units for gid is not empty,
  -- so we can take an element, no need to check that we took Nothing
  i <- getIndex
  w <- state $ removeWork gid i
  return (gid, threads, nd, w)


-- select a compute node, work unit and number of threads to use
selectNodeWork :: Scheduler (Maybe (GroupId,NumThreads,AnyNode,AnyWork))
selectNodeWork = do
  gtn_maybe <- selectGroupThreadsNode
  case gtn_maybe of
    Nothing  -> return Nothing
    Just gtn -> Just <$> takeWork gtn


------------------------------------------------------------

showStatusNodes :: Scheduler String
showStatusNodes = do
  nts <- gets $ Map.toList . sdNodeThreads
  let ss = map (uncurry $ printf "%s<%d>") nts
  return $ printf "*** available threads on nodes = %s" $ intercalate ", " ss


showWorkQueue :: Scheduler String
showWorkQueue = do
  gwsl <- gets $ Map.toList . sdWorkSet
  let wss = map (\(gs, ws) -> printf "   group = %s, work %s " (showGroups gs)
                               (intercalate ", " $ map workName ws)) gwsl

  return $ printf "work units not executed: \n%s" (intercalate "\n" wss)


showJobList :: Scheduler String
showJobList = do
  js <- schedulerJobIds
  ts <- mapM schedulerJobsThreads js
  ns <- mapM (nodeName <-< schedulerJobsNode) js
  let jtns = zipWith3 (printf "%d<%s:%d>") js ns ts
  return $ intercalate ", " jtns


showStatusWork :: Scheduler String
showStatusWork = do
  gwsl <- gets $ Map.toList . sdWorkSet
  let showGroupStatus (gs,wsl) = printf "%s<%d>" (showGroups gs) (length wsl)
      strRemain = intercalate ", " $ map showGroupStatus gwsl
  return $ if null gwsl
             then "*** no work units left"
             else "*** remaining work units (groupdId<#work>): " ++ strRemain


showStatusAsync :: Scheduler String
showStatusAsync = do
  as <- schedulerJobsAsyncs
  ps <- liftIO $ mapM poll as
  if all isNothing ps
    then printf "*** waiting for (jobId<node:threads>): %s" <$> showJobList
    else return $ printf "*** handling further %d finished work unit(s)"
                         (length $ catMaybes ps)


showStatusWaitForJobs :: Scheduler [String]
showStatusWaitForJobs =
  sequence [showStatusNodes, showStatusWork, showStatusAsync]


------------------------------------------------------------

-- check that the queue is not empty, this indicates a misconfiguration
-- or network problems (ssh errors)
checkJobQueue :: Scheduler ()
checkJobQueue = do
  isJobQueueEmpty <- gets $ Map.null . sdJobs
  when isJobQueueEmpty $ do
    workQueueOut <- showWorkQueue
    let pf = printf "%s"
        msg =    "there are work units which cannot be executed\n"
              ++ "(no groups specified, number of threads too large, ssh failures, etc)\n"
              ++ workQueueOut
    lockIO <- reader envLockIO
    liftIO $ atomicPutStrLn lockIO $ pf msg
    error "revise configuration setup"


-- if an ssh error occurs, remove the node and reschedule the work unit
updateSSHError :: JobId -> Scheduler ()
updateSSHError jid = do
  jobs <- gets sdJobs
  let jd = fromJust $ Map.lookup jid jobs
      nd = jNode jd
      w = jWork jd
      jobs' = Map.delete jid jobs

  nts' <- Map.delete (nodeName nd) <$> gets sdNodeThreads
  wm' <- insertMapLists (workGroups w) w <$> gets sdWorkSet

  modify (\sd -> sd { sdNodeThreads = nts'
                    , sdJobs = jobs'
                    , sdWorkSet = wm'
                    })




-- update scheduler if an ssh error occurs and print some messages
-- of what has been done
runWorkSSHError :: JobId -> Scheduler ()
runWorkSSHError jid = do
  nd <- schedulerJobsNode jid
  w <- schedulerJobsWork jid

  nm <- gets sdNodeThreads
  lockIO <- reader envLockIO
  liftIO $ when (Map.member (nodeName nd) nm) $
    withMVar_ lockIO $
      putStrLn $ printf "*** removing remote node due to ssh error: %s" (nodeName nd)

  liftIO $ atomicPutStrLn lockIO $
    printf "*** rescheduling work unit due to ssh error: %s" (workName w)

  updateSSHError jid


------------------------------------------------------------

runWorkLoopWait :: Scheduler ()
runWorkLoopWait = do
  checkJobQueue

  lockIO <- reader envLockIO
  strInfo <- showStatusWaitForJobs
  liftIO $ withMVar_ lockIO $ mapM_ putStrLn strInfo

  aj <- schedulerJobsAsyncs
  (_,(jid,status)) <- liftIO $ waitAny aj
--  liftIO $ putStrLn $ printf "&&& status = %d, %s" jid (show status)
  case status of
    SSH_ERROR   -> runWorkSSHError jid
    CTRLC_ERROR -> liftIO $ raiseSignal sigINT
    _           -> removeJob jid
  return ()


runWorkLoopNewJob :: (GroupId,NumThreads,AnyNode,AnyWork) -> Scheduler ()
runWorkLoopNewJob (gid,threads,nd,w) = do
  lockIO <- reader envLockIO
  jid <- getJobId

  pstoreVar <- gets sdProcessStore
  timeout_maybe <- Map.lookup gid <$> reader envTimeout
  -- scale timeout if we use less than the maximal number of threads for a group
  -- (we assume that the timeout value is with respect to the maximal number of threads for this group)
  let (_, tmax) = getThreadsForGroup gid nd
      scale = (fromIntegral tmax) / (fromIntegral threads)
      timeout_maybe' = scaleInteger scale <$> timeout_maybe  -- within the Maybe monad

  let pf1 = printf "%s jobId = %4d, node = %s, work unit = %s"
  liftIO $ atomicPutStrLn lockIO $ pf1 "*** fire up" jid (nodeName nd) (workName w)

  let pd = ProcessData { pdNode = nd
                       , pdThreads = threads
                       , pdTimeout = timeout_maybe'
                       , pdWork = w
                       , pdJobId = jid
                       , pdProcessStore = pstoreVar
                       }

  -- create a new async job and store it in the jobs map
  r <- ask
  jh <- liftIO (async $ runProcess executeWork pd r)
  addJob jid JobData { jNode = nd
                     , jThreads = threads
                     , jWork = w
                     , jHandle = jh
                     }


-- get a work unit and a compute unit, if both were obtained, fire up the next job
-- if no compute unit is available the wait until one becomes available and
-- remove the finished job from the jobs list and
-- return the updated loop data
runWorkLoopBody :: Scheduler ()
runWorkLoopBody = do
  gtnw_maybe <- selectNodeWork
  case gtnw_maybe of
    -- not enough execution threads available, wait for another job to finish
    Nothing   -> runWorkLoopWait
    -- fire up a new job
    Just gtnw -> runWorkLoopNewJob gtnw


runWorkLoop :: Scheduler ()
runWorkLoop = whileM_ (gets $ not . isFinished) runWorkLoopBody


------------------------------------------------------------

cleanupProcesses :: MVar (Map.Map JobId Process.ProcessHandle) -> IO ()
cleanupProcesses pstoreVar =
  modifyMVar_ pstoreVar
    (\pstore -> do
        -- get list of all currently active processes
        let phs = Map.elems pstore
            jsout = intercalate ", " $ map show (Map.keys pstore)
        putStrLn $ printf "\nkilling jobs: %s" jsout
        forM_ phs Process.interruptProcessGroupOf
--        forM_ phs Process.terminateProcess
        return Map.empty
    )


sigIntTermHandler :: MVar (Map.Map JobId Process.ProcessHandle) -> Handler
sigIntTermHandler = CatchOnce . cleanupProcesses


------------------------------------------------------------

runWork :: Scheduler ()
runWork = do
  initNodeResources

  pstoreVar <- gets sdProcessStore

  old_sigIntHandler  <- liftIO $ installHandler sigINT (sigIntTermHandler pstoreVar) Nothing
  old_sigTermHandler <- liftIO $ installHandler sigTERM (sigIntTermHandler pstoreVar) Nothing

  -- we want to measure time for profiling purposes
  timeStart <- liftIO getCurrentTime

  -- this is the main loop, which starts all jobs and waits until all
  -- work units are executed and all jobs are finished
  local (setTimeBase timeStart) runWorkLoop

  timeAll <- liftIO $ relativeTime timeStart

  let outFinished = printf "*** work queue finished, overall execution time = %8.1f" timeAll
  liftIO $ putStrLn outFinished

  -- reinstall the original signal handlers
  _ <- liftIO $ installHandler sigINT old_sigIntHandler Nothing
  _ <- liftIO $ installHandler sigTERM old_sigTermHandler Nothing

  return ()
