{-# LANGUAGE TemplateHaskell #-}
{-# OPTIONS_HADDOCK hide #-}
-- | Run.hs
--
-- Functions for starting DNA programs
module DNA.Run (
    dnaRun
  ) where

import Control.Applicative
import Control.Monad
import Control.Exception (onException,SomeException)
import Control.Distributed.Process      hiding (onException)
import Control.Distributed.Process.Closure
import Control.Distributed.Process.Node (initRemoteTable)
import Data.Char
import Network.BSD        (getHostName)
import System.Environment (getExecutablePath,getEnv,lookupEnv)
import System.Directory
import System.Process
import System.FilePath    ((</>))
import System.Exit
import System.Posix.Process (getProcessID,executeFile)
import Text.ParserCombinators.ReadP hiding (many)
import Text.Printf

import DNA.SlurmBackend (initializeBackend,startMaster,startSlave,terminateAllSlaves)
import qualified DNA.SlurmBackend as CH
import DNA.CmdOpts
import DNA.Logging
import DNA.Types
import DNA.DSL ( DNA(..) )
import DNA.Interpreter
import DNA.Interpreter.Run
import DNA.Interpreter.Types



-- |
-- Execute DNA program. First parameter is list of remote tables. Each
-- invocation of 'remotable' generate @__remoteTable@ top level
-- identifier with type @RemoteTable -> RemoteTable@. All such remote
-- tables must composed using '.' and passed to @dnaRun@ as in
-- following example:
--
-- > dnaRun (ModuleA.__remoteTable . ModuleB.__remoteTable) program
--
-- UNIX startup. If command line parameter '--nprocs=N' is
-- given. Program will create N processes on same machine and execute
-- program using these processes as cloud haskell's nodes.
--
-- SLURM startup. Jobs of starting processes is handled to SLURM and
-- processes learn addresses of other processes from environment
-- variables set by SLURM. No command line parameters is required in
-- this case.
dnaRun
    :: (RemoteTable -> RemoteTable) -- ^ Cloud haskell's remote tablse
    -> DNA ()                       -- ^ DNA program
    -> IO ()
dnaRun remoteTable dna = do
    (opts,common) <- dnaParseOptions
    initLogging (dnaLogger common)
    case opts of
      Unix n        -> runUnix n common
      UnixWorker o  -> runUnixWorker rtable o common dna
      Slurm | dnaNoReexec common -> runSlurmWorker rtable "." common dna
            | otherwise          -> runSlurm common
      SlurmWorker d -> runSlurmWorker rtable d common dna
  where
    rtable = ( remoteTable
             . DNA.Interpreter.__remoteTable
             . DNA.Interpreter.Run.__remoteTable
             ) initRemoteTable


----------------------------------------------------------------
-- SLURM
----------------------------------------------------------------

runSlurm :: CommonOpt -> IO ()
runSlurm common = do
    dir       <- getCurrentDirectory
    mslurmJID <- lookupEnv "SLURM_JOBID"
    slurmRnk  <- slurmRank
    localID   <- slurmLocalID
    let slurmJID = case mslurmJID of
            Just s  -> s ++ "-s"
            Nothing -> error "SLURM_JOBID is not set!"
    -- Cd to log dir and reexec
    home <- getEnv "HOME"
    let rankDir | localID == 0  = show slurmRnk
                | otherwise     = show slurmRnk ++ "-" ++ show localID
        logDir = home </> "_dna" </> "logs" </> slurmJID </> rankDir
    createDirectoryIfMissing True logDir
    setCurrentDirectory logDir
    program <- getExecutablePath
    executeFile program True
        [ "+RTS", "-l-au", "-T", "-RTS"
        , "--base-port", show (dnaBasePort common)
        , "--workdir", dir
        , "-v",              show $ logOptVerbose (dnaLogger common)
        , "--measure",       logOptMeasure (dnaLogger common)
        ] Nothing

runSlurmWorker :: RemoteTable -> FilePath -> CommonOpt -> DNA () -> IO ()
runSlurmWorker rtable dir common dna = do
    setCurrentDirectory dir
    -- Obtain data from SLURM
    rank     <- slurmRank
    localID  <- slurmLocalID
    hostList <- slurmHosts
    host     <- getHostName
    let port = dnaBasePort common + localID
    -- FIXME: treat several tasks per node correctly
    backend <- initializeBackend
                 (CH.SLURM (dnaBasePort common) hostList)
                 host (show port) rtable 10
    case (rank, localID) of
      (0,0) -> startMaster backend $ \n -> do
               synchronizationPoint "CH"
               r <- executeDNA (dnaLogger common) dir dna n
               terminateAllSlaves backend
               case r of
                 Just e  -> do errorMsg $ "main actor died: " ++ show e
                               -- FIXME: a bit klunky
                               liftIO exitFailure
                 Nothing -> return ()
      _ -> startSlave backend

slurmRank :: IO Int
slurmRank = do
    mslurmRnk <- (safeRead =<<) <$> lookupEnv "SLURM_PROCID" :: IO (Maybe Int)
    case mslurmRnk of
      Just r  -> return r
      Nothing -> error "SLURM_PROCID is not set! Try starting with --nprocs [n] if you want a local configuration!"

slurmLocalID :: IO Int
slurmLocalID = do

    -- Instead, use PMI_RANK and SLURM_TASKS_PER_NODE to derive
    -- it. This is set if we get launched by mpirun.
    mMpiRnk <- (safeRead =<<) <$> lookupEnv "PMI_RANK" :: IO (Maybe Int)
    case mMpiRnk of
      Just r -> do
        numTasksStr <- getEnv "SLURM_TASKS_PER_NODE"
        nodeStr <- getEnv "SLURM_NODEID"
        let numTasks = runReadP parseNumNode  =<< split ',' numTasksStr
            -- Number of tasks on nodes with lower ranks
            tasksLower = sum $ take (read nodeStr) numTasks
        return (r - tasksLower)

      Nothing -> do

        -- Fall back to the SLURM rank. This only works when we are
        -- run by srun - for Intel's mpirun it apparently always
        -- returns "0".
        mslurmRnk <- (safeRead =<<) <$> lookupEnv "SLURM_LOCALID"
        case mslurmRnk of
         Just r -> return r
         Nothing -> error "Neither PMI_RANK nor SLURM_LOCALID are set!"

-- | Obtain list of hosts from SLURM together with number of tasks per
--   host
slurmHosts :: IO [(String,Int)]
slurmHosts = do
    nodeListStr <- getEnv "SLURM_NODELIST"
    numNodesStr <- getEnv "SLURM_TASKS_PER_NODE"
    --eventMessage $ "SLURM_NODELIST=" ++ nodeListStr
    --eventMessage $ "SLURM_TASKS_PER_NODE=" ++ numNodesStr
    let nodeList = runReadP parseNodeList nodeListStr
        numNodes = runReadP parseNumNode  =<< split ','numNodesStr
    when (length nodeList /= length numNodes) $
        error "Length of SLURM_NODELIST and SLURM_TASKS_PER_NODE does not match"
    return $ nodeList `zip` numNodes

-- Parser for SLURM_TASKS_PER_NODE envvar
parseNumNode :: ReadP [Int]
parseNumNode =  (pure <$> readS_to_P reads <* eof)
     <|> do n <- readS_to_P reads
            _ <- char '('
            _ <-  char 'x'
            r <- readS_to_P reads
            _ <- char ')'
            eof
            return $ replicate r n

-- Parser for SLURM_NODELIST
parseNodeList :: ReadP [String]
parseNodeList = do
    chunks <- list <* eof
    let toStr (Str s)           = [s]
        toStr (Range len (n,m)) = [printf (printf "%%0%ii" len) i | i <- [n .. m]]
        toStr (List rs)         = toStr =<< rs
        toStr (Seq  xs)         = foldl (liftA2 (++)) [""] $ map toStr xs
    return $ toStr =<< chunks
    -- return chunks
  where
    -- Comma-separated list
    list = do r <- many ((Str <$> str) <|> range) `sepBy1` char ','
              return $ map Seq r
    -- Range expression
    range = do _  <- char '['
               rs <- sepBy1 (rangeInt <|> singleInt) (char ',')
               _  <- char ']'
               return $ List rs
    -- Integer ranges
    leadingInt :: ReadP (Int,Int)
    leadingInt = do
        sn <- munch1 isDigit
        n  <- case safeRead sn of
                Just n  -> return n
                Nothing -> empty
        let len = case sn of
                    '0':_ -> length sn
                    _     -> 0
        return (len,n)
    singleInt :: ReadP Chunk
    singleInt = do
        (len,n) <- leadingInt
        return $ Range len (n,n)
    rangeInt :: ReadP Chunk
    rangeInt = do
        (len,n) <- leadingInt
        _ <- char '-'
        m <- int
        return $ Range len (n,m)
    -- String
    str = munch1 $ \c -> or [ isAsciiLower c
                            , isAsciiUpper c
                            , isDigit c
                            , c == '-'
                            ]
    -- Integer
    int = do s <- munch1 isDigit
             return (read s :: Int)

-- Chunk of NODELIST entry
data Chunk = Str String          -- String chunk
           | Range Int (Int,Int) -- Range chunk: width, (range)
           | List [Chunk]
           | Seq  [Chunk]
           deriving (Show)

runReadP :: ReadP a -> String -> a
runReadP p s = case readP_to_S p s of
    [(a,"")] -> a
    _        -> error $ "Cannot parse string: '" ++ s ++ "'"



----------------------------------------------------------------
-- UNIX start
----------------------------------------------------------------

-- Startup of program
runUnix :: Int -> CommonOpt -> IO ()
runUnix n common = do
    -- Check if SLURM envvars are set and fail they're
    do v <- lookupEnv "SLURM_JOBID"
       case v of
         Nothing -> return ()
         _       -> error "SLURM environment variable is set will not use UNIX startup"
    --
    pid     <- getProcessID
    program <- getExecutablePath
    workdir <- getCurrentDirectory
    -- Create log dirs for all processess
    home    <- getEnv "HOME"
    let dnaPID = show pid ++ "-u"
        logDir = home </> "_dna" </> "logs" </> dnaPID
    forM_ [0 .. n-1] $ \k ->
        createDirectoryIfMissing True (logDir </> show k)
    -- Reexec program with enabled eventlog
    setCurrentDirectory (logDir </> "0")
    executeFile program True
        [ "+RTS", "-l-au", "-T", "-RTS"
        , "--nprocs",        show n
        , "--internal-rank", "0"
        , "--internal-pid",  dnaPID
        , "--workdir",       workdir
        , "--base-port",     show (dnaBasePort common)
        , "-v",              show $ logOptVerbose (dnaLogger common)
        , "--measure",       logOptMeasure (dnaLogger common)
        ] Nothing

runUnixWorker :: RemoteTable -> UnixStart -> CommonOpt -> DNA () -> IO ()
runUnixWorker rtable opts common dna = do
    synchronizationPoint "UNIX"
    home    <- getEnv "HOME"
    let basePort = dnaBasePort common
        rank     = dnaUnixRank opts
        port     = basePort + rank
        dnaPID   = dnaUnixPID opts
        nProc    = dnaUnixNProc opts
    let logDir = home </> "_dna" </> "logs" </> dnaPID
    -- Spawn child processes
    pids <- case rank of
        0 -> do program <- getExecutablePath
                forM [1 .. nProc - 1] $ \rnk -> do
                    let dir = logDir </> show rnk
                    createDirectoryIfMissing True dir
                    setCurrentDirectory dir
                    spawnProcess program
                        [ "+RTS", "-l-au", "-T", "-RTS"
                        , "--base-port",     show basePort
                        , "--nprocs",        show nProc
                        , "--internal-rank", show rnk
                        , "--internal-pid",  dnaPID
                        , "--workdir",       dnaUnixWDir opts
                        , "-v",              show $ logOptVerbose (dnaLogger common)
                        , "--measure",       logOptMeasure (dnaLogger common)
                        ]
        _ -> return []
    -- Go back to working dir
    setCurrentDirectory (dnaUnixWDir opts)
    let killChild pid = do
            st <- getProcessExitCode pid
            case st of
              Just _  -> return ()
              Nothing -> terminateProcess pid
    -- Initialize backend
    let ports = map (+ basePort) [0 .. nProc - 1]
    backend <- initializeBackend (CH.Local ports) "localhost" (show port) rtable 10
    -- FIXME: this is source of potential problems if we get exception
    --        during checking children state
    let reapChildren = mapM_ killChild pids
        start = startMaster backend $ \n -> do
            synchronizationPoint "CH"
            r <- executeDNA (dnaLogger common) (dnaUnixWDir opts) dna n
            terminateAllSlaves backend
            case r of
              Just e  -> do errorMsg $ "main actor died: " ++ show e
                            liftIO exitFailure
              Nothing -> return ()
    -- Start master or slave program
    case rank of
      0 -> start `onException` reapChildren
      _ -> startSlave backend


executeDNA :: LoggerOpt -> FilePath -> DNA () -> [NodeId] -> Process (Maybe SomeException)
executeDNA logopt workDir dna nodes = do
    nid <- getSelfNode
    -- FIXME: correctly treat parent-related data
    let param = ActorParam
              { actorParent      = Nothing
              , actorInterpreter = $(mkStaticClosure 'theInterpreter)
              , actorRank        = Rank 0
              , actorGroupSize   = GroupSize 1
              , actorNodes       = VirtualCAD (NodeInfo nid) (map NodeInfo nodes)
              , actorDebugFlags  = []
              , actorSendBack    = error "Impossible: should not evaluate ever"
              , actorLogOpt      = logopt
              , actorWorkDir     = workDir
              }
    r <- try $ runDnaParam param dna
    return $ either Just (const Nothing) r


safeRead :: Read a => String -> Maybe a
safeRead s = do
    [(a,"")] <- Just $ reads s
    return a

split :: Char -> String -> [String]
split c str =
    case break (==c) str of
      (a,"")  -> [a]
      (a,_:s) -> a : split c s
