{-# LANGUAGE GADTs #-}

module RunWork
  (
    -- types
    WorkExitStatus(..)
  , Work(..)
  , AnyWork(..)
    -- functions regarding exit status
  , readWorkExitStatus
  , isExitCodeTimeout
  , isExitCodeSSHError
  , isExitCodeCTRLCError
  , exitCodeToStatus
  , showExitCodeMsg
  , exitCodeColour
  , rerunStatusStandard
  , isWorkRerun
  , isWorkGrep
  , isWorkStatus
  , isOneOfWorkStatus
    -- other functions
  , reduceWorkGroups
  , composeWorkCommand
  , showWorkWithStatus
  , logFileNames
  , composeLogFileNames
  , removeWorkLogs
  , resetWorkStatus
  ) where

import           RunAuxiliary

import           Data.Char (toLower)
import           Data.List (intercalate)
import qualified Data.Set as Set
import           Data.Text (pack)
import           Data.Text.IO (writeFile)
import           Prelude hiding (writeFile)

import           System.Console.ANSI
import           Text.Printf (printf)

import           Control.Exception (IOException, catch)
import           System.Directory (removeFile)
import           System.Exit (ExitCode(ExitSuccess, ExitFailure))
import           System.FilePath ((</>), (<.>))


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

-- exit status (e.g. used stored in the group file for update runs)
data WorkExitStatus = NONE
                    | SUCCEEDED
                    | TIMEOUT
                    | FAILED
                    | SSH_ERROR
                    | CTRLC_ERROR
                    deriving (Eq, Show)


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

readWorkExitStatus :: String -> WorkExitStatus
readWorkExitStatus err
  | "FAILED"      `elem` ls = FAILED
  | "TIMEOUT"     `elem` ls = TIMEOUT
  | "SUCCEEDED"   `elem` ls = SUCCEEDED
  | "CTRLC_ERROR" `elem` ls = CTRLC_ERROR
  | "SSH_ERROR"   `elem` ls = SSH_ERROR
  | otherwise               = NONE
  where ls = lines err


-- see man timeout to find out which exit code is returned by the timeout command
isExitCodeTimeout :: Int -> Bool
isExitCodeTimeout err = err == 124

isExitCodeSSHError :: Int -> Bool
isExitCodeSSHError err = err == 255

isExitCodeCTRLCError :: Int -> Bool
isExitCodeCTRLCError err = err == -2

-- map exit code to a single keyword
-- (used as parameter for the runGrep script and written to the group file
-- to be used for update runs)
exitCodeToStatus :: ExitCode -> WorkExitStatus
exitCodeToStatus ExitSuccess  = SUCCEEDED
exitCodeToStatus (ExitFailure i)
  | isExitCodeTimeout i    = TIMEOUT
  | isExitCodeSSHError i   = SSH_ERROR
  | isExitCodeCTRLCError i = CTRLC_ERROR
  | otherwise              = FAILED


-- show exit code for a short message (this is more verbose than show, as it
-- shows the exit code number)
showExitCodeMsg :: ExitCode -> String
showExitCodeMsg ExitSuccess = "succeeded"
showExitCodeMsg (ExitFailure i)
  | isExitCodeTimeout i    = "timeout"
  | isExitCodeSSHError i   = "ssh error"
  | isExitCodeCTRLCError i = "ctrl-c"
  | otherwise              = printf "failure %d" i


exitCodeColour :: ExitCode -> [SGR]
exitCodeColour ExitSuccess = [ SetColor Foreground Dull Blue ]
exitCodeColour (ExitFailure i)
  | isExitCodeTimeout i    = [ SetColor Foreground Dull Yellow ]
  | isExitCodeSSHError i   = [ SetColor Foreground Dull Cyan ]
  | isExitCodeCTRLCError i = [ SetColor Foreground Dull Cyan ]
  | otherwise              = [ SetColor Foreground Dull Red ]


-- should a work unit be rerun, this should be done if status is FAILED,
-- but also if it is NONE, moreover add CTRCL_ERROR and SSH_ERROR
-- (ssh error should be caught and not appear as status in the status file)
rerunStatusStandard :: [WorkExitStatus]
rerunStatusStandard = [ NONE
                      , FAILED
                      , CTRLC_ERROR
                      , SSH_ERROR
                      ]


isWorkRerun :: Work a => [WorkExitStatus] -> a -> Bool
isWorkRerun = isOneOfWorkStatus


-- should a work unit be included in a grep run,
-- this should be done if status is SUCCEEDED, TIMEOUT or FAILED
grepStatus :: [WorkExitStatus]
grepStatus = [SUCCEEDED, TIMEOUT, FAILED]

isWorkGrep :: Work a => a -> Bool
isWorkGrep w = status `elem` grepStatus
  where status = workExitStatus w


-- check whether a work unit has a given work status
isWorkStatus :: Work a => WorkExitStatus -> a -> Bool
isWorkStatus status w = workExitStatus w == status

-- check whether the status of a work unit is among a given list of status
isOneOfWorkStatus :: Work a => [WorkExitStatus] -> a -> Bool
isOneOfWorkStatus ss w = workExitStatus w `elem` ss


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

class Work a where
  -- identifying name of the work (e.g. filename or part of the workDir directory)
  workName       :: a -> String
  -- working directory
  workDir        :: a -> FilePath
  -- function to execute the work (locally on the master without any timeout
  -- or threads settings, which are added later on)
  workCommand    :: a -> FilePath
  -- group number
  workGroups     :: a -> Set.Set GroupId
  -- ".status" file with complete path
  workStatusFile :: a -> FilePath
  -- exit status (usually read from the group file)
  workExitStatus :: a -> WorkExitStatus
  -- set the work groups a work unit belongs to
  setWorkGroups  :: Set.Set GroupId -> a -> a
  -- set the exit status of a work (e.g. overwritten by a clean)
  setExitStatus  :: WorkExitStatus -> a -> a


data AnyWork where { AnyWork :: Work a => a -> AnyWork }


instance Work AnyWork where
  workName       (AnyWork w) = workName w
  workDir        (AnyWork w) = workDir w
  workCommand    (AnyWork w) = workCommand w
  workGroups     (AnyWork w) = workGroups w
  workStatusFile (AnyWork w) = workStatusFile w
  workExitStatus (AnyWork w) = workExitStatus w
  setWorkGroups gs (AnyWork w) = AnyWork $ setWorkGroups gs w
  setExitStatus s  (AnyWork w) = AnyWork $ setExitStatus s w



-- function which composes the command based on the workCommand and further data,
-- it takes number of threads, an optional timeout value in seconds,
-- a function which prefixes any local directory such that a node can access the directory
-- on ther master (i.e. prefix with some /net/xyz),
-- a function composing the command on a node (which takes work directory and the command itself
-- on the master node, see nodeCommand as part of Node)
composeWorkCommand :: Work a => a -> NumThreads
                   -> (FilePath -> FilePath)
                   -> String
composeWorkCommand w nbThreads ndPrefix =
  let cmdBase = ndPrefix $ workCommand w
  in printf "env OMP_NUM_THREADS=%d %s" nbThreads cmdBase


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

instance Eq AnyWork where
  (==) x y = workDir x == workDir y


instance Ord AnyWork where
  compare x y = compare (map toLower $ workDir x) (map toLower $ workDir y)


instance Show AnyWork where
  show w = "AnyWork: " ++ (intercalate ", " $ map ($w) [workName, show . workGroups])


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

-- intersect the given set of group id of a work unit with a selected set of group ids
reduceWorkGroups :: Work a => Set.Set GroupId -> a -> a
reduceWorkGroups gs w = let is = Set.intersection gs (workGroups w)
                        in setWorkGroups is w


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

-- just show relative directory and previous status,
-- this function is used in printing work (within blocks grouped by group id,
-- so no need to print the group id)
showWorkWithStatus :: Work a => a -> String
showWorkWithStatus w = printf "{%-11s} \"%s\"" status name
  where name = workName w
        status = show $ workExitStatus w


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

logFileNames :: FilePath -> FilePath -> (FilePath, FilePath)
logFileNames logDir logFile = (logDirFile <.> "out", logDirFile <.> "err")
  where logDirFile = logDir </> logFile


composeLogFileNames :: Work a => a -> FilePath -> (FilePath, FilePath)
composeLogFileNames w = logFileNames (workDir w)


handleRemoveLogException :: IOException -> IO ()
handleRemoveLogException _ = return ()


removeWorkLogs :: Work a => FilePath -> a -> IO ()
removeWorkLogs logFile w = do
  let (logout, logerr) = logFileNames (workDir w) logFile
  removeFile logout `catch` handleRemoveLogException
  removeFile logerr `catch` handleRemoveLogException


resetWorkStatus :: Work a => a -> IO ()
resetWorkStatus w = writeFile (workStatusFile w) (pack $ printf "%s\n" $ show NONE)
