{- Copyright 2015-2019 NGLess Authors
 - License: MIT
 -}

{-# LANGUAGE FlexibleContexts, OverloadedStrings #-}

module StandardModules.Samtools
    ( loadModule
    ) where

import Control.Monad.Trans.Resource
import Control.Exception
import Control.Concurrent
import Control.Monad

import qualified Data.Conduit as C
import           Data.Conduit ((.|))
import qualified Data.Text as T
import qualified Control.Concurrent.Async as A
import           Data.List (uncons, tails)
import           Data.Maybe (mapMaybe)


import System.FilePath
import System.Process
import System.Exit
import System.IO

import Control.Monad.IO.Class (liftIO)
import Data.Default (def)

import FileManagement
import FileOrStream
import Transform
import Language
import Modules
import Output
import NGLess

import Utils.Conduit (byteLineVSinkHandle)
import Utils.Utils (passthrough)


headtails :: [a] -> [(a,[a])]
headtails = mapMaybe uncons . tails

-- This function is necessary because we cannot call readProcessWithErrorCode directly (see comment below)
callSamtools :: FileOrStream -> [String] -> Handle -> NGLessIO ExitCode
callSamtools istream cmdargs hout = do
    samtoolsPath <- samtoolsBin
    outputListLno' TraceOutput ["Calling binary ", samtoolsPath, " with args: ", unwords cmdargs]
    let (_, istream') = asSamStream istream
        cp = (proc samtoolsPath cmdargs) { std_in = CreatePipe, std_out = UseHandle hout, std_err = CreatePipe }
    (Just pipe_out, Nothing, Just herr, jHandle) <- liftIO $ createProcess cp
    samP <- liftIO . A.async $ waitForProcess jHandle
    errP <- liftIO . A.async $ do
                -- In a separate thread, consume all the error input.
                -- the same pattern is used in the implementation of
                -- readProcessWithErrorCode (which cannot be used here
                -- as we want control over stdin/stdout)
                err <- hGetContents herr
                void (evaluate (length err))
                hClose herr
                return err
    C.runConduit $ istream' .| byteLineVSinkHandle pipe_out
    (err,exitCode) <- liftIO $ do
        hClose pipe_out
        A.waitBoth errP samP
    outputListLno' DebugOutput ["Samtools err output: ", err]
    liftIO $ hClose hout
    return exitCode

executeSort :: NGLessObject -> KwArgsValues -> NGLessIO NGLessObject
executeSort (NGOMappedReadSet name istream rinfo) args = do
    outBam <- lookupBoolOrScriptErrorDef (return False) "samtools_sort" "__output_bam" args
    sortByName <- lookupSymbolOrScriptErrorDef (return "coordinate") "arguments to samtools_sort" "by" args >>= \case
        "coordinate" -> return False
        "name" -> return True
        other -> throwShouldNotOccur ("Check failed. No samtool_sort option: " ++ T.unpack other)
    let oformat = if outBam then "bam" else "sam"
        fname = case istream of
                    File f -> f
                    Stream _ f _ -> f
    (rk, (newfp, hout)) <- openNGLTempFile' fname "sorted_" oformat
    (trk, tdirectory) <- createTempDir "samtools_sort_temp"

    numCapabilities <- liftIO getNumCapabilities
    let cmdargs = ["sort"] <> ["-n" | sortByName] <> ["-@", show numCapabilities, "-O", oformat, "-T", tdirectory </> "samruntmp"]
    exitCode <- callSamtools istream cmdargs hout
    release trk
    case exitCode of
        ExitSuccess -> do
            outputListLno' InfoOutput ["Done samtools sort"]
            return (NGOMappedReadSet name (File newfp) rinfo)
        ExitFailure code -> do
            release rk
            samtoolsPath <- samtoolsBin
            throwSystemError $ concat ["Failed samtools sort\n",
                            "Executable used::\t", samtoolsPath,"\n",
                            "Command line was::\n\t", unwords cmdargs, "\n",
                            "Samtools exit code was ", show code, "."]
executeSort _ _ = throwScriptError "Unexpected arguments for samtools_sort function"

sortOFormat :: [(Int, Expression)] -> NGLessIO [(Int, Expression)]
sortOFormat = return . sortOFormat'

-- If samtools_sort is called and its return value is only used in a write()
-- call which is in BAM format, then we should immediately use BAM as the
-- output format.
sortOFormat' :: [(Int, Expression)] -> [(Int, Expression)]
sortOFormat' [] = []
sortOFormat' ((lno,e):es) = (lno,e'):sortOFormat' es
    where
        e' = case e of
            Assignment v (FunctionCall fname@(FuncName "samtools_sort") expr args Nothing)
                | outputBam v es -> Assignment v (FunctionCall fname expr ((Variable "__output_bam", ConstBool True):args) Nothing)
            _ -> e
        outputBam _ [] = False
        outputBam v ((_,c):rest) = case c of
            FunctionCall (FuncName "write") (Lookup _ v') args Nothing
                | v == v' -> isOBam args && not (isVarUsed v rest)
            _
                | isVarUsed1 v c -> False
                | otherwise -> outputBam v rest
        -- The rules are
        --  1. if a format argument is present, it takes priority;
        --     else, infer from ofile argument
        --  2. In doubt, False is the safe option
        isOBam args = case oFormat args of
            Just "bam" -> True
            Just _ -> False
            Nothing -> case lookup (Variable "ofile") args of
                Nothing -> False
                Just oname -> stringWillEndWith oname ".bam" == Just True
        oFormat :: [(Variable, Expression)] -> Maybe String
        oFormat args = case lookup (Variable "format") args of
            Just (ConstSymbol "bam") -> Just "bam"
            Just _ -> Just "?"
            Nothing -> Nothing

        stringWillEndWith :: Expression -> T.Text -> Maybe Bool
        stringWillEndWith (ConstStr b) post
            | T.length b >= T.length post = Just $ T.takeEnd (T.length post) b == post
            | otherwise = Nothing
        stringWillEndWith (BinaryOp BOpAdd _ left) post = stringWillEndWith left post
        stringWillEndWith _ _ = Nothing


checkUnique :: [(Int, Expression)] -> NGLessIO [(Int, Expression)]
checkUnique = passthrough $ \sc ->
        forM_ (headtails $ reverse sc) $ \((lno, e), rest) -> case e of
            Assignment _ (FunctionCall (FuncName "select") (Lookup _ v) kwargs Nothing)
                        | selectsUnique kwargs -> checkSorted lno v (map snd rest)
            _ -> return ()
    where
        checkSorted _ _ [] = return ()
        checkSorted lno v (r:rs) = case r of
            Assignment v' expr
                | v == v' -> case expr of
                    FunctionCall (FuncName "samtools_sort") _ _ _ ->
                            throwScriptError ("Cannot select unique reads from a sorted mappedreadset (at line " ++ show lno ++ ").\n\tConsider selecting, *then* sorting.")
                    Lookup _ v'' -> checkSorted lno v'' rs
                    _ -> return ()
            _ -> checkSorted lno v rs
        selectsUnique :: [(Variable, Expression)] -> Bool
        selectsUnique = any $ \(Variable k,v) -> k == "keep_if" && hasUnique v
        hasUnique (ConstSymbol "unique") = True
        hasUnique (ListExpression vs) = any hasUnique vs
        hasUnique _ = False

samtools_sort_function = Function
    { funcName = FuncName "samtools_sort"
    , funcArgType = Just NGLMappedReadSet
    , funcArgChecks = []
    , funcRetType = NGLMappedReadSet
    , funcKwArgs = [ArgInformation "by" False NGLSymbol [ArgCheckSymbol ["coordinate", "name"]]]
    , funcAllowsAutoComprehension = False
    , funcChecks = []
    }


executeView :: NGLessObject -> KwArgsValues -> NGLessIO NGLessObject
executeView (NGOMappedReadSet name istream rinfo) args = do
    outBam <- lookupBoolOrScriptErrorDef (return False) "samtools_view" "__output_bam" args
    bedFile <- lookupStringOrScriptError "coordinates in BED format" "bed_file" args

    let oformat = if outBam then "bam" else "sam"
        fname = case istream of
                    File f -> f
                    Stream _ f _ -> f

    (rk, (newfp, hout)) <- openNGLTempFile' fname "subset_" oformat

    numCapabilities <- liftIO getNumCapabilities
    let cmdargs = ["view", "-h", "-@", show numCapabilities, "-O", oformat, "-L", T.unpack bedFile]
    exitCode <- callSamtools istream cmdargs hout
    case exitCode of
        ExitSuccess -> do
            outputListLno' InfoOutput ["Done samtools view"]
            return (NGOMappedReadSet name (File newfp) rinfo)
        ExitFailure code -> do
            release rk
            samtoolsPath <- samtoolsBin
            throwSystemError $ concat ["Failed samtools view\n",
                            "Executable used::\t", samtoolsPath,"\n",
                            "Command line was::\n\t", unwords cmdargs, "\n",
                            "Samtools exit code was ", show code, "."]
executeView _ _ = throwScriptError "Unexpected arguments for samtools_view function"

samtools_view_function = Function
    { funcName = FuncName "samtools_view"
    , funcArgType = Just NGLMappedReadSet
    , funcArgChecks = []
    , funcRetType = NGLMappedReadSet
    , funcKwArgs = [ArgInformation "bed_file" True NGLString [ArgCheckFileReadable]]
    , funcAllowsAutoComprehension = False
    , funcChecks = []
    }

loadModule :: T.Text -> NGLessIO Module
loadModule v
    | v == "0.0" = return def
                   { modInfo = ModInfo "stdlib.samtools" "0.0"
                   , modCitations = [citation]
                   , modFunctions = [samtools_sort_function]
                   , modTransform = sortOFormat >=> checkUnique
                   , runFunction = \case
                                   "samtools_sort" -> executeSort
                                   other -> \_ _ -> throwShouldNotOccur ("samtools runction called with wrong arguments: " ++ show other)
                   }
    | v `elem` ["1.0", "0.1"] = return def
                   { modInfo = ModInfo "stdlib.samtools" "1.0"
                   , modCitations = [citation]
                   , modFunctions = [samtools_sort_function, samtools_view_function]
                   , modTransform = sortOFormat >=> checkUnique
                   , runFunction = \case
                                   "samtools_sort" -> executeSort
                                   "samtools_view" -> executeView
                                   other -> \_ _ -> throwShouldNotOccur ("samtools runction called with wrong arguments: " ++ show other)
                                   }
    | otherwise = throwScriptError ("samtools module version " ++ T.unpack v ++ " doesn't exist")
    where
        citation = T.concat
            ["'The Sequence Alignment/Map format and SAMtools' "
            ,"by Heng Li, Bob Handsaker, Alec Wysoker, Tim Fennell, Jue Ruan, Nils Homer, Gabor Marth, Goncalo Abecasis, Richard Durbin, and 1000 Genome Project Data Processing Subgroup "
            ,"in Bioinformatics (2009) 25 (16): 2078-2079 (2009) DOI:10.1093/bioinformatics/btp352"]