    {-# LANGUAGE ScopedTypeVariables #-}
    module What4.Utils.HandleReader where

    import           Control.Monad (unless)
    import           Data.IORef
    import           Data.Text (Text)
    import qualified Data.Text as Text
    import qualified Data.Text.Lazy as LazyText
    import qualified Data.Text.IO as Text
    import           Control.Exception(bracket,catch,IOException)
    import           Control.Concurrent(ThreadId,forkIO,killThread)
    import           Control.Concurrent.Chan(Chan,newChan,readChan,writeChan)
    import           System.IO(Handle,hClose)
    import           System.IO.Streams( OutputStream, InputStream )
    import qualified System.IO.Streams as Streams

    teeInputStream :: InputStream a -> OutputStream a -> IO (InputStream a)
    teeInputStream :: forall a. InputStream a -> OutputStream a -> IO (InputStream a)
    teeInputStream InputStream a
    i OutputStream a
    o = IO (Maybe a) -> IO (InputStream a)
    forall a. IO (Maybe a) -> IO (InputStream a)
    Streams.makeInputStream IO (Maybe a)
    go
      where
      go :: IO (Maybe a)
    go = do x <- InputStream a -> IO (Maybe a)
    forall a. InputStream a -> IO (Maybe a)
    Streams.read InputStream a
    i
              Streams.write x o
              return x

    teeOutputStream :: OutputStream a -> OutputStream a -> IO (OutputStream a)
    teeOutputStream :: forall a. OutputStream a -> OutputStream a -> IO (OutputStream a)
    teeOutputStream OutputStream a
    o OutputStream a
    aux = (Maybe a -> IO ()) -> IO (OutputStream a)
    forall a. (Maybe a -> IO ()) -> IO (OutputStream a)
    Streams.makeOutputStream Maybe a -> IO ()
    go
      where
      go :: Maybe a -> IO ()
    go Maybe a
    x =
        do Maybe a -> OutputStream a -> IO ()
    forall a. Maybe a -> OutputStream a -> IO ()
    Streams.write Maybe a
    x OutputStream a
    aux
           Maybe a -> OutputStream a -> IO ()
    forall a. Maybe a -> OutputStream a -> IO ()
    Streams.write Maybe a
    x OutputStream a
    o

    lineBufferedOutputStream :: Text -> OutputStream Text -> IO (OutputStream Text)
    lineBufferedOutputStream :: Text -> OutputStream Text -> IO (OutputStream Text)
    lineBufferedOutputStream Text
    prefix OutputStream Text
    out =
        do ref <- Text -> IO (IORef Text)
    forall a. a -> IO (IORef a)
    newIORef Text
    forall a. Monoid a => a
    mempty
           Streams.makeOutputStream (con ref)
     where
     newl :: Text
    newl = String -> Text
    Text.pack String
    "\n"

     con :: IORef Text -> Maybe Text -> IO ()
    con IORef Text
    ref Maybe Text
    mx =
       do start <- IORef Text -> IO Text
    forall a. IORef a -> IO a
    readIORef IORef Text
    ref
          case mx of
            Maybe Text
    Nothing ->
              do Bool -> IO () -> IO ()
    forall (f :: Type -> Type). Applicative f => Bool -> f () -> f ()
    unless (Text -> Bool
    Text.null Text
    start) (Maybe Text -> OutputStream Text -> IO ()
    forall a. Maybe a -> OutputStream a -> IO ()
    Streams.write (Text -> Maybe Text
    forall a. a -> Maybe a
    Just (Text
    prefix Text -> Text -> Text
    forall a. Semigroup a => a -> a -> a
    <> Text
    start)) OutputStream Text
    out)
                 Maybe Text -> OutputStream Text -> IO ()
    forall a. Maybe a -> OutputStream a -> IO ()
    Streams.write Maybe Text
    forall a. Maybe a
    Nothing OutputStream Text
    out
            Just Text
    x -> IORef Text -> Text -> IO ()
    go IORef Text
    ref (Text
    start Text -> Text -> Text
    forall a. Semigroup a => a -> a -> a
    <> Text
    x)

     go :: IORef Text -> Text -> IO ()
    go IORef Text
    ref Text
    x =
       let (Text
    ln, Text
    x') = (Char -> Bool) -> Text -> (Text, Text)
    Text.break (Char -> Char -> Bool
    forall a. Eq a => a -> a -> Bool
    == Char
    '\n') Text
    x in
       if Text -> Bool
    Text.null Text
    x' then
         -- Flush
         do Maybe Text -> OutputStream Text -> IO ()
    forall a. Maybe a -> OutputStream a -> IO ()
    Streams.write (Text -> Maybe Text
    forall a. a -> Maybe a
    Just Text
    forall a. Monoid a => a
    mempty) OutputStream Text
    out
            IORef Text -> Text -> IO ()
    forall a. IORef a -> a -> IO ()
    writeIORef IORef Text
    ref Text
    x
       else
         do Maybe Text -> OutputStream Text -> IO ()
    forall a. Maybe a -> OutputStream a -> IO ()
    Streams.write (Text -> Maybe Text
    forall a. a -> Maybe a
    Just (Text
    prefix Text -> Text -> Text
    forall a. Semigroup a => a -> a -> a
    <> Text
    ln Text -> Text -> Text
    forall a. Semigroup a => a -> a -> a
    <> Text
    newl)) OutputStream Text
    out
            IORef Text -> Text -> IO ()
    go IORef Text
    ref (Int -> Text -> Text
    Text.drop Int
    1 Text
    x')

    demuxProcessHandles ::
      Handle {- ^ stdin for process -} ->
      Handle {- ^ stdout for process -} ->
      Handle {- ^ stderr for process -} ->
      Maybe (Text, Handle) {- optional handle to echo ouput; text argument is a line-comment prefix  -} ->
      IO ( OutputStream Text, InputStream Text, HandleReader )
    demuxProcessHandles :: Handle
    -> Handle
    -> Handle
    -> Maybe (Text, Handle)
    -> IO (OutputStream Text, InputStream Text, HandleReader)
    demuxProcessHandles Handle
    in_h Handle
    out_h Handle
    err_h Maybe (Text, Handle)
    Nothing =
      do in_str  <- OutputStream ByteString -> IO (OutputStream Text)
    Streams.encodeUtf8 (OutputStream ByteString -> IO (OutputStream Text))
    -> IO (OutputStream ByteString) -> IO (OutputStream Text)
    forall (m :: Type -> Type) a b. Monad m => (a -> m b) -> m a -> m b
    =<< Handle -> IO (OutputStream ByteString)
    Streams.handleToOutputStream Handle
    in_h
         out_str <- Streams.decodeUtf8 =<< Streams.handleToInputStream out_h
         err_reader <- startHandleReader err_h Nothing
         return (in_str, out_str, err_reader)
    demuxProcessHandles Handle
    in_h Handle
    out_h Handle
    err_h (Just (Text
    comment_prefix, Handle
    aux_h)) =
      do aux_str <- OutputStream Text -> IO (OutputStream Text)
    forall a. OutputStream a -> IO (OutputStream a)
    Streams.lockingOutputStream (OutputStream Text -> IO (OutputStream Text))
    -> IO (OutputStream Text) -> IO (OutputStream Text)
    forall (m :: Type -> Type) a b. Monad m => (a -> m b) -> m a -> m b
    =<< OutputStream ByteString -> IO (OutputStream Text)
    Streams.encodeUtf8 (OutputStream ByteString -> IO (OutputStream Text))
    -> IO (OutputStream ByteString) -> IO (OutputStream Text)
    forall (m :: Type -> Type) a b. Monad m => (a -> m b) -> m a -> m b
    =<< Handle -> IO (OutputStream ByteString)
    Streams.handleToOutputStream Handle
    aux_h
         in_str  <- Streams.encodeUtf8 =<< Streams.handleToOutputStream in_h
         out_str <- Streams.decodeUtf8 =<< Streams.handleToInputStream out_h

         in_aux <- lineBufferedOutputStream mempty aux_str
         in_str' <- teeOutputStream in_str in_aux

         out_aux <- lineBufferedOutputStream comment_prefix aux_str
         out_str' <- teeInputStream out_str out_aux

         err_reader <- startHandleReader err_h . Just
                        =<< lineBufferedOutputStream comment_prefix aux_str

         return (in_str', out_str', err_reader)

    {- | Wrapper to help with reading from another process's
         standard out and stderr.

    We want to be able to read from another process's stderr and stdout without
    causing the process to stall because 'stdout' or 'stderr' becomes full.  This
    data type will read from either of the handles, and buffer as much data
    as needed in the queue.  It then provides a line-based method for reading
    that data as strict bytestrings. -}
    data HandleReader = HandleReader { HandleReader -> Chan (Maybe Text)
    hrChan :: !(Chan (Maybe Text))
                                     , HandleReader -> Handle
    hrHandle :: !Handle
                                     , HandleReader -> ThreadId
    hrThreadId :: !ThreadId
                                     }

    streamLines :: Chan (Maybe Text) -> Handle -> Maybe (OutputStream Text) -> IO ()
    streamLines :: Chan (Maybe Text) -> Handle -> Maybe (OutputStream Text) -> IO ()
    streamLines Chan (Maybe Text)
    c Handle
    h Maybe (OutputStream Text)
    Nothing = IO ()
    forall {b}. IO b
    go
     where
     go :: IO b
    go = do ln <- Handle -> IO Text
    Text.hGetLine Handle
    h
             writeChan c (Just ln)
             go
    streamLines Chan (Maybe Text)
    c Handle
    h (Just OutputStream Text
    auxstr) = IO ()
    forall {b}. IO b
    go
     where
     go :: IO b
    go = do ln <- Handle -> IO Text
    Text.hGetLine Handle
    h
             Streams.write (Just ln) auxstr
             writeChan c (Just ln)
             go

    -- | Create a new handle reader for reading the given handle.
    startHandleReader :: Handle -> Maybe (OutputStream Text) -> IO HandleReader
    startHandleReader :: Handle -> Maybe (OutputStream Text) -> IO HandleReader
    startHandleReader Handle
    h Maybe (OutputStream Text)
    auxOutput = do
      c <- IO (Chan (Maybe Text))
    forall a. IO (Chan a)
    newChan
      let handle_err (IOException
    _e :: IOException) = Chan (Maybe Text) -> Maybe Text -> IO ()
    forall a. Chan a -> a -> IO ()
    writeChan Chan (Maybe Text)
    c Maybe Text
    forall a. Maybe a
    Nothing
      tid <- forkIO $ streamLines c h auxOutput `catch` handle_err

      return $! HandleReader { hrChan     = c
                             , hrHandle   = h
                             , hrThreadId = tid
                             }

    -- | Stop the handle reader; cannot be used afterwards.
    stopHandleReader :: HandleReader -> IO ()
    stopHandleReader :: HandleReader -> IO ()
    stopHandleReader HandleReader
    hr = do
      ThreadId -> IO ()
    killThread (HandleReader -> ThreadId
    hrThreadId HandleReader
    hr)
      Handle -> IO ()
    hClose (HandleReader -> Handle
    hrHandle HandleReader
    hr)

    -- | Run an execution with a handle reader and stop it wheen down
    withHandleReader :: Handle -> Maybe (OutputStream Text) -> (HandleReader -> IO a) -> IO a
    withHandleReader :: forall a.
    Handle
    -> Maybe (OutputStream Text) -> (HandleReader -> IO a) -> IO a
    withHandleReader Handle
    h Maybe (OutputStream Text)
    auxOut = IO HandleReader
    -> (HandleReader -> IO ()) -> (HandleReader -> IO a) -> IO a
    forall a b c. IO a -> (a -> IO b) -> (a -> IO c) -> IO c
    bracket (Handle -> Maybe (OutputStream Text) -> IO HandleReader
    startHandleReader Handle
    h Maybe (OutputStream Text)
    auxOut) HandleReader -> IO ()
    stopHandleReader

    readNextLine :: HandleReader -> IO (Maybe Text)
    readNextLine :: HandleReader -> IO (Maybe Text)
    readNextLine HandleReader
    hr = do
      mr <- Chan (Maybe Text) -> IO (Maybe Text)
    forall a. Chan a -> IO a
    readChan (HandleReader -> Chan (Maybe Text)
    hrChan HandleReader
    hr)
      case mr of
        -- Write back 'Nothing' because thread should have terminated.
        Maybe Text
    Nothing -> Chan (Maybe Text) -> Maybe Text -> IO ()
    forall a. Chan a -> a -> IO ()
    writeChan (HandleReader -> Chan (Maybe Text)
    hrChan HandleReader
    hr) Maybe Text
    forall a. Maybe a
    Nothing
        Just{} -> () -> IO ()
    forall a. a -> IO a
    forall (m :: Type -> Type) a. Monad m => a -> m a
    return()
      return mr

    readAllLines :: HandleReader -> IO LazyText.Text
    readAllLines :: HandleReader -> IO Text
    readAllLines HandleReader
    hr = Text -> IO Text
    go Text
    LazyText.empty
      where go :: LazyText.Text -> IO LazyText.Text
            go :: Text -> IO Text
    go Text
    prev = do
              mr <- HandleReader -> IO (Maybe Text)
    readNextLine HandleReader
    hr
              case mr of
                Maybe Text
    Nothing -> Text -> IO Text
    forall a. a -> IO a
    forall (m :: Type -> Type) a. Monad m => a -> m a
    return Text
    prev
                Just Text
    e -> Text -> IO Text
    go (Text -> IO Text) -> Text -> IO Text
    forall a b. (a -> b) -> a -> b
    $! Text
    prev Text -> Text -> Text
    `LazyText.append` (Text -> Text
    LazyText.fromStrict Text
    e)
                                     Text -> Char -> Text
    `LazyText.snoc` Char
    '\n'