{-# LANGUAGE RecordWildCards #-}

module Main where

import           Control.Concurrent.Async
import           Control.Concurrent.STM
import           Control.Monad
import           Data.ByteString (ByteString)
import qualified Data.ByteString as BS
import           Data.List (elemIndex)
import           Options.Applicative as Opt
import           System.IO
import           System.Process

data Options = Options
    { verbose  :: Bool
    , commands :: [(String,String)]
    }
    deriving Show

splitOn :: Char -> String -> (String, String)
splitOn c xs =
    maybe ("unnamed", xs) (\(x, _:y) -> (x, y))
        $ flip splitAt xs <$> elemIndex c xs

options :: Opt.Parser Options
options = Options
    <$> switch (short 'v' <> long "verbose" <> help "Display verbose output")
    <*> some (option (splitOn '=' <$> str)
                     (short 'c' <>
                      long "command" <>
                      help "Display verbose output"))

data MsgType = MsgError | MsgOutput
    deriving (Eq, Show)

data Message = Message
    { messageType    :: MsgType
    , messageOrigin  :: String
    , messageContent :: ByteString
    }
    deriving Show

loggingActor :: Bool -> TChan Message -> IO ()
loggingActor verb msgs = forever $ do
    when verb $ putStrLn "Waiting for message..."
    msg@Message {..} <- atomically $ readTChan msgs
    when verb $ putStrLn $ "Read message: " ++ show msg
    putStr $ messageOrigin ++ ": "
          ++ case messageType of
              MsgError -> "ERROR: "
              MsgOutput -> ""
    BS.putStr messageContent
    putStrLn ""

readOutput :: Bool -> TChan Message -> MsgType -> String -> Handle -> IO ()
readOutput verb msgs msgType label h = go
  where
      go = do
          eof <- hIsEOF h
          unless eof $ do
              line <- BS.hGetLine h
              when verb $ putStrLn $ "read: " ++ label ++ ": " ++ show line
              atomically $ writeTChan msgs (Message msgType label line)
              go

runCommands :: Options -> IO ()
runCommands opts = do
    msgs <- newTChanIO
    withAsync (loggingActor verb msgs) $ \_ -> do
        hs <- forM (commands opts) $ \(label, cmd) -> do
            let p = CreateProcess
                    { cmdspec            = ShellCommand cmd
                    , cwd                = Nothing
                    , env                = Nothing
                    , std_in             = Inherit
                    , std_out            = CreatePipe
                    , std_err            = CreatePipe
                    , close_fds          = False
                    , create_group       = False
                    , delegate_ctlc      = False
                    , detach_console     = False
                    , create_new_console = False
                    , new_session        = False
                    , child_group        = Nothing
                    , child_user         = Nothing
                    }
            when verb $ putStrLn $ "Executing: " ++ cmd
            (_, Just hout, Just herr, h) <- createProcess p

            t1 <- async $ readOutput verb msgs MsgOutput label hout
            link t1
            t2 <- async $ readOutput verb msgs MsgError label herr
            link t2

            return h

        when verb $ putStrLn "Waiting for processes to finish..."
        forM_ hs waitForProcess

        when verb $ putStrLn "Waiting for logging queue to flush..."
        atomically $ do
            e <- isEmptyTChan msgs
            unless e retry

        when verb $ putStrLn "Exiting..."
  where
    verb = verbose opts

main :: IO ()
main = execParser opts >>= runCommands
  where
    opts = info (helper <*> options)
                (fullDesc
                 <> progDesc "Run multiple commands, interleaving output and errors"
                 <> header "runmany - run many commands at once")