{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE LambdaCase #-}
{-# LANGUAGE MultiWayIf #-}
{-# LANGUAGE OverloadedStrings #-}
{-# LANGUAGE RecordWildCards #-}
module Flameshot.Internal.Process (Callbacks(..),run,runAway) where

import Control.Concurrent (forkIO,threadDelay)
import Control.Concurrent.Async (race)
import Control.Exception
import Control.Monad.Extended (untilM_,unless)
import Data.ByteString (ByteString)
import qualified Data.ByteString as BS
import Data.Text (Text)
import qualified Data.Text.IO as T
import System.Exit
import System.IO (BufferMode(LineBuffering),hSetBuffering)
import System.IO (Handle,hClose,hIsEOF)
import System.IO.Error (catchIOError,isDoesNotExistError)
import System.Process
import System.Posix.Process (getProcessGroupIDOf)
import System.Posix.Signals (Signal,signalProcessGroup,killProcess)
import System.Posix.Types (ProcessGroupID)


data Callbacks = Callbacks
    { onOutLine :: Pid -> Text -> IO ()
    , onErrLine :: Pid -> Text -> IO ()
    , onError :: Pid -> SomeException -> IO ()
    , onExit :: Pid -> ExitCode -> IO ()
    , onStart :: Pid -> IO ()
    }

instance Semigroup Callbacks where
  a <> b =
      Callbacks
        { onOutLine = \pid s -> do
          onOutLine a pid s
          onOutLine b pid s

        , onErrLine = \pid s -> do
          onErrLine a pid s
          onErrLine b pid s

        , onError = \pid err -> do
          onError a pid err
          onError b pid err

        , onExit = \pid status -> do
          onExit a pid status
          onExit b pid status

        , onStart = \pid -> do
          onStart a pid
          onStart b pid
        }

instance Monoid Callbacks where
  mempty =
    Callbacks
      { onOutLine = \_pid _s -> return ()
      , onErrLine = \_pid _s -> return ()
      , onError = \_pid _err -> return ()
      , onExit = \_pid _status -> return ()
      , onStart = \_pid -> return ()
      }

run :: FilePath
     -> [String]
     -> Maybe FilePath
     -> Maybe [(String, String)]
     -> ByteString
     -> Maybe Int
     -> Callbacks
     -> IO ()
run path args cwd env input hTimeout Callbacks{..} =
    f `catch` onError (-1)
  where
    f = withCreateProcess p $ \(Just inh) (Just outh) (Just errh) ph -> do
      Just pid <- getPid ph
      pgid <- getProcessGroupIDOf pid

      onStart pid

      mapM_ forkIO [
          BS.hPutStr inh input `finally` hClose inh,
          hWithLines outh (onOutLine pid),
          hWithLines errh (onErrLine pid)
        ]

      case hTimeout of
        Just time ->
          race (threadDelay time) (waitForProcess ph) >>= \case
            Left () -> onError pid (SomeException (ErrorCall "timeout"))
            Right code -> onExit pid code
        Nothing ->
          waitForProcess ph >>= onExit pid

      killProcessGroup pgid

    p = (proc path args)
          { cwd = cwd
          , env = env
          , std_in = CreatePipe
          , std_out = CreatePipe
          , std_err = CreatePipe
          , close_fds = True
          , create_group = True
          , new_session = True
          }

runAway :: FilePath
     -> [String]
     -> Maybe FilePath
     -> Maybe [(String, String)]
     -> ByteString
     -> Maybe Int
     -> Callbacks
     -> IO ()
runAway path args cwd env input hTimeout Callbacks{..} =
    f `catch` onError (-1)
  where
    f = do
      (Just inh, Just outh, Just errh, ph) <- createProcess p

      Just pid <- getPid ph

      onStart pid

      mapM_ forkIO [
          BS.hPutStr inh input `finally` hClose inh,
          hWithLines outh (onOutLine pid),
          hWithLines errh (onErrLine pid)
        ]

      case hTimeout of
        Just time ->
          race (threadDelay time) (waitForProcess ph) >>= \case
            Left () -> onError pid (SomeException (ErrorCall "timeout"))
            Right code -> onExit pid code
        Nothing ->
          waitForProcess ph >>= onExit pid

    p = (proc path args)
          { cwd = cwd
          , env = env
          , std_in = CreatePipe
          , std_out = CreatePipe
          , std_err = CreatePipe
          , close_fds = True
          , create_group = True
          , new_session = True
          }

killProcessGroup :: ProcessGroupID -> IO ()
killProcessGroup = signalProcessGroup' killProcess

signalProcessGroup' :: Signal -> ProcessGroupID -> IO ()
signalProcessGroup' sig pgid =
    catchIOError
      (signalProcessGroup sig pgid)
      (\e -> unless (isDoesNotExistError e) $ ioError e)

hWithLines :: Handle -> (Text -> IO ()) -> IO ()
hWithLines h f = do
    hSetBuffering h LineBuffering
    untilM_ (hIsEOF h) (T.hGetLine h >>= f) `finally` hClose h