{-# LANGUAGE RankNTypes #-}
module Data.Conduit.Process (
  -- * Run process
  sourceProcess,
  sourceProcessE,
  OutAndErr,
  PConduitInp,
  PConduitE,
  conduitProcessE,

  -- * Run shell command
  sourceCmd,
  sourceCmdE,
  conduitCmd,

  -- * Conduit manipulation
  onStdOut,
  onStdErr,

  -- * Convenience re-exports
  shell,
  proc,
  CreateProcess(..),
  CmdSpec(..),
  StdStream(..),
  ProcessHandle,
  ) where

import qualified Control.Exception as E
import Control.Monad
import Control.Monad.Trans
import Control.Concurrent (forkIO)
import Control.Concurrent.Chan
import qualified Data.ByteString as S
import qualified Data.ByteString.Char8 as SB
import Data.Conduit
import qualified Data.Conduit.List as CL
import Data.Void
import System.Exit (ExitCode(..))
import System.IO
import System.Process

bufSize :: Int
bufSize = 64 * 1024


-- | Both stdout and stderr are provided by the OutAndErr type
type OutAndErr = Either S.ByteString S.ByteString

-- | The type of a process conduit providing OutAndErr and also the ExitCode result
type PConduitE m = Pipe Void () OutAndErr () m ExitCode

-- | The type of a process conduit input pipe
type PConduitInp = Pipe Void () S.ByteString () IO ()


-- internal channel data
data MiddleChan = MidOut S.ByteString | MidErr S.ByteString | MidOutEnd | MidErrEnd


conduitProcessE :: MonadIO m => PConduitInp -> CreateProcess -> PConduitE m
conduitProcessE inpPipe cp =
    do (Just cin, Just cout, Just cerr, ph) <- liftIO createp
       mChan <- liftIO newChan
       liftIO $ do void $ forkIO $ runPipe $ readFrom MidOutEnd MidOut cout mChan
                   void $ forkIO $ runPipe $ readFrom MidErrEnd MidErr cerr mChan
                   void $ forkIO $ runPipe $ inpPipe >+> sinkInp cin mChan
       chanSrcMid ph mChan
  where
    createp = createProcess cp
      { std_in  = CreatePipe
      , std_out = CreatePipe
      , std_err = CreatePipe
      }

    readFrom e i d c = readRet e i c d >> liftIO (hClose d)
    readRet e i c d = do rd <- liftIO $ S.hGetSome d bufSize
                         if S.null rd then liftIO $ writeChan c e else
                             do liftIO $ writeChan c $ i rd
                                readRet e i c d

    sinkInp d c = await >>= \r ->
                  case r of
                    Just v  -> (liftIO $ S.hPut d v >> hFlush d) >> sinkInp d c
                    Nothing -> liftIO $ hClose d

    -- chanSrcMid will wait for and complete with the process exit
    -- code after both stderr and stdout are closed (which the process
    -- must do to exit).  This does not wait for stdin to close since
    -- that can be dealt with asynchronously, but it does completely
    -- drain the process' output and error streams before completing.
    chanSrcMid ph c = chanSrcMd' ph 2 c
    chanSrcMd' ph 0 _ = liftIO $ waitForProcess' ph
    chanSrcMd' ph n c = do rd <- liftIO $ readChan c
                           case rd of
                              MidOut s -> yield (Right s) >> chanSrcMd' ph n c
                              MidErr s -> yield (Left  s) >> chanSrcMd' ph n c
                              MidOutEnd -> chanSrcMd' ph (n-1) c
                              MidErrEnd -> chanSrcMd' ph (n-1) c

    waitForProcess' ph =
      waitForProcess ph `E.catch` \(E.SomeException _) -> return ExitSuccess


-- | Source of process
sourceProcess :: MonadResource m => CreateProcess -> Pipe Void () S.ByteString () m ()
sourceProcess cp = conduitProcessE CL.sourceNull cp >+> oobErr >+> nullRval


oobErr :: MonadResource m => forall l r. Pipe l OutAndErr S.ByteString r m r
oobErr = awaitForever $ either showErr passOut
showErr, passOut :: MonadResource m =>
                    forall l r. S.ByteString -> Pipe l OutAndErr S.ByteString r m r
showErr e = liftIO (SB.hPutStrLn stderr e) >> oobErr
passOut o = yield o >> oobErr
nullRval :: MonadResource m => forall l i u. Pipe l i i u m ()
nullRval  = await >>= maybe (return ()) (\v -> yield v >> nullRval)


-- | Source of process with both stdout (Right) and stderr (Left) output and exit status
sourceProcessE :: MonadResource m => CreateProcess -> PConduitE m
sourceProcessE = conduitProcessE CL.sourceNull

-- | Conduit of shell command
conduitCmd :: MonadResource m => String -> Pipe Void () S.ByteString () m ()
conduitCmd scmd = conduitProcessE CL.sourceNull (shell scmd) >+> oobErr >+> nullRval


-- | Source of shell command
sourceCmd :: MonadResource m => String -> Pipe Void () S.ByteString () m ()
sourceCmd = sourceProcess . shell

-- | Source of shell command providing output, error, and exit status
sourceCmdE :: MonadResource m => String -> PConduitE m
sourceCmdE = sourceProcessE . shell


-- | onStdOut can be used on conduitProcessE output to specify a set
-- of conduit operations to be performed only on the output (Right)
-- stream, leaving the error (Left) stream untouched.
onStdOut :: MonadIO m => forall l r e .
            Pipe Void p q () m ()
         -> Pipe l (Either e p) (Either e q) r m r
onStdOut inner = awaitForever $ \t ->
                 case t of
                   Left e -> yield $ Left e
                   Right o -> yield o
                              >+> inner
                              >+> awaitForever (yield . Right)


-- | onStdErr can be used on conduitProcessE output to specify a set
-- of conduit operations to be performed only on the error (Left)
-- stream, leaving the output (Right) stream untouched.
onStdErr :: MonadIO m => forall l r o.
            Pipe Void p q () m ()
         -> Pipe l (Either p o) (Either q o) r m r
onStdErr inner = awaitForever $ \t ->
                 case t of
                   Right o -> yield $ Right o
                   Left e -> yield e
                             >+> inner
                             >+> awaitForever (yield . Left)