{-# LANGUAGE ScopedTypeVariables #-}

module System.IO.Capture (
  -- * As usual, You specify import only this function.
  capture

  -- * If you want to use lazy reading in the action, I recomend you use these.
, getContents
, hGetContents
, readFile
) where

import Prelude hiding (getContents,readFile,catch)
import System.IO hiding (getContents,hGetContents,readFile)
import qualified System.IO.Strict as SIO (getContents,hGetContents,readFile)
    -- in package strict (not in strict-io)
import Control.Exception
import Control.Applicative
import Control.Monad
import System.Directory
import System.Posix

{-|

  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","")

  WARNING: If Lazy Reading such as @getContents@ was contained in the
  action, its behavior is very strange. For detail, see
  tests/Tests.hs.

 -}

capture :: IO a -> String -> IO (String,String)
capture action givenInput = 
  withTempFile $ \(outh,outpath) ->
  withTempFile $ \(errh,errpath) -> ignoringException ("","") $ do
    let stds =  [stdin,stdout,stderr]
    stdbufs <- mapM hGetBuffering stds
    mapM_ (flip hSetBuffering NoBuffering) stds
    (inr,inw) <- createPipe
    fd_in  <- dup 0
    fd_out <- dup 1
    fd_err <- dup 2
    let capture' = do
          inr `dupTo` 0
          flip dupTo 1 =<< handleToFd outh
          flip dupTo 2 =<< handleToFd errh
          fdWrite inw givenInput
          closeFd inw
          ignoringException () (action >> return ())
          (,) <$> readFile outpath <*> readFile errpath
        recover = do
          fd_in  `dupTo` 0
          fd_out `dupTo` 1
          fd_err `dupTo` 2
          forM_ [inr,inw,fd_in,fd_out,fd_err] $ \fd -> ignoringException () $ closeFd fd
          zipWithM_ ((ignoringException () .) . hSetBuffering) stds stdbufs
    capture' `finally` recover

-- override for forcing strictness

-- if exception was thrown, they returns empty string.

readFile     :: FilePath -> IO String
getContents  :: IO String
hGetContents :: Handle -> IO String

readFile     = ignoringException "" . SIO.readFile
getContents  = ignoringException "" SIO.getContents
hGetContents = ignoringException "" . SIO.hGetContents

-- utilities

withTempFile :: ((Handle,FilePath) -> IO a) -> IO a
withTempFile action = do
  (path,hdl) <- flip openTempFile "tmpXXXXX" =<< getTemporaryDirectory
  action (hdl,path) `finally` (hClose hdl >> removeFile path)

ignoringException :: a -> IO a -> IO a
ignoringException value action = action `catch` do \(_::SomeException) -> return value