{- Copyright 2016-2022 NGLess Authors
 - License: MIT
 -}

{-# LANGUAGE FlexibleContexts, CPP #-}

module StandardModules.Parallel
    ( loadModule
    , pasteCounts
    ) where

import qualified Data.ByteString as B
import qualified Data.ByteString.Char8 as B8
import qualified Data.ByteString.Lazy as BL
import qualified Data.ByteString.Builder as BB
import qualified Data.Text as T
import qualified Data.Text.IO as T
import qualified Data.Vector as V
import qualified Data.Vector.Mutable as VM
import           Data.Time (getZonedTime)
import           Data.Time.Format (formatTime, defaultTimeLocale)
import           Data.List.Extra (snoc, chunksOf)

#ifndef WINDOWS
import           System.Posix.Unistd (fileSynchronise)
import           System.Posix.IO (openFd, defaultFileFlags, closeFd, OpenMode(..))
import           Control.Exception (bracket)
#endif


import           System.FilePath
import           GHC.Conc (getNumCapabilities, atomically)
import qualified Control.Concurrent.Async as A
import qualified Control.Concurrent.STM.TBMQueue as TQ
import qualified Data.Conduit.List as CL
import qualified Data.Conduit.TQueue as CA
import qualified Data.Conduit.Algorithms as CAlg
import qualified Data.Conduit.Algorithms.Async as CAlg
import           Data.Conduit.Algorithms.Async (conduitPossiblyCompressedFile)
import           Control.Monad.ST (runST)
import           Control.Monad.Except (throwError)
import           Control.Monad.Extra (allM, unlessM)
import           Control.DeepSeq
import           Data.Traversable
import           Control.Monad.Trans.Class
import           System.AtomicWrite.Writer.Text (atomicWriteFile)
import           System.Random.Shuffle (shuffleM)


import Control.Monad.Trans.Resource
import Control.Monad.State.Lazy
import System.IO
import Data.Default
import Data.Maybe (fromMaybe)
import System.Directory (createDirectoryIfMissing, doesFileExist, getDirectoryContents)

import qualified Data.Hash.MD5 as MD5
import qualified Data.Conduit as C
import qualified Data.Conduit.Combinators as C
import qualified Data.Conduit.Combinators as CC
import qualified Data.Conduit.Binary as CB
import           Data.Conduit ((.|), (.|), ($$+))

import Output
import NGLess
import Modules
import Language
import Transform
import FileOrStream
import Configuration
import FileManagement
import NGLess.NGError
import NGLess.NGLEnvironment

import Interpretation.Write (moveOrCopyCompress, WriteOptions(..))

import Utils.Utils (fmapMaybeM, allSame, moveOrCopy)
import Utils.Conduit
import qualified Utils.LockFile as LockFile
import           Utils.LockFile (LockParameters(..))

syncFile :: FilePath -> IO ()
#ifndef WINDOWS
syncFile fname = do
    bracket (openFd fname ReadWrite Nothing defaultFileFlags)
        closeFd
        fileSynchronise
    -- The code below will not work on Windows
    bracket (openFd (takeDirectory fname) ReadOnly Nothing defaultFileFlags)
        closeFd
        fileSynchronise

#else
syncFile _ = return ()
#endif


setupHashDirectory :: String -> FilePath -> T.Text -> NGLessIO FilePath
setupHashDirectory prefix basename hash = do
    isSubsample <- nConfSubsample <$> nglConfiguration
    let actiondir = basename </> prefix ++ take 8 (T.unpack hash) ++ (if isSubsample then "-subsample" else "")
        scriptfile = actiondir </> "script.ngl"
    liftIO $ createDirectoryIfMissing True actiondir
    unlessM (liftIO $ doesFileExist scriptfile) $ do
        sct <- ngleScriptText <$> nglEnvironment
        liftIO $ atomicWriteFile scriptfile sct
    return actiondir

-- Beware that addition of characters here can lead to lock collisions
-- as is "project/sample" clashes with "project_sample" but ... uncommon case
unsafeCharMap = [('/', '_'),
                 ('\\', '_')]

-- | Remove '/' and '\' from filenames
sanitizePath :: T.Text -> T.Text
sanitizePath = T.map (\x -> fromMaybe x (lookup x unsafeCharMap))

executeLock1OrForAll funcname (NGOList entries) kwargs  = do
    let readSetOrTypeError (NGOReadSet name _) = return name
        readSetOrTypeError _ = throwShouldNotOccur "Expected a readset"
    entries' <- case entries of
        [] -> throwDataError "Cannot run on empty list"
        (NGOString _:_) -> mapM (stringOrTypeError funcname) entries
        (NGOReadSet _ _:_) -> mapM (readSetOrTypeError) entries
        _ -> throwScriptError ("Unsupported type for function " ++ funcname)

    hash <- lookupStringOrScriptError funcname "__hash" kwargs
    tag <- lookupStringOrScriptErrorDef (return "") "collect arguments (hidden tag)"
                (if funcname == "lock1" then "__parallel_tag" else "tag") kwargs
    let prefix
            | T.null tag = ""
            | otherwise = T.unpack tag ++ "-"
    lockdir <- setupHashDirectory prefix "ngless-locks" hash
    -- Keep a map of 'sane -> original' names used for locks to backtrace
    -- what file was locked and return the unsanitized name
    -- See also https://github.com/ngless-toolkit/ngless/issues/68
    let saneentries = sanitizePath <$> entries'
    (e,rk) <- getLock lockdir saneentries
    outputListLno' InfoOutput [funcname, ": Obtained lock file: '", lockdir </> T.unpack e ++ ".lock", "'"]
    reportbase <- setupHashDirectory prefix "ngless-stats" hash
    let reportdir = reportbase </> T.unpack e
    outputListLno' InfoOutput ["Writing stats to '", reportdir, "'"]
    let setReportDir c = c { nConfReportDirectory = reportdir }
    updateNglEnvironment $ \env -> env { ngleConfiguration = setReportDir (ngleConfiguration env) }
    registerHook FinishOkHook $ do
        let receiptfile = lockdir </> T.unpack e ++ ".finished"
        liftIO $ withFile receiptfile WriteMode $ \h -> do
            t <- getZonedTime
            let tformat = "%a %d-%m-%Y %R"
                tstr = formatTime defaultTimeLocale tformat t
            hPutStrLn h (concat ["Finished ", T.unpack e, " at ", tstr])
        release rk
    registerFailHook $ do
        let logfile = lockdir </> T.unpack e ++ ".failed"
        withFile logfile WriteMode $ \h -> do
            hPutStrLn h "Execution failed. Execution log:"
            writeOutputTo h
    case entries of
        (NGOString _:_) -> return . NGOString $! fromMaybe e $ lookup e (zip saneentries entries')
        _ -> case lookup e (zip saneentries entries) of
                Just r -> return r
                Nothing -> throwShouldNotOccur "Could not find entry in map (should not happen)"

executeLock1OrForAll func arg _ = throwScriptError ("Wrong argument for " ++ func ++ " (expected a list of strings, got `" ++ show arg ++ "`")


lockName = (++ ".lock") . T.unpack
finishedName = (++ ".finished") . T.unpack
failedName = (++ ".failed") . T.unpack

-- | Create a lock file
getLock :: FilePath
                -- ^  directory where to create locks
                -> [T.Text]
                -- ^ keys to attempt to lock
                -> NGLessIO (T.Text, ReleaseKey)
getLock basedir fs = do
    existing <- liftIO $ getDirectoryContents basedir
    let notfinished = flip filter fs $ \fname -> finishedName fname `notElem` existing
        notlocked = flip filter notfinished $ \fname -> lockName fname `notElem` existing
        notfailed = flip filter notlocked $ \fname -> failedName fname `notElem` existing
        failed = flip filter notfinished $ \fname -> failedName fname `elem` existing
        locked = flip filter notfinished $ \fname -> lockName fname `elem` existing
    when (null notfinished) $ do
        outputListLno' InfoOutput ["All jobs are finished"]
        throwError $ NGError NoErrorExit "All jobs are finished"

    outputListLno' TraceOutput ["Looking for a lock in '", basedir, "'"]
    outputListLno' TraceOutput [
                    "Total number of tasks to run is ", show (length fs),
                    " (total not finished (including locked & failed): ", show (length notfinished),
                    " ; locked: ", show (length locked),
                    " ; failed: ", show (length failed),
                    ")."]
    -- first try all the tasks that are not locked and have not failed
    -- if that fails, try the locked tasks in the hope that some may be stale
    -- Finally, try the unlocked but failed (in random order)
    getLock' basedir notfailed >>= \case
        Just v -> return v
        Nothing -> do
            outputListLno' InfoOutput ["All tasks locked or failed. Checking for stale locks..."]
            getLock' basedir locked >>= \case
                Just v -> return v
                Nothing -> do
                    when (null failed) $ do
                        outputListLno' InfoOutput ["All jobs appear to be finished or running"]
                        throwError $ NGError NoErrorExit "All jobs are finished or running"
                    -- randomizing the order maximizes the possibilities to get a lock
                    failed' <- liftIO $ shuffleM failed
                    outputListLno' InfoOutput ["All tasks locked or failed and there are no stale locks."]
                    outputListLno' InfoOutput ["Will retry some failed tasks, but it is possible that this will fail again."]
                    outputListLno' InfoOutput ["Failed logs are in directory '", basedir, "'"]
                    getLock' basedir failed' >>= \case
                        Just v -> return v
                        Nothing -> do
                            let msg
                                 | null (notfailed ++ notfailed) = "All jobs are finished"
                                 | null failed = "Jobs appear to be running"
                                 | otherwise = "Jobs are either locked or failed. Check directory '" ++ basedir ++ "' for more information"
                            outputListLno' WarningOutput ["Could get a lock for any file: ", msg]
                            throwError $ NGError NoErrorExit msg

getLock' _ [] = return Nothing
getLock' basedir (f:fs) =
    LockFile.acquireLock LockFile.LockParameters
                        { lockFname = basedir </> lockName f
                        , maxAge = fromInteger (60*60)
                            -- one hour. Given that lock files are touched
                            -- every ten minutes if things are good (see
                            -- thread below), this is an indication that
                            -- the process has crashed
                        , whenExistsStrategy = LockFile.IfLockedNothing
                        , mtimeUpdate = True } >>= \case
        Nothing -> getLock' basedir fs
        Just rk -> do
            isFinished <- liftIO $ doesFileExist (basedir </> finishedName f)
            if isFinished
                then do
                    release rk
                    getLock' basedir fs
                else return $ Just (f, rk)

executeCollect :: NGLessObject -> [(T.Text, NGLessObject)] -> NGLessIO NGLessObject
executeCollect (NGOCounts istream) kwargs = do
    isSubsample <- nConfSubsample <$> nglConfiguration
    current <- case lookup "current" kwargs of
        Nothing -> throwScriptError "current not specified in collect call"
        Just (NGOString c) -> return c
        Just (NGOReadSet n _) -> return n
        Just _ -> throwScriptError "current argument (in collect()) must be a string or a readset"

    allentries <- case lookup "allneeded" kwargs of
        Nothing -> throwScriptError "collect() called without 'allneeded' argument"
        Just (NGOList ells) -> forM ells $ \case
            NGOString s -> return s
            NGOReadSet n _ -> return n
            _ -> throwScriptError "collect() called with 'allneeded' argument, but not all elements are strings or readsets"
        Just _ -> throwScriptError "collect() called with 'allneeded' argument that is not a list"

    ofile <- lookupStringOrScriptError "collect arguments" "ofile" kwargs
    hash <- lookupStringOrScriptError "collect arguments" "__hash" kwargs
    tag <- lookupStringOrScriptErrorDef (return "") "collect arguments (hidden tag)" "__parallel_tag" kwargs
    let prefix
            | T.null tag = ""
            | otherwise = T.unpack tag ++ "-"
    hashdir <- setupHashDirectory prefix "ngless-partials" hash
    (gzfp,gzout) <- openNGLTempFile "compress" "partial." "tsv.gz"
    C.runConduit $
        (snd . asStream $ istream)
        .| CC.concat
        .| CL.map unwrapByteLine
        .| C.unlinesAscii
        .| CAlg.asyncGzipTo gzout
    let partialfile entry = hashdir </> "partial." ++ T.unpack (sanitizePath entry) <.> "tsv.gz"
    outputListLno' TraceOutput ["Collect will write partial file to ", partialfile current]
    liftIO $ do
        hClose gzout
        syncFile gzfp
        moveOrCopy gzfp (partialfile current)
    canCollect <- liftIO $ allM (doesFileExist . partialfile)  (reverse allentries)
                 -- ^ checking in reverse order makes it more likely that ngless notices a missing file early on

    -- It seems wasteful to build the comment string even if `canCollect` is False.
    -- However, these operations are very cheap and provide some basic error checking:
    manualComment <- fmapMaybeM (stringOrTypeError "comment argument to collect() function") (lookup "comment" kwargs)

    autoComments <- case lookup "auto_comments" kwargs of
                        Nothing -> return []
                        Just (NGOList cs) -> mapM (\s -> do
                                                        let errmsg = "auto_comments argument in collect() call"
                                                        symbolOrTypeError errmsg s >>=
                                                            decodeSymbolOrError errmsg
                                                                [("date", AutoDate)
                                                                ,("script", AutoScript)
                                                                ,("hash", AutoResultHash)]) cs
                        _ -> throwScriptError "auto_comments argument to collect() call must be a list of symbols"
    comment <- buildComment manualComment autoComments hash

    if canCollect
        then do
            outputListLno' TraceOutput ["Can collect"]
            newfp <- pasteCounts comment False allentries (map partialfile allentries)
            outputListLno' TraceOutput ["Pasted. Will move result to ", T.unpack ofile]
            moveOrCopyCompress (def
                                    { woCompressLevel = Nothing
                                    , woCanMove = True
                                    , woOFile = T.unpack ofile ++ (if isSubsample then ".subsample" else "")
                                    }) newfp
        else do
            outputListLno' TraceOutput ["Cannot collect (not all files present yet), wrote partial file to ", partialfile current]
            Just lno <- ngleLno <$> nglEnvironment
            registerHook FinishOkHook $
                outputListLno InfoOutput Nothing
                        ["The collect() call at line ", show lno, " could not be executed as there are partial results missing.\n"
                        ,"When you use the parallel module and the collect() function,\n"
                        ,"you typically need to run ngless *multiple times* (once per sample)!\n"
                        ,"\n\n"
                        ,"For more information, see https://ngless.embl.de/stdlib.html#parallel-module"]
    return NGOVoid
executeCollect arg _ = throwScriptError ("collect got unexpected argument: " ++ show arg)

executeSetTag :: NGLessObject -> [(T.Text, NGLessObject)] -> NGLessIO NGLessObject
executeSetTag _ _ = throwShouldNotOccur "set_parallel_tag should have been transformed away!"


-- | split a list into a given number of (roughly) equally sized chunks
nChunks :: Int -- ^ number of chunks
            -> [a] -> [[a]]
nChunks 1 xs = [xs]
nChunks n xs = chunksOf p xs
    where
        p = 1 + (length xs `div` n)

splitAtTab ell = case B.elemIndex 9 ell of -- 9 is TAB
    Nothing -> throwDataError "Line does not have a TAB character"
    Just tix -> return $ B.splitAt tix ell

-- partialPaste pastes a set of inputs, returning both the indices (row
-- headers) and the pasted row content
partialPaste :: [V.Vector B.ByteString] -> NGLess (V.Vector B.ByteString, V.Vector B.ByteString)
partialPaste [] = throwShouldNotOccur "partialPaste called with empty vector"
partialPaste vs
    | not (allSame $ V.length <$> vs) = throwDataError $ "collect(): inputs have differing number of rows"
partialPaste (first_ell:ells) = runST $ do
        indices <- VM.new n
        contents <- VM.new n
        fillData 0 indices contents
    where
        n = V.length first_ell
        splitCheck :: Int -> B.ByteString -> V.Vector B.ByteString -> NGLess B.ByteString
        splitCheck !ix !h e
            | B.isPrefixOf h (e V.! ix) = return $! B.drop (B.length h) (e V.! ix)
            | otherwise = throwDataError $
                                "Inconsistent row index in files for collect() [expected index entry '"++B8.unpack h++"', saw '"++B8.unpack (e V.! ix)++"']."
        fillData !ix indices contents
            | ix == n = do
                indices' <- V.unsafeFreeze indices
                contents' <- V.unsafeFreeze contents
                return . Right $ (indices', contents')
            | otherwise = case splitAtTab (first_ell V.! ix) of
                    Left err -> return $ Left err
                    Right (!h,!c) -> case forM ells (splitCheck ix h) of
                        Left err -> return $ Left err
                        Right cs -> do
                            VM.write indices ix h
                            VM.write contents ix $! B.concat (c:cs)
                            fillData (ix + 1) indices contents

concatPartials :: [(V.Vector B.ByteString, V.Vector B.ByteString)] -> NGLess BL.ByteString
concatPartials [] = throwShouldNotOccur "concatPartials of empty set"
concatPartials groups
    | not (allSame (fst <$> groups)) = throwDataError "indices do not match"
    | otherwise = do
        let contents = snd <$> groups
            header = fst (head groups)
        return . BL.fromChunks $ concatMap (\ix -> (header V.! ix):(map (V.! ix) contents ++ ["\n"])) [0 .. V.length header - 1]

-- | strict variation of sinkTBMQueue
sinkTBMQueue' q shouldClose = do
        C.awaitForever $ \ !v -> liftSTM (TQ.writeTBMQueue q v)
        when shouldClose (liftSTM $ TQ.closeTBMQueue q)
    where
        liftSTM = liftIO . atomically

-- If the number of input files is very large (>1024, typically), we risk
-- hitting the limit on open files by a process, so we work in batches of 512.
maxNrOpenFiles = 512 :: Int


{- Now there is a whole lot of complicated code to efficiently solve the
 - following problem:
 -
 - INPUT 0
 -
 -       h0
 - row1 c01
 - row2 c02
 - row3 c03
 -
 - INPUT 1
 -
 -       h1
 - row0 c11
 - row2 c12
 - row4 c14
 -
 - INPUT 2
 -
 -       h2
 - row0 c21
 - row1 c21
 - row3 c23
 -
 - should produce
 -
 - OUTPUT
 -
 -       h0  h1  h2
 - row0 c00 c10 c20
 - row1 c01 c11 c21
 - row2 c02 c12 c22
 - row3 c03 c13 c23
 - row4 c04 c14 c24
 -
 - Where the missing values (e.g., c00) are assumed to be 0
 -
 - A further complication is that each individual input file may itself have
 - multiple columns.
 -}

data SparseCountData = SparseCountData
                            { spdHeader :: !B.ByteString
                            , _spdIndex :: {-# UNPACK #-} !Int
                            , _spdPayload :: !B.ByteString
                            }
    deriving (Eq)

instance Ord SparseCountData where
    compare (SparseCountData ah ai _) (SparseCountData bh bi _) = case compare ah bh of
        EQ -> compare ai bi
        LT -> LT
        GT -> GT

tagSource :: Int -> C.ConduitT ByteLine SparseCountData NGLessIO ()
tagSource ix = C.awaitForever $ \(ByteLine v) -> case splitAtTab v of
    Left err -> lift $ throwError err
    Right (h, pay) -> C.yield $ SparseCountData h ix pay

-- complete takes a set of SparseCountData and fills in missing columns from
-- the placeholder set.
--
-- Example
--      placeholder = [p_0, p_1, p_2, p_3, p_4]
--      hinput = ...
--      inputs = [h_0 0 l_0, h_2 2 l_2, h_3 3 l_3]
--
--      output = [l_0, p_1, l_2, l_3, p_4]
--
complete :: [B.ByteString] -> (SparseCountData,[SparseCountData]) -> ByteLine
complete placeholders (hinput,inputs) = ByteLine $ B.concat merged
    where
        header = spdHeader hinput
        merged = header:complete' 0 placeholders (hinput:inputs)
        complete' _ [] [] = []
        complete' _ [] (_:_) = error "Logic error in StandardModules/parallel//complete"
        complete' ix (p:ps) xs@(SparseCountData _ ix' pay:rest)
            | ix == ix' = pay:complete' (ix+1) ps rest
            | otherwise = p:complete' (ix+1) ps xs
        complete' _ ps [] = ps

-- Merge input lines by index (first element). The input lines are assumed to
-- be sorted, but not necessary identical (i.e., some may be missing).
mergeCounts :: [C.ConduitT () ByteLine NGLessIO ()] -> C.ConduitT () ByteLine NGLessIO ()
mergeCounts [] = throwShouldNotOccur "Attempt to merge empty sources"
mergeCounts ss = do
        start <- forM ss $ \s -> do
            (s', v) <- lift $ s $$+ (CL.mapM (splitAtTab . unwrapByteLine) .| CC.head)
            case v of
                Nothing -> do
                    lift $ outputListLno' WarningOutput ["Merging empty file"]
                    return (s', placeholder 1)
                Just (_,hs) -> do
                    let p = placeholder (B8.count '\t' hs)
                    return (s', p)
        let (ss', placeholders) = unzip start
            ss'' = map C.unsealConduitT ss'
        CAlg.mergeC [s .| tagSource ix | (s,ix) <- zip ss'' [0..]]
            .| CL.groupOn1 spdHeader
            .| CL.map (complete placeholders)
    where
        placeholder :: Int -> B.ByteString
        placeholder n = B.intercalate "\t" ("":["0" | _ <- [1..n]])


{- There are two modes:
 -
 - 1) The rows match (i.e., row headers are always identical). This is the easy
 - case, and it is like the `paste` command
 -
 - 2) The rows do not match. In this case, NGLess assumes that they are sorted
 - [in "C" locale] and merges them.
 -}
pasteCounts :: [T.Text]
            -- ^ comment text
            -> Bool
            -- ^ whether rows match
            -> [T.Text]
            -- ^ headers
            -> [FilePath]
            -- ^ input files
            -> NGLessIO FilePath
pasteCounts comments matchingRows headers inputs
    | length inputs > maxNrOpenFiles = do
        let current = take maxNrOpenFiles inputs
            currenth = take maxNrOpenFiles headers
            rest = drop maxNrOpenFiles inputs
            resth = drop maxNrOpenFiles headers
        first <- pasteCounts [] matchingRows currenth current
        pasteCounts comments matchingRows (snoc resth $ T.intercalate "\t" currenth) (snoc rest first)
    | otherwise = makeNGLTempFile "collected" "collected.counts." "txt" $ \hout -> do
        C.runConduit (commentC "# " comments .| CB.sinkHandle hout)
        liftIO $ T.hPutStrLn hout (T.intercalate "\t" ("":headers))
        numCapabilities <- liftIO getNumCapabilities
        if matchingRows
            then do
                let sources =
                        [conduitPossiblyCompressedFile f
                            .| CB.lines
                            .| (CC.drop 1 >>
                                C.conduitVector 2048 :: C.ConduitT B.ByteString (V.Vector B.ByteString) (ResourceT IO) ())
                            | f <- inputs]
                    sourcesplits = nChunks numCapabilities sources
                channels <- liftIO $ forM sourcesplits $ \ss -> do
                    ch <- TQ.newTBMQueueIO 4
                    a <- A.async $ C.runConduitRes (C.sequenceSources ss .| CL.map (force . partialPaste) .| sinkTBMQueue' ch True)
                    A.link a
                    return (CA.sourceTBMQueue ch, a)
                C.runConduit $
                    C.sequenceSources (fst <$> channels)
                    .| CAlg.asyncMapEitherC numCapabilities (sequence >=> concatPartials)
                    .| CL.map BB.lazyByteString
                    .| CB.sinkHandleBuilder hout
                forM_ (snd <$> channels) (liftIO . A.wait)
            else C.runConduit $
                    mergeCounts [conduitPossiblyCompressedFile f .|  linesC | f <- inputs]
                    .| byteLineSinkHandle hout


executePaste :: NGLessObject -> [(T.Text, NGLessObject)] -> NGLessIO NGLessObject
executePaste (NGOList ifiles) kwargs = do
    outputListLno' WarningOutput ["Calling __paste which is an internal function, exposed for testing only"]
    ofile <- lookupStringOrScriptError "__paste arguments" "ofile" kwargs
    headers <- lookupStringListOrScriptError "__paste arguments" "headers" kwargs
    matchingRows <- lookupBoolOrScriptErrorDef (return False) "__paste arguments" "matching_rows" kwargs
    ifiles' <- forM ifiles (stringOrTypeError "__concat argument")
    newfp <- pasteCounts [] matchingRows headers (map T.unpack ifiles')
    liftIO $ moveOrCopy newfp (T.unpack ofile)
    return NGOVoid
executePaste _ _ = throwScriptError "Bad call to test function __paste"

lock1 = Function
    { funcName = FuncName "lock1"
    , funcArgType = Just (NGLUnion [NGList NGLString, NGList NGLReadSet])
    , funcArgChecks = []
    , funcRetType = NGLString
    , funcKwArgs = []
    , funcAllowsAutoComprehension = False
    , funcChecks = []
    }

collectFunction isV11 = Function
    { funcName = FuncName "collect"
    , funcArgType = Just NGLCounts
    , funcArgChecks = []
    , funcRetType = NGLVoid
    , funcKwArgs =
        [ArgInformation "current" (not isV11) NGLString []
        ,ArgInformation "allneeded" (not isV11) (NGList NGLString) []
        ,ArgInformation "ofile" True NGLString [ArgCheckFileWritable]
        ,ArgInformation "__can_move" False NGLBool []
        ,ArgInformation "comment" False NGLString []
        ,ArgInformation "auto_comments" False (NGList NGLSymbol) [ArgCheckSymbol ["date", "script", "hash"]]
        ]
    , funcAllowsAutoComprehension = False
    , funcChecks = []
    }

setTagFunction = Function
    { funcName = FuncName "set_parallel_tag"
    , funcArgType = Just NGLString
    , funcArgChecks = []
    , funcRetType = NGLString
    , funcKwArgs = []
    , funcAllowsAutoComprehension = False
    , funcChecks = []
    }


pasteHiddenFunction = Function
    { funcName = FuncName "__paste"
    , funcArgType = Just (NGList NGLString)
    , funcArgChecks = []
    , funcRetType = NGLString
    , funcKwArgs =
        [ ArgInformation "ofile" True NGLString [ArgCheckFileWritable]
        , ArgInformation "headers" True (NGList NGLString) []
        , ArgInformation "matching_rows" False NGLBool []
        ]
    , funcAllowsAutoComprehension = False
    , funcChecks = []
    }



runForAllFunctions =
    [ Function
        { funcName = FuncName "run_for_all"
        , funcArgType = Just (NGList NGLString)
        , funcArgChecks = []
        , funcRetType = NGLString
        , funcKwArgs =
            [ ArgInformation "tag" False NGLString []
            ]
        , funcAllowsAutoComprehension = False
        , funcChecks = []
        }
    , Function
        { funcName = FuncName "run_for_all_samples"
        , funcArgType = Just (NGList NGLReadSet)
        , funcArgChecks = []
        , funcRetType = NGLReadSet
        , funcKwArgs =
            [ ArgInformation "tag" False NGLString []
            ]
        , funcAllowsAutoComprehension = False
        , funcChecks = []
        }
    ]


parallelTransform :: Bool -> [(Int, Expression)] -> NGLessIO [(Int, Expression)]
parallelTransform includeForAll = processRunForAll includeForAll >=> processSetParallelTag >=> addLockHash

addLockHash :: [(Int, Expression)] -> NGLessIO [(Int, Expression)]
addLockHash script = pureTransform addLockHash' script
    where
        addLockHash' :: Expression -> Expression
        addLockHash' (FunctionCall fn@(FuncName fname) expr kwargs block)
            | fname `elem` ["lock1", "run_for_all", "run_for_all_samples"] =
                FunctionCall fn expr ((Variable "__hash", ConstStr h):kwargs) block
                where
                    h = T.pack . MD5.md5s . MD5.Str . show $ map snd script
        addLockHash' e = e

processRunForAll :: Bool -> [(Int, Expression)] -> NGLessIO [(Int, Expression)]
processRunForAll False = checkNoRunForAll
processRunForAll True = processRunForAll' Nothing

processRunForAll' _ [] = return []
processRunForAll' Nothing ((lno,expr):rest) = case expr of
    Assignment v (FunctionCall (FuncName fname) slist kwargs _)
        | fname `elem` ["run_for_all", "run_for_all_samples"] -> do
            let save_match = Assignment (Variable "$parallel$iterator") (Lookup (Just NGLString) v)
                save_list  = Assignment (Variable "$parallel$list") slist
                set_tag = do
                    tag <- lookup (Variable "tag") kwargs
                    return (lno,
                            FunctionCall (FuncName "set_parallel_tag") tag [] Nothing)
            rest' <- processRunForAll' (Just (lno, slist)) rest
            let res = ((lno,expr):(lno,save_match):(lno,save_list):rest')
            case set_tag of
                Nothing -> return res
                Just t -> return (t:res)
    _ -> do
        ((lno,expr):) <$> processRunForAll' Nothing rest
processRunForAll' (Just prev) ((lno,e):rest) = case e of
    Assignment _ (FunctionCall (FuncName fname) _ _ _)
        | fname `elem` ["run_for_all", "run_for_all_samples"] -> do
            throwScriptError ("The functions 'run_for_all'/'run_for_all_samples' can only be called once (seen on lines "++show prev++" and "++show lno++")")
    FunctionCall fn@(FuncName "collect") expr kwargs block -> do
        let kwargs' =  (Variable "allneeded", Lookup (Just NGLString) (Variable "$parallel$list"))
                      :(Variable "current", Lookup (Just NGLString) (Variable "$parallel$iterator"))
                      :kwargs
            e' = FunctionCall fn expr kwargs' block
        rest' <- processRunForAll' (Just prev) rest
        return ((lno,e'):rest')
    _ -> do
        rest' <- processRunForAll' (Just prev) rest
        return ((lno,e):rest')

checkNoRunForAll = mapM checkNoRunForAll1
    where
        checkNoRunForAll1 (_,Assignment _ (FunctionCall (FuncName fname) _ _ _))
            | fname `elem` ["run_for_all", "run_for_all_samples"] =
                throwScriptError ("Function '"++T.unpack fname++"' is only available in parallel module version 1.1+. Please upgrade your import")
        checkNoRunForAll1 e = return e

processSetParallelTag :: [(Int, Expression)] -> NGLessIO [(Int, Expression)]
processSetParallelTag = return . processSetParallelTag' False
    where
        processSetParallelTag' :: Bool -> [(Int, Expression)] -> [(Int, Expression)]
        processSetParallelTag' _ [] = []
        processSetParallelTag' hasTag ((lno, e):rest) = let
                (e',ch) = case e of
                    FunctionCall (FuncName "set_parallel_tag") expr [] Nothing
                            -> (Assignment (Variable "$parallel$tag") expr, True)
                    FunctionCall fn@(FuncName fname) expr kwargs block
                        | hasTag && fname `elem` ["lock1", "collect"]
                            -> (FunctionCall fn expr ((Variable "__parallel_tag", Lookup (Just NGLString) (Variable "$parallel$tag")):kwargs) block, True)
                    _ -> (e, False)
                rest' = processSetParallelTag' (hasTag || ch) rest
            in (lno, e'):rest'


loadModule :: T.Text -> NGLessIO Module
loadModule v
    | v `notElem` ["1.1", "1.0", "0.6"] = throwScriptError ("The behaviour of the parallel module changed.\n"++
                                    "Only versions 1.1/1.0/0.6 are now supported (currently attempting to import version '"++T.unpack v++"')")
    | otherwise = do
        let includeForAll = v == "1.1"
        return def
            { modInfo = ModInfo "stdlib.parallel" v
            , modFunctions =
                [ lock1
                , collectFunction includeForAll
                , setTagFunction
                , pasteHiddenFunction
                ] ++ (if includeForAll then runForAllFunctions else [])
            , modTransform = parallelTransform includeForAll
            , runFunction = \case
                "lock1" -> executeLock1OrForAll "lock1"
                "collect" -> executeCollect
                "set_parallel_tag" -> executeSetTag
                "run_for_all" -> executeLock1OrForAll "run_for_all"
                "run_for_all_samples" -> executeLock1OrForAll "run_for_all_samples"
                "__paste" -> executePaste
                _ -> error "Bad function name"
            }