-- ghc-mod: Making Haskell development *more* fun
-- Copyright (C) 2015  Daniel Gröber <dxld ÄT darkboxed DOT org>
--
-- This program is free software: you can redistribute it and/or modify
-- it under the terms of the GNU Affero General Public License as published by
-- the Free Software Foundation, either version 3 of the License, or
-- (at your option) any later version.
--
-- This program is distributed in the hope that it will be useful,
-- but WITHOUT ANY WARRANTY; without even the implied warranty of
-- MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
-- GNU Affero General Public License for more details.
--
-- You should have received a copy of the GNU Affero General Public License
-- along with this program.  If not, see <http://www.gnu.org/licenses/>.

-- Derived from process:System.Process
-- Copyright (c) The University of Glasgow 2004-2008

{-# LANGUAGE FlexibleInstances #-}
module Language.Haskell.GhcMod.Output (
    gmPutStr
  , gmErrStr
  , gmPutStrLn
  , gmErrStrLn

  , gmPutStrIO
  , gmErrStrIO

  , gmReadProcess
  , gmReadProcess'

  , stdoutGateway
  , flushStdoutGateway
  ) where

import Data.List
import Data.ByteString.Lazy (ByteString)
import qualified Data.ByteString.Lazy as BS
import qualified Data.Label as L
import qualified Data.Label.Base as LB
import System.IO
import System.Exit
import System.Process
import Control.Monad
import Control.Monad.State.Strict
import Control.DeepSeq
import Control.Exception
import Control.Concurrent (forkIO, killThread)
import Control.Concurrent.MVar
import Control.Concurrent.Chan
import Pipes
import Pipes.Lift
import Prelude

import Language.Haskell.GhcMod.Types hiding (LineSeparator, MonadIO(..))
import Language.Haskell.GhcMod.Monad.Types hiding (MonadIO(..))
import Language.Haskell.GhcMod.Gap ()

class ProcessOutput a where
  hGetContents' :: Handle -> IO a

instance ProcessOutput String where
  hGetContents' = hGetContents

instance ProcessOutput ByteString where
  hGetContents' = BS.hGetContents

outputFns :: (GmOut m, MonadIO m')
          => m (String -> m' (), String -> m' ())
outputFns =
  outputFns' `liftM` gmoAsk

outputFns' ::
  MonadIO m => GhcModOut -> (String -> m (), String -> m ())
outputFns' (GhcModOut oopts c)  = let
  OutputOpts {..} = oopts
 in
  case ooptLinePrefix of
    Nothing -> stdioOutputFns
    Just _  -> chanOutputFns c

stdioOutputFns :: MonadIO m => (String -> m (), String -> m ())
stdioOutputFns =
    ( liftIO . putStr
    , liftIO . hPutStr stderr
    )

chanOutputFns :: MonadIO m
              => Chan (Either (MVar ()) (GmStream, String))
              -> (String -> m (), String -> m ())
chanOutputFns c = (write GmOutStream, write GmErrStream)
 where
   write stream s = liftIO $ writeChan c $ Right $ (stream,s)

gmPutStr, gmPutStrLn, gmErrStr, gmErrStrLn
    :: (MonadIO m, GmOut m) => String -> m ()

gmPutStr str = do
  putOut <- gmPutStrIO
  putOut str

gmErrStr str = do
  putErr <- gmErrStrIO
  putErr str

gmPutStrLn = gmPutStr . (++"\n")
gmErrStrLn = gmErrStr . (++"\n")

gmPutStrIO, gmErrStrIO :: (GmOut m, MonadIO mi) => m (String -> mi ())

gmPutStrIO = fst `liftM` outputFns
gmErrStrIO = snd `liftM` outputFns


gmReadProcess :: GmOut m => m (FilePath -> [String] -> String -> IO String)
gmReadProcess = do
  GhcModOut {..} <- gmoAsk
  case ooptLinePrefix gmoOptions of
    Just _ ->
        readProcessStderrChan
    Nothing ->
        return $ readProcess

gmReadProcess' :: GmOut m => m (FilePath -> [String] -> String -> IO ByteString)
gmReadProcess' = readProcessStderrChan

flushStdoutGateway :: Chan (Either (MVar ()) (GmStream, String)) -> IO ()
flushStdoutGateway c = do
  mv <- newEmptyMVar
  writeChan c $ Left mv
  takeMVar mv

type Line = String

stdoutGateway :: (String, String) -> Chan (Either (MVar ()) (GmStream, String)) -> IO ()
stdoutGateway (outPf, errPf) chan = do
  runEffect $ commandProc >-> evalStateP ("","") seperateStreams
 where
   commandProc :: Producer (Either (MVar ()) (GmStream, String)) IO ()
   commandProc = do
     cmd <- liftIO $ readChan chan
     case cmd of
       Left mv -> do
           yield $ Left mv
       Right input -> do
           yield $ Right input
           commandProc

   seperateStreams :: Consumer (Either (MVar ()) (GmStream, String)) (StateT (String, String) IO) ()
   seperateStreams  = do
       ecmd <- await
       case ecmd of
         Left mv -> do
           -- flush buffers
           (\s -> lift $ zoom (streamLens s) $ sGetLine Nothing)
               `mapM_` [GmOutStream, GmErrStream]

           liftIO $ putMVar mv ()
         Right (stream, str) -> do
           ls <- lift $ zoom (streamLens stream) $ sGetLine (Just str)
           case ls of
               [] -> return ()
               _  -> liftIO $ putStr $ unlines $ map (streamPf stream++) ls

           liftIO $ hFlush stdout
           seperateStreams

   sGetLine :: (Maybe String) -> StateT String IO [Line]
   sGetLine mstr' = do
     buf <- get
     let mstr = (buf++) `liftM` mstr'
     case mstr of
       Nothing -> put "" >> return [buf]
       Just "" -> return []
       Just s | last s == '\n' -> put "" >> return (lines s)
              | otherwise -> do
                  let (p:ls') = reverse $ lines s
                  put p
                  return $ reverse $ ls'

   streamLens GmOutStream = LB.fst
   streamLens GmErrStream = LB.snd

   streamPf GmOutStream = outPf
   streamPf GmErrStream = errPf

zoom :: Monad m => (f L.:-> o) -> StateT o m a -> StateT f m a
zoom l (StateT a) =
    StateT $ \f -> do
      (a', s') <- a $ L.get l f
      return (a', L.set l s' f)

readProcessStderrChan ::
    (GmOut m, ProcessOutput a, NFData a) => m (FilePath -> [String] -> String -> IO a)
readProcessStderrChan = do
  (_, e :: String -> IO ()) <- outputFns
  return $ readProcessStderrChan' e

readProcessStderrChan' :: (ProcessOutput a, NFData a) =>
    (String -> IO ()) -> FilePath -> [String] -> String -> IO a
readProcessStderrChan' putErr exe args input = do
     let cp = (proc exe args) {
                std_out = CreatePipe
              , std_err = CreatePipe
              , std_in  = CreatePipe
              }
     (Just i, Just o, Just e, h) <- createProcess cp

     _ <- forkIO $ reader e

     output  <- hGetContents' o
     withForkWait (evaluate $ rnf output) $ \waitOut -> do

       -- now write any input
       unless (null input) $
         ignoreSEx $ hPutStr i input
       -- hClose performs implicit hFlush, and thus may trigger a SIGPIPE
       ignoreSEx $ hClose i

       -- wait on the output
       waitOut
       hClose o

     res <- waitForProcess h
     case res of
       ExitFailure rv ->
           throw $ GMEProcess "readProcessStderrChan" exe args $ Left rv
       ExitSuccess ->
           return output
    where
      ignoreSEx = handle (\(SomeException _) -> return ())
      reader h = ignoreSEx $ do
        putErr . (++"\n") =<< hGetLine h
        reader h

withForkWait :: IO () -> (IO () ->  IO a) -> IO a
withForkWait async body = do
  waitVar <- newEmptyMVar :: IO (MVar (Either SomeException ()))
  mask $ \restore -> do
    tid <- forkIO $ try (restore async) >>= putMVar waitVar
    let wait = takeMVar waitVar >>= either throwIO return
    restore (body wait) `onException` killThread tid