content: {- Copyright 2016-2025 NGLess Authors
   - License: MIT
   -}
  {-# LANGUAGE FlexibleContexts, RankNTypes #-}
  
  module StandardModules.Mappers.Soap
      ( hasValidIndex
      , createIndex
      , callMapper
      ) where
  
  import System.Process
  import System.IO
  import System.Exit
  import System.Directory
  import           Control.Monad (forM_)
  import           Control.Monad.IO.Class (liftIO)
  import qualified Data.ByteString as B
  import qualified Data.ByteString.Char8 as B8
  import qualified Data.ByteString.Lazy.Char8 as BL8
  
  import qualified UnliftIO.Async as A
  import qualified Control.Concurrent.STM.TBMQueue as TQ
  
  import qualified Data.Conduit.Binary as CB
  import qualified Data.Conduit.TQueue as CA
  import qualified Data.Conduit.Process as CP
  import qualified Data.Conduit.List as CL
  import qualified Data.Conduit as C
  import qualified UnliftIO as U
  import           Data.Conduit ((.|))
  import           Data.Conduit.Algorithms.Utils (awaitJust)
  import           Data.Conduit.Algorithms.Async (conduitPossiblyCompressedFile)
  import           Control.Monad.Extra (guard, allM, whenM)
  import           GHC.Conc (getNumCapabilities, setNumCapabilities)
  import           Data.List (isSuffixOf)
  import           Control.Monad.Except (MonadError)
  import           Control.Monad.Trans.Resource
  
  import Output
  import NGLess
  import Data.FastQ
  import Utils.Conduit (ByteLine(..), linesC)
  import FileManagement
  import Configuration
  import NGLess.NGLEnvironment
  import Utils.Utils (dropEnd, safeHead)
  
  hasValidIndex :: FilePath -> NGLessIO Bool
  hasValidIndex basepath = liftIO $ flip allM indexRequiredFormats $ \ext -> doesFileExist (basepath' ++ ext)
      where
          basepath'
              | ".gz" `isSuffixOf` basepath = dropEnd 3 basepath
              | otherwise = basepath
          indexRequiredFormats =
                  [".index.amb"
                  ,".index.ann"
                  ,".index.bwt"
                  ,".index.fmv"
                  ,".index.hot"
                  ,".index.lkt"
                  ,".index.pac"
                  ,".index.rev.bwt"
                  ,".index.rev.fmv"
                  ,".index.rev.lkt"
                  ,".index.rev.pac"
                  ,".index.sa"
                  ,".index.sai"
                  ]
  
  
  createIndex :: FilePath -> NGLessIO ()
  createIndex fafile = do
      outputListLno' InfoOutput ["Start SOAP index creation for ", fafile]
      fafile' <- if ".gz" `isSuffixOf` fafile
                      then do
                          let gunzipped = dropEnd 3 fafile
                          outputListLno' WarningOutput ["SOAP indexing does not work on gzipped files directly. Gunzipping '", fafile, "'..."]
                          whenM (liftIO $ doesFileExist gunzipped) $
                              throwDataError ("SOAP indexing does not work on gzipped files (got argument '" ++ fafile ++ "' and gunzipped version already exists (so refusing to overwrite).")
                          C.runConduit $ conduitPossiblyCompressedFile fafile .| CB.sinkFile gunzipped
                          return gunzipped
                      else return fafile
      (exitCode, out, err) <- liftIO $
          readProcessWithExitCode "2bwt-builder" [fafile'] []
      outputListLno' DebugOutput ["SOAP-index stderr: ", err]
      outputListLno' DebugOutput ["SOAP-index stdout: ", out]
      case exitCode of
          ExitSuccess -> return ()
          ExitFailure _err -> throwSystemError err
  
  callMapper :: FilePath -> ReadSet -> [String] -> C.ConduitT B.ByteString C.Void NGLessIO a -> NGLessIO a
  callMapper refIndex (ReadSet rs1 rs2) extraArgs outC = do
      outputListLno' InfoOutput ["Starting mapping to ", refIndex]
      numCapabilities <- liftIO getNumCapabilities
      strictThreads <- nConfStrictThreads <$> nglConfiguration
      let fps1 = map (\(FastQFilePath _ a, FastQFilePath _ b) -> [a, b]) rs1
          fps2 = map (\(FastQFilePath _ a) -> [a]) rs2
          with1Thread act
              | strictThreads = U.bracket_
                                  (liftIO $ setNumCapabilities 1)
                                  (liftIO $ setNumCapabilities numCapabilities)
                                  act
              | otherwise = act
          soapThreads
              | strictThreads && numCapabilities > 1 = numCapabilities - 1
              | otherwise = numCapabilities
          soapPath = "soap2.21"
  
      -- We are going to create a thread which simply consumes input and feeds it
      -- to `outC`. Then we feed it both the SAM header and the result of
      -- soup2sam for both runs.  This is possibly not the simplest way to
      -- achieve this goal, but without having a resumable sink, I do not see
      -- exactly how to get it all to work.
      q <- liftIO $ TQ.newTBMQueueIO 4
      out <- A.async (C.runConduit $ CA.sourceTBMQueue q .| outC)
      C.runConduitRes $
          makeSAMHeader refIndex
              .| CA.sinkTBMQueue q
      forM_ (filter (not . null) (fps1 ++ fps2)) $ \fps -> do
          (rk,(otemp,htemp)) <- openNGLTempFile' refIndex "map_" "soap"
          (rk2,(otemp2,htemp2)) <- openNGLTempFile' refIndex "map2_" "soap"
          liftIO $ hClose htemp
          liftIO $ hClose htemp2
          let fps' = case fps of
                          [fp] -> ["-a", fp]
                          [fp, fp'] -> ["-a", fp, "-b", fp', "-2", otemp2]
                          _ -> error "SOAP Multiple errors"
              cmdargs =  concat [fps', ["-p", show soapThreads, "-D", refIndex ++ ".index", "-o", otemp], extraArgs]
          outputListLno' TraceOutput ["Calling binary ", soapPath, " with args: ", unwords cmdargs]
          let cp = proc soapPath cmdargs
          (exitCode, (), err) <- liftIO $ with1Thread $
                  CP.sourceProcessWithStreams cp
                      (return ()) -- stdin
                      (return ()) -- stdout
                      CL.consume -- stderr
          outputListLno' DebugOutput ["SOAP info: ", BL8.unpack $ BL8.fromChunks err]
          case exitCode of
              ExitFailure code ->
                  throwSystemError $ concat ["Failed mapping\nCommand line was::\n\t",
                                  soapPath, " ", unwords cmdargs, "\n",
                                  "SOAP error code was ", show code, "."]
              ExitSuccess -> return ()
          outputListLno' InfoOutput ["Done mapping to ", refIndex, ". Converting to SAM..."]
          let cmdargs' =  ["-p" | length fps' == 4]
              soap2samPath = "soap2sam.pl"
              soup2sam = proc soap2samPath cmdargs'
          (exitCode', (), err') <- liftIO $
              withFile otemp ReadMode $ \otemph ->
                  withFile otemp2 ReadMode $ \otemp2h ->
                      CP.sourceProcessWithStreams soup2sam
                          (CB.sourceHandle otemph >> CB.sourceHandle otemp2h) -- stdin
                          (C.toConsumer $ CA.sinkTBMQueue q) -- stdout
                          (CL.consume :: C.ConduitT B.ByteString C.Void IO [B.ByteString])
          release rk
          release rk2
          case exitCode' of
              ExitSuccess -> do
                  outputListLno' InfoOutput ["Done soap2sam."]
              ExitFailure code ->
                  throwSystemError $ concat ["Failed conversion to SAM\nCommand line was::\n\t",
                                  soap2samPath, " ", unwords cmdargs', "\n",
                                  "SOAP2SAM error code was ", show code, ".\n",
                                  "Error output: ", B8.unpack (B8.intercalate "\n\t" err')]
  
      liftIO $ do
          U.atomically (TQ.closeTBMQueue q)
          A.wait out
  
  makeSAMHeader :: (MonadResource m, MonadUnliftIO m, MonadThrow m, MonadError NGError m) => FilePath -> C.ConduitT () B.ByteString m ()
  makeSAMHeader fafile = conduitPossiblyCompressedFile fafile .| linesC .| asSamHeader
      where
          -- asSamHeader :: C.Conduit ByteLine IO B.ByteString
          asSamHeader = awaitJust $ \(ByteLine line) -> case readId line of
                                                  Just rid -> processRead rid 0
                                                  Nothing -> return ()
          -- processRead :: B.ByteString -> Int -> C.Conduit ByteLine IO B.ByteString
          processRead rid !n = C.await >>= \case
              Nothing -> emit rid n
              Just (ByteLine line) -> case readId line of
                  Just rid' -> do
                      emit rid n
                      processRead rid' 0
                  Nothing -> processRead rid (n + B.length line)
          readId :: B.ByteString -> Maybe B.ByteString
          readId line = do
              (mark,rest) <- B8.uncons line
              guard (mark == '>')
              safeHead $ B8.split ' ' rest
          emit rid n = C.yield $ B.concat ["@SQ\tSN:", rid, "\t", B8.pack (show n), "\n"]
  