{-# LANGUAGE OverloadedStrings #-}
{-# LANGUAGE RecordWildCards   #-}

module Hyperion.Config where

import qualified Data.Text              as T
import           Data.Time.Format       (defaultTimeLocale, formatTime)
import           Data.Time.LocalTime    (getZonedTime)
import           Hyperion.Cluster
import qualified Hyperion.Database      as DB
import           Hyperion.HoldServer    (HoldMap)
import qualified Hyperion.Log           as Log
import           Hyperion.ProgramId
import qualified Hyperion.Slurm         as Slurm
import           Hyperion.Util          (savedExecutable)
import           Hyperion.WorkerCpuPool (SSHCommand)
import           Hyperion.TokenPool     (newTokenPool)
import           System.Directory       (copyFile, createDirectoryIfMissing)
import           System.FilePath.Posix  (takeBaseName, takeDirectory, (<.>),
                                         (</>))
import Hyperion.LockMap (newLockMap)

-- | Global configuration for "Hyperion" cluster. 
data HyperionConfig = HyperionConfig
  { -- | Default options to use for @sbatch@ submissions
    HyperionConfig -> SbatchOptions
defaultSbatchOptions :: Slurm.SbatchOptions
  , -- | Maximum number of jobs to submit at a time
    HyperionConfig -> Maybe Int
maxSlurmJobs         :: Maybe Int
    -- | Base directory for working dirs produced by 'newWorkDir'
  , HyperionConfig -> FilePath
dataDir              :: FilePath
    -- | Base directory for all the log files 
  , HyperionConfig -> FilePath
logDir               :: FilePath
    -- | Base directory for databases
  , HyperionConfig -> FilePath
databaseDir          :: FilePath
    -- | Base directory for copies of the main executable
  , HyperionConfig -> FilePath
execDir              :: FilePath
    -- | Base directory for SLURM job files
  , HyperionConfig -> FilePath
jobDir              :: FilePath
    -- | The command to run the main executable. Automatic if 'Nothing' (see 'newClusterEnv')
  , HyperionConfig -> Maybe FilePath
hyperionCommand      :: Maybe FilePath
    -- | The database from which to initiate the program database
  , HyperionConfig -> Maybe FilePath
initialDatabase      :: Maybe FilePath
    -- | The command used to run @ssh@ on nodes. Usually can be safely set to
    -- 'Nothing'. See 'SSHCommand' for details.
  , HyperionConfig -> SSHCommand
sshRunCommand        :: SSHCommand
    -- | Email address for cluster notifications from
    -- hyperion. Nothing means no emails will be sent. Note that this
    -- setting can be different from the one in defaultSbatchOptions,
    -- which controls notifications from SLURM.
  , HyperionConfig -> Maybe Text
emailAddr            :: Maybe T.Text
  }

-- | Default configuration, with all paths built form a single
-- 'baseDirectory'
defaultHyperionConfig :: FilePath -> HyperionConfig
defaultHyperionConfig :: FilePath -> HyperionConfig
defaultHyperionConfig FilePath
baseDirectory = HyperionConfig :: SbatchOptions
-> Maybe Int
-> FilePath
-> FilePath
-> FilePath
-> FilePath
-> FilePath
-> Maybe FilePath
-> Maybe FilePath
-> SSHCommand
-> Maybe Text
-> HyperionConfig
HyperionConfig
  { defaultSbatchOptions :: SbatchOptions
defaultSbatchOptions = SbatchOptions
Slurm.defaultSbatchOptions
  , maxSlurmJobs :: Maybe Int
maxSlurmJobs         = Maybe Int
forall a. Maybe a
Nothing
  , dataDir :: FilePath
dataDir              = FilePath
baseDirectory FilePath -> FilePath -> FilePath
</> FilePath
"data"
  , logDir :: FilePath
logDir               = FilePath
baseDirectory FilePath -> FilePath -> FilePath
</> FilePath
"logs"
  , databaseDir :: FilePath
databaseDir          = FilePath
baseDirectory FilePath -> FilePath -> FilePath
</> FilePath
"databases"
  , execDir :: FilePath
execDir              = FilePath
baseDirectory FilePath -> FilePath -> FilePath
</> FilePath
"executables"
  , jobDir :: FilePath
jobDir               = FilePath
baseDirectory FilePath -> FilePath -> FilePath
</> FilePath
"jobs"
  , hyperionCommand :: Maybe FilePath
hyperionCommand      = Maybe FilePath
forall a. Maybe a
Nothing
  , initialDatabase :: Maybe FilePath
initialDatabase      = Maybe FilePath
forall a. Maybe a
Nothing
  , sshRunCommand :: SSHCommand
sshRunCommand        = SSHCommand
forall a. Maybe a
Nothing
  , emailAddr :: Maybe Text
emailAddr            = Maybe Text
forall a. Maybe a
Nothing
  }

-- | Takes 'HyperionConfig' and returns 'ClusterEnv', the path to the executable,
-- and a new 'HoldMap.
--
-- Things to note: 
--
--     * 'programId' is generated randomly.
--     * If 'hyperionCommand' is specified in 'HyperionConfig', then
--       'hyperionExec' == 'hyperionCommand'. Otherwise the running executable 
--       is copied to 'execDir' with a unique name, and that is used as 'hyperionExec'.
--     * 'newDatabasePath' is used to determine 'programDatabase' from 'initialDatabase'
--       and 'databaseDir', 'programId'.
--     * 'timedProgramDir' is used to determine 'programLogDir' and 'programDataDir' 
--       from the values in 'HyperionConfig' and 'programId'.
--     * 'slurmWorkerLauncher' is used for 'clusterWorkerLauncher'
--     * 'clusterDatabaseRetries' is set to 'defaultDBRetries'.
newClusterEnv :: HyperionConfig -> HoldMap -> Int -> IO (ClusterEnv, FilePath)
newClusterEnv :: HyperionConfig -> HoldMap -> Int -> IO (ClusterEnv, FilePath)
newClusterEnv HyperionConfig{FilePath
Maybe Int
Maybe FilePath
SSHCommand
Maybe Text
SbatchOptions
emailAddr :: Maybe Text
sshRunCommand :: SSHCommand
initialDatabase :: Maybe FilePath
hyperionCommand :: Maybe FilePath
jobDir :: FilePath
execDir :: FilePath
databaseDir :: FilePath
logDir :: FilePath
dataDir :: FilePath
maxSlurmJobs :: Maybe Int
defaultSbatchOptions :: SbatchOptions
emailAddr :: HyperionConfig -> Maybe Text
sshRunCommand :: HyperionConfig -> SSHCommand
initialDatabase :: HyperionConfig -> Maybe FilePath
hyperionCommand :: HyperionConfig -> Maybe FilePath
jobDir :: HyperionConfig -> FilePath
execDir :: HyperionConfig -> FilePath
databaseDir :: HyperionConfig -> FilePath
logDir :: HyperionConfig -> FilePath
dataDir :: HyperionConfig -> FilePath
maxSlurmJobs :: HyperionConfig -> Maybe Int
defaultSbatchOptions :: HyperionConfig -> SbatchOptions
..} HoldMap
holdMap Int
holdPort = do
  ProgramId
programId    <- IO ProgramId
newProgramId
  FilePath
hyperionExec <- IO FilePath
-> (FilePath -> IO FilePath) -> Maybe FilePath -> IO FilePath
forall b a. b -> (a -> b) -> Maybe a -> b
maybe
    (FilePath -> FilePath -> IO FilePath
savedExecutable FilePath
execDir (Text -> FilePath
T.unpack (ProgramId -> Text
programIdToText ProgramId
programId)))
    FilePath -> IO FilePath
forall (m :: * -> *) a. Monad m => a -> m a
return
    Maybe FilePath
hyperionCommand
  FilePath
programDatabase <- Maybe FilePath -> FilePath -> ProgramId -> IO FilePath
newDatabasePath Maybe FilePath
initialDatabase FilePath
databaseDir ProgramId
programId
  FilePath
programLogDir <- FilePath -> ProgramId -> IO FilePath
timedProgramDir FilePath
logDir ProgramId
programId
  FilePath
programDataDir <- FilePath -> ProgramId -> IO FilePath
timedProgramDir FilePath
dataDir ProgramId
programId
  TokenPool
sbatchTokenPool <- Maybe Int -> IO TokenPool
newTokenPool Maybe Int
maxSlurmJobs
  let clusterJobOptions :: SbatchOptions
clusterJobOptions = SbatchOptions
defaultSbatchOptions { chdir :: Maybe FilePath
Slurm.chdir = FilePath -> Maybe FilePath
forall a. a -> Maybe a
Just FilePath
jobDir }
      programSSHCommand :: SSHCommand
programSSHCommand = SSHCommand
sshRunCommand
      clusterProgramInfo :: ProgramInfo
clusterProgramInfo = ProgramInfo :: ProgramId
-> FilePath -> FilePath -> FilePath -> SSHCommand -> ProgramInfo
ProgramInfo {FilePath
SSHCommand
ProgramId
programSSHCommand :: SSHCommand
programDataDir :: FilePath
programLogDir :: FilePath
programDatabase :: FilePath
programId :: ProgramId
programSSHCommand :: SSHCommand
programDataDir :: FilePath
programLogDir :: FilePath
programDatabase :: FilePath
programId :: ProgramId
..}
      clusterWorkerLauncher :: SbatchOptions -> ProgramInfo -> WorkerLauncher JobId
clusterWorkerLauncher = Maybe Text
-> FilePath
-> HoldMap
-> Int
-> TokenPool
-> SbatchOptions
-> ProgramInfo
-> WorkerLauncher JobId
slurmWorkerLauncher Maybe Text
emailAddr FilePath
hyperionExec HoldMap
holdMap Int
holdPort TokenPool
sbatchTokenPool
      clusterDatabaseRetries :: Int
clusterDatabaseRetries = Int
defaultDBRetries
  Pool Connection
clusterDatabasePool <- FilePath -> IO (Pool Connection)
DB.newDefaultPool FilePath
programDatabase
  LockMap
clusterLockMap <- IO LockMap
newLockMap
  (ClusterEnv, FilePath) -> IO (ClusterEnv, FilePath)
forall (m :: * -> *) a. Monad m => a -> m a
return (ClusterEnv :: (SbatchOptions -> ProgramInfo -> WorkerLauncher JobId)
-> ProgramInfo
-> SbatchOptions
-> Pool Connection
-> Int
-> LockMap
-> ClusterEnv
ClusterEnv{Int
LockMap
SbatchOptions
Pool Connection
ProgramInfo
SbatchOptions -> ProgramInfo -> WorkerLauncher JobId
clusterLockMap :: LockMap
clusterDatabaseRetries :: Int
clusterDatabasePool :: Pool Connection
clusterJobOptions :: SbatchOptions
clusterProgramInfo :: ProgramInfo
clusterWorkerLauncher :: SbatchOptions -> ProgramInfo -> WorkerLauncher JobId
clusterLockMap :: LockMap
clusterDatabasePool :: Pool Connection
clusterDatabaseRetries :: Int
clusterWorkerLauncher :: SbatchOptions -> ProgramInfo -> WorkerLauncher JobId
clusterProgramInfo :: ProgramInfo
clusterJobOptions :: SbatchOptions
..}, FilePath
hyperionExec)

-- | Returns the path to a new database, given 'Maybe' inital database filepath
-- and base directory
--
-- If 'ProgramId' id is @XXXXX@ and initial database filename is @original.sqlite@,
-- then the new filename is @original-XXXXX.sqlite@. If initial database path is
-- 'Nothing', then the filename is @XXXXX.sqlite@.
--
-- The path is in subdirectory @YYYY-mm@ (determined by current date) of base directory. 
--
-- If inital database is given, then the new database is initilized with its contents.
newDatabasePath :: Maybe FilePath -> FilePath -> ProgramId -> IO FilePath
newDatabasePath :: Maybe FilePath -> FilePath -> ProgramId -> IO FilePath
newDatabasePath Maybe FilePath
mOldDb FilePath
baseDir ProgramId
progId = do
  let base :: FilePath
base = case Maybe FilePath
mOldDb of
        Maybe FilePath
Nothing -> FilePath
""
        Just FilePath
f  -> FilePath -> FilePath
takeBaseName FilePath
f FilePath -> FilePath -> FilePath
forall a. [a] -> [a] -> [a]
++ FilePath
"-"
  FilePath
date <- TimeLocale -> FilePath -> ZonedTime -> FilePath
forall t. FormatTime t => TimeLocale -> FilePath -> t -> FilePath
formatTime TimeLocale
defaultTimeLocale FilePath
"%Y-%m" (ZonedTime -> FilePath) -> IO ZonedTime -> IO FilePath
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> IO ZonedTime
getZonedTime
  let newDb :: FilePath
newDb = FilePath
baseDir FilePath -> FilePath -> FilePath
</> FilePath
date FilePath -> FilePath -> FilePath
</> (FilePath
base FilePath -> FilePath -> FilePath
forall a. [a] -> [a] -> [a]
++ Text -> FilePath
T.unpack (ProgramId -> Text
programIdToText ProgramId
progId)) FilePath -> FilePath -> FilePath
<.> FilePath
"sqlite"
  Bool -> FilePath -> IO ()
createDirectoryIfMissing Bool
True (FilePath -> FilePath
takeDirectory FilePath
newDb)
  case Maybe FilePath
mOldDb of
    Maybe FilePath
Nothing -> () -> IO ()
forall (m :: * -> *) a. Monad m => a -> m a
return ()
    Just FilePath
f -> do
      Text -> (FilePath, FilePath) -> IO ()
forall a (m :: * -> *). (Show a, MonadIO m) => Text -> a -> m ()
Log.info Text
"Copying database" (FilePath
f, FilePath
newDb)
      FilePath -> FilePath -> IO ()
copyFile FilePath
f FilePath
newDb
  FilePath -> IO FilePath
forall (m :: * -> *) a. Monad m => a -> m a
return FilePath
newDb

-- | Given base directory and 'ProgramId' (@==XXXXX@), returns the @YYYY-mm/XXXXX@ 
-- subdirectory of the base directory (determined by current date).
timedProgramDir :: FilePath -> ProgramId -> IO FilePath
timedProgramDir :: FilePath -> ProgramId -> IO FilePath
timedProgramDir FilePath
baseDir ProgramId
progId = do
  FilePath
date <- TimeLocale -> FilePath -> ZonedTime -> FilePath
forall t. FormatTime t => TimeLocale -> FilePath -> t -> FilePath
formatTime TimeLocale
defaultTimeLocale FilePath
"%Y-%m" (ZonedTime -> FilePath) -> IO ZonedTime -> IO FilePath
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> IO ZonedTime
getZonedTime
  FilePath -> IO FilePath
forall (m :: * -> *) a. Monad m => a -> m a
return (FilePath -> IO FilePath) -> FilePath -> IO FilePath
forall a b. (a -> b) -> a -> b
$ FilePath
baseDir FilePath -> FilePath -> FilePath
</> FilePath
date FilePath -> FilePath -> FilePath
</> Text -> FilePath
T.unpack (ProgramId -> Text
programIdToText ProgramId
progId)