{-# LANGUAGE ScopedTypeVariables #-}

module System.IO.Capture (
  -- * Capturing std(out|err)
  capture
) where

import Prelude hiding (catch)
import System.IO
import Control.Exception
import Control.Applicative
import Control.Monad
import System.Posix
import System.Posix.IO

{-|

  Takes an IO as action to run, and a String as given stdin,
  then returns whole stdout and stderr as String of tuple.

  > import System.IO.Capture (capture)
  >
  > main = print =<< capture (getLine >>= putStr) "foobar"
  >   -- prints ("foobar","")

 -}

capture :: IO a -> String -> IO (String,String)
capture action givenInput = do
  (p,o,e) <- forkHandle (action >> return ()) givenInput
  (o_str,e_str) <- (,) <$> hGetContents o <*> hGetContents e
  getProcessStatus True True p -- >>= hPutStrLn stderr . show
  return (o_str,e_str)

forkHandle :: IO () -> String -> IO (ProcessID,Handle,Handle)
forkHandle action givenInput = do
  (inr,inw)   <- createPipe
  (outr,outw) <- createPipe
  (errr,errw) <- createPipe
  pid <- forkProcess $ do
           zipWithM_ dupTo [inr,outw,errw] [0,1,2]
           inh <- fdToHandle inw
           hPutStr inh givenInput >> hClose inh
           catchSome action $ hPutStrLn stderr . ("*** Exception: "++) . show
           hClose stdout >> hClose stderr
           mapM_ closeFd [inr,outr,errr]
  mapM_ closeFd [inr,inw,outw,errw]
  [outh,errh] <- mapM fdToHandle [outr,errr]
  return (pid,outh,errh)

catchSome :: IO a -> (SomeException -> IO a) -> IO a
catchSome = catch