{-# LANGUAGE FlexibleContexts, RankNTypes, ScopedTypeVariables, TypeFamilies #-}
{-# OPTIONS_GHC -fno-warn-name-shadowing #-}
module System.Process.Read.Convenience
    ( -- * Predicates
      isResult
    , isStdout
    , isStderr
    , isOutput
    , isException
    , isHandle
    -- * Filters
    , discardStdout
    , discardStderr
    , discardOutput
    , discardExceptions
    , discardResult
    , keepStdout
    , keepStderr
    , keepOutput
    , keepExceptions
    , keepResult
    -- * Transformers
    , mergeToStdout
    , mergeToStderr
    , mapMaybeResult
    , mapMaybeStdout
    , mapMaybeStderr
    , mapMaybeException
    -- * Collectors
    , collectOutputs
    -- * IO operations
    , ePutStr
    , ePutStrLn
    , eMessage
    , eMessageLn

    , foldException
    , foldChars
    , foldStdout
    , foldStderr
    , foldResult
    , foldSuccess
    , foldFailure
    , foldFailure'

    , doException
    , doOutput
    , doStdout
    , doStderr
    , doExit
    , doAll

    , dots
    , prefixed
    ) where

import Control.Exception (throw)
import Control.Monad.Trans (MonadIO, liftIO)
import Data.ListLike (ListLike(empty, append, concat, span, tail, singleton, null), ListLikeIO(hPutStr))
import Data.Maybe (mapMaybe)
import Prelude hiding (length, rem, concat, span, tail, null)
import System.Exit (ExitCode(..), exitWith)
import System.IO (stdout, stderr)
import qualified System.IO as IO (hPutStr, hPutStrLn)
import System.Process (ProcessHandle)
import System.Process.ListLike (ListLikePlus(..))
import System.Process.Chunks (Chunk(..), foldChunks, foldChunk, putDots)

isHandle :: ListLikePlus a c => Chunk a -> Bool
isHandle = foldChunk (const True) (const False) (const False) (const False) (const False)
isResult :: ListLikePlus a c => Chunk a -> Bool
isResult = foldChunk (const False) (const False) (const False) (const False) (const True)
isStdout :: ListLikePlus a c => Chunk a -> Bool
isStdout = foldChunk (const False) (const True) (const False) (const False) (const False)
isStderr :: ListLikePlus a c => Chunk a -> Bool
isStderr = foldChunk (const False) (const False) (const True) (const False) (const False)
isOutput :: ListLikePlus a c => Chunk a -> Bool
isOutput = foldChunk (const False) (const True) (const True) (const False) (const False)
isException :: ListLikePlus a c => Chunk a -> Bool
isException = foldChunk (const False) (const False) (const False) (const True) (const False)

toStdout :: ListLikePlus a c => Chunk a -> Chunk a
toStdout = foldChunk ProcessHandle Stdout Stdout Exception Result
toStderr :: ListLikePlus a c => Chunk a -> Chunk a
toStderr = foldChunk ProcessHandle Stderr Stderr Exception Result

mergeToStdout :: ListLikePlus a c => [Chunk a] -> [Chunk a]
mergeToStdout = map toStdout
mergeToStderr :: ListLikePlus a c => [Chunk a] -> [Chunk a]
mergeToStderr = map toStderr

discardStdout :: ListLikePlus a c => [Chunk a] -> [Chunk a]
discardStdout = filter (not . isStdout)
discardStderr :: ListLikePlus a c => [Chunk a] -> [Chunk a]
discardStderr = filter (not . isStderr)
discardOutput :: ListLikePlus a c => [Chunk a] -> [Chunk a]
discardOutput = filter (\ x -> not (isStdout x || isStderr x))
discardExceptions :: ListLikePlus a c => [Chunk a] -> [Chunk a]
discardExceptions = filter (not . isException)
discardResult :: ListLikePlus a c => [Chunk a] -> [Chunk a]
discardResult = filter (not . isResult)

keepStdout :: ListLikePlus a c => [Chunk a] -> [a]
keepStdout = mapMaybe $ foldChunk (const Nothing) Just (const Nothing) (const Nothing) (const Nothing)
keepStderr :: ListLikePlus a c => [Chunk a] -> [a]
keepStderr = mapMaybe $ foldChunk (const Nothing) (const Nothing) Just (const Nothing) (const Nothing)
keepOutput :: ListLikePlus a c => [Chunk a] -> [a]
keepOutput = mapMaybe $ foldChunk (const Nothing) Just Just (const Nothing) (const Nothing)
keepExceptions :: ListLikePlus a c => [Chunk a] -> [IOError]
keepExceptions = mapMaybe $ foldChunk (const Nothing) (const Nothing) (const Nothing) Just (const Nothing)
keepResult :: ListLikePlus a c => [Chunk a] -> [ExitCode]
keepResult = mapMaybe $ foldChunk (const Nothing) (const Nothing) (const Nothing) (const Nothing) Just

mapMaybeResult :: ListLikePlus a c => (ExitCode -> Maybe (Chunk a)) -> [Chunk a] -> [Chunk a]
mapMaybeResult f = mapMaybe (foldChunk (Just . ProcessHandle) (Just . Stdout) (Just . Stderr) (Just . Exception) f)
mapMaybeStdout :: ListLikePlus a c => (a -> Maybe (Chunk a)) -> [Chunk a] -> [Chunk a]
mapMaybeStdout f = mapMaybe (foldChunk (Just . ProcessHandle) f (Just . Stderr) (Just . Exception) (Just . Result))
mapMaybeStderr :: ListLikePlus a c => (a -> Maybe (Chunk a)) -> [Chunk a] -> [Chunk a]
mapMaybeStderr f = mapMaybe (foldChunk (Just . ProcessHandle) (Just . Stdout) f (Just . Exception) (Just . Result))
mapMaybeException :: ListLikePlus a c => (IOError -> Maybe (Chunk a)) -> [Chunk a] -> [Chunk a]
mapMaybeException f = mapMaybe (foldChunk (Just . ProcessHandle) (Just . Stdout) (Just . Stderr) f (Just . Result))

collectOutputs :: forall a c. ListLikePlus a c => [Chunk a] -> ([ExitCode], a, a, [IOError])
collectOutputs xs =
    foldChunks (\ r -> foldChunk (pidfn r) (outfn r) (errfn r) (exnfn r) (codefn r)) result0 xs
    where
      result0 :: ([ExitCode], a, a, [IOError])
      result0 = ([], empty, empty, [])
      pidfn :: ([ExitCode], a, a, [IOError]) -> ProcessHandle -> ([ExitCode], a, a, [IOError])
      pidfn (codes, outs, errs, exns) _ = (codes, outs, errs, exns)
      codefn :: ([ExitCode], a, a, [IOError]) -> ExitCode -> ([ExitCode], a, a, [IOError])
      codefn (codes, outs, errs, exns) code = (code : codes, outs, errs, exns)
      outfn :: ([ExitCode], a, a, [IOError]) -> a -> ([ExitCode], a, a, [IOError])
      outfn (codes, outs, errs, exns) out = (codes, append out outs, errs, exns)
      errfn :: ([ExitCode], a, a, [IOError]) -> a -> ([ExitCode], a, a, [IOError])
      errfn (codes, outs, errs, exns) err = (codes, outs, append err errs, exns)
      exnfn :: ([ExitCode], a, a, [IOError]) -> IOError -> ([ExitCode], a, a, [IOError])
      exnfn (codes, outs, errs, exns) exn = (codes, outs, errs, exn : exns)

ePutStr :: MonadIO m => String -> m ()
ePutStr s = liftIO $ IO.hPutStr stderr s

ePutStrLn :: MonadIO m => String -> m ()
ePutStrLn s = liftIO $ IO.hPutStrLn stderr s

eMessage :: MonadIO m => String -> a -> m a
eMessage s x = ePutStr s >> return x

eMessageLn :: MonadIO m => String -> a -> m a
eMessageLn s x = ePutStrLn s >> return x

foldException :: ListLikePlus a c => (IOError -> IO (Chunk a)) -> [Chunk a] -> IO [Chunk a]
foldException exnfn = mapM (foldChunk (return . ProcessHandle) (return . Stdout) (return . Stderr) exnfn (return . Result))

foldChars :: ListLikePlus a c => (a -> IO (Chunk a)) -> (a -> IO (Chunk a)) -> [Chunk a] -> IO [Chunk a]
foldChars outfn errfn = mapM (foldChunk (return . ProcessHandle) outfn errfn (return . Exception)(return . Result))

foldStdout :: ListLikePlus a c => (a -> IO (Chunk a)) -> [Chunk a] -> IO [Chunk a]
foldStdout outfn = foldChars outfn (return . Stderr)

foldStderr :: ListLikePlus a c => (a -> IO (Chunk a)) -> [Chunk a] -> IO [Chunk a]
foldStderr errfn = foldChars (return . Stdout) errfn

foldResult :: ListLikePlus a c => (ExitCode -> IO (Chunk a)) -> [Chunk a] -> IO [Chunk a]
foldResult codefn = mapM (foldChunk (return . ProcessHandle) (return . Stdout) (return . Stderr) (return . Exception) codefn)

foldFailure :: ListLikePlus a c => (Int -> IO (Chunk a)) -> [Chunk a] -> IO [Chunk a]
foldFailure failfn = foldResult codefn
    where codefn (ExitFailure n) = failfn n
          codefn x = return (Result x)

foldResult' :: ListLikePlus a c => ([Chunk a] -> ExitCode -> IO (Chunk a)) -> [Chunk a] -> IO [Chunk a]
foldResult' codefn outputs = mapM (foldChunk (return . ProcessHandle) (return . Stdout) (return . Stderr) (return . Exception) (codefn outputs)) outputs

foldFailure' :: ListLikePlus a c => (Int -> IO (Chunk a)) -> [Chunk a] -> IO [Chunk a]
foldFailure' failfn outputs = foldResult' codefn outputs
    where codefn _outputs (ExitFailure n) = failfn n
          codefn _ x = return (Result x)

foldSuccess :: ListLikePlus a c => IO (Chunk a) -> [Chunk a] -> IO [Chunk a]
foldSuccess successfn = foldResult codefn
    where codefn ExitSuccess = successfn
          codefn x = return (Result x)

doException :: ListLikePlus a c => [Chunk a] -> IO [Chunk a]
doException = foldException throw

doOutput :: ListLikePlus a c => [Chunk a] -> IO [Chunk a]
doOutput = foldChars (\ cs -> hPutStr stdout cs >> return (Stdout cs)) (\ cs -> hPutStr stderr cs >> return (Stderr cs))

doStdout :: ListLikePlus a c => [Chunk a] -> IO [Chunk a]
doStdout = foldStdout (\ cs -> hPutStr stdout cs >> return (Stdout cs))

doStderr :: ListLikePlus a c => [Chunk a] -> IO [Chunk a]
doStderr = foldStderr (\ cs -> hPutStr stderr cs >> return (Stderr cs))

-- | I don't see much use for this.
doExit :: ListLikePlus a c => [Chunk a] -> IO [Chunk a]
doExit = foldResult (\ code -> exitWith code >> return (Result code))

doAll :: ListLikePlus a c => [Chunk a] -> IO [Chunk a]
doAll = mapM (foldChunk (\ pid -> return (ProcessHandle pid))
                         (\ cs -> hPutStr stdout cs >> return (Stdout cs))
                         (\ cs -> hPutStr stderr cs >> return (Stderr cs))
                         throw
                         (\ code -> return (Result code)))

dots :: (ListLikePlus a c, c ~ Char) => Int -> {- (Int -> IO ()) -> -} [Chunk a] -> IO [Chunk a]
dots charsPerDot chunks = putDots charsPerDot '.' chunks
{-
dots :: forall a c. ListLikePlus a c => Int -> (Int -> IO ()) -> [Chunk a] -> IO [Chunk a]
dots charsPerDot nDots outputs =
    nDots 1 >> dots' 0 outputs >>= eMessage "\n"
    where
      dots' _ [] = nDots 1 >> return []
      dots' rem (x : xs) =
          do let (count', rem') = divMod (rem + foldChunk (const 0) length length (const 0) (const 0) x) charsPerDot
             when (count' > 0) (nDots count')
             xs' <- dots' rem' xs
             return (x : xs')
-}

-- | Output the stream with a prefix added at the beginning of each
-- line of stdout and stderr.
prefixed :: forall a c. (Enum c, ListLikePlus a c) => a -> a -> [Chunk a] -> IO [Chunk a]
prefixed opre epre output =
    mapM (\ (old, new) -> doOutput [new] >> return old) (prefixes opre epre output)

-- | Return the original stream of outputs zipped with one that has
-- had prefixes for stdout and stderr inserted.  For the prefixed
-- stream only, apply @map snd@.
prefixes :: forall a c. (Enum c, ListLikePlus a c) => a -> a -> [Chunk a] -> [(Chunk a, Chunk a)]
prefixes opre epre output =
    loop True output
    where
      loop :: (Enum c, ListLike a c) => Bool -> [Chunk a] -> [(Chunk a, Chunk a)]
      loop _ [] = []
      loop bol (x@(Stdout s) : xs) = let (s', bol') = step bol opre s in (x, Stdout s') : loop bol' xs
      loop bol (x@(Stderr s) : xs) = let (s', bol') = step bol epre s in (x, Stderr s') : loop bol' xs
      loop bol (x : xs) = (x, Stdout empty) : loop bol xs

      step :: (Enum c, ListLike a c) => Bool -> a -> a -> (a, Bool)
      step bol pre s =
          let (a, b) = Data.ListLike.span (\ c -> fromEnum c /= fromEnum '\n') s in
          if null a
          then if null b
               then (empty, bol)
               else let x = (if bol then pre else empty)
                        (s', bol') = step True pre (tail b) in
                    (concat [x, singleton (toEnum . fromEnum $ '\n'), s'], bol')
          -- There is some text before a possible newline
          else let x = (if bol then append pre a else a)
                   (s', bol') = step False pre b in
               (append x s', bol')