{-# LANGUAGE DataKinds #-}

-- | The "System.Process.Typed" module from @typed-process@, but with

-- added conduit helpers.

module Data.Conduit.Process.Typed

  ( -- * Conduit specific stuff

    createSink

  , createSource

    -- * Generalized functions

  , withProcess

  , withProcess_

  , withLoggedProcess_

    -- * Reexports

  , module System.Process.Typed

  ) where

import System.Process.Typed hiding (withProcess, withProcess_)

import qualified System.Process.Typed as P

import Data.Conduit (ConduitM, (.|), runConduit)

import qualified Data.Conduit as C

import qualified Data.Conduit.Binary as CB

import Control.Monad.IO.Unlift

import qualified Data.ByteString as S

import System.IO (hClose)

import qualified Data.Conduit.List as CL

import qualified Data.ByteString.Lazy as BL

import Data.IORef (IORef, newIORef, readIORef, modifyIORef)

import Control.Exception (throwIO, catch)

import Control.Concurrent.Async (concurrently)

-- | Provide input to a process by writing to a conduit.

--

-- @since 1.2.1

createSink :: MonadIO m => StreamSpec 'STInput (ConduitM S.ByteString o m ())

createSink =

    (\h -> C.addCleanup (\_ -> liftIO $ hClose h) (CB.sinkHandle h))

    `fmap` createPipe

-- | Read output from a process by read from a conduit.

--

-- @since 1.2.1

createSource :: MonadIO m => StreamSpec 'STOutput (ConduitM i S.ByteString m ())

createSource =

    (\h -> C.addCleanup (\_ -> liftIO $ hClose h) (CB.sourceHandle h))

    `fmap` createPipe

-- | Internal function: like 'createSource', but stick all chunks into

-- the 'IORef'.

createSourceLogged

  :: MonadIO m

  => IORef ([S.ByteString] -> [S.ByteString])

  -> StreamSpec 'STOutput (ConduitM i S.ByteString m ())

createSourceLogged ref =

    -- We do not add a cleanup action to close the handle, since in

    -- withLoggedProcess_ we attempt to read from the handle twice

    (\h ->

       (  CB.sourceHandle h

       .| CL.iterM (\bs -> liftIO $ modifyIORef ref (. (bs:))))

    )

    `fmap` createPipe

-- | Same as 'P.withProcess', but generalized to 'MonadUnliftIO'.

--

-- @since 1.2.1

withProcess

  :: MonadUnliftIO m

  => ProcessConfig stdin stdout stderr

  -> (Process stdin stdout stderr -> m a)

  -> m a

withProcess pc f = withRunInIO $ \run -> P.withProcess pc (run . f)

-- | Same as 'P.withProcess_', but generalized to 'MonadUnliftIO'.

--

-- @since 1.2.1

withProcess_

  :: MonadUnliftIO m

  => ProcessConfig stdin stdout stderr

  -> (Process stdin stdout stderr -> m a)

  -> m a

withProcess_ pc f = withRunInIO $ \run -> P.withProcess_ pc (run . f)

-- | Run a process, throwing an exception on a failure exit code. This

-- will store all output from stdout and stderr in memory for better

-- error messages. Note that this will require unbounded memory usage,

-- so caveat emptor.

--

-- This will ignore any previous settings for the stdout and stderr

-- streams, and instead force them to use 'createSource'.

--

-- @since 1.2.3

withLoggedProcess_

  :: MonadUnliftIO m

  => ProcessConfig stdin stdoutIgnored stderrIgnored

  -> (Process stdin (ConduitM () S.ByteString m ()) (ConduitM () S.ByteString m ()) -> m a)

  -> m a

withLoggedProcess_ pc inner = withUnliftIO $ \u -> do

  stdoutBuffer <- newIORef id

  stderrBuffer <- newIORef id

  let pc' = setStdout (createSourceLogged stdoutBuffer)

          $ setStderr (createSourceLogged stderrBuffer) pc

  P.withProcess pc' $ \p -> do

    a <- unliftIO u $ inner p

    let drain src = unliftIO u (runConduit (src .| CL.sinkNull))

    ((), ()) <- drain (getStdout p) `concurrently`

                drain (getStderr p)

    checkExitCode p `catch` \ece -> do

      stdout <- readIORef stdoutBuffer

      stderr <- readIORef stderrBuffer

      throwIO ece

        { eceStdout = BL.fromChunks $ stdout []

        , eceStderr = BL.fromChunks $ stderr []

        }

    return a