path: root/src/Process/Shell.hs

{-# LANGUAGE BlockArguments #-}
{-# LANGUAGE DefaultSignatures #-}
{-# LANGUAGE IncoherentInstances #-}
{-# LANGUAGE QuasiQuotes #-}
{-# LANGUAGE TemplateHaskell #-}
{-# LANGUAGE TypeFamilies #-}
{-# LANGUAGE UndecidableInstances #-}

module Process.Shell
  ( sh,
    Quotable (..),
    Outputable (..),
    Inputable (..),
    ExitCodeException (..),
    DecodeException (..),
  )
where

import Conduit
import Control.Exception (Exception, throw)
import Control.Monad
import Data.Aeson
import Data.ByteString.Char8 qualified as B
import Data.ByteString.Lazy.Char8 qualified as LB
import Data.ByteString.Lazy.UTF8 qualified as LB
import Data.ByteString.UTF8 qualified as B
import Data.Conduit.Process.Typed
import Data.Function
import Data.Maybe
import Data.String
import Data.Text qualified as T
import Data.Text.Encoding qualified as T
import Data.Text.Lazy qualified as LT
import Data.Text.Lazy.Encoding qualified as LT
import Language.Haskell.TH hiding (Type)
import Language.Haskell.TH.Quote
import Text.Megaparsec
import Text.Megaparsec.Char
import Prelude hiding (exp)

data Script = Script (LB.ByteString -> LB.ByteString) String

class Processable m a where
  sh_ :: m Script -> m a

instance (MonadIO m) => Processable m () where
  sh_ = ((\(Script _ s) -> liftIO (runProcess_ (fromString s))) =<<)

instance (MonadIO m, Outputable stdoutAndStderr) => Processable m stdoutAndStderr where
  sh_ = ((\(Script strip s) -> fmap (fromLBS . strip) (liftIO (readProcessInterleaved_ (fromString s)))) =<<)

instance (MonadIO m, Outputable stdout, Outputable stderr) => Processable m (stdout, stderr) where
  sh_ = ((\(Script stripNL s) -> fmap (\(out, err) -> (fromLBS (stripNL out), fromLBS (stripNL err))) (liftIO (readProcess_ (fromString s)))) =<<)

instance (MonadIO m, Outputable stdout) => Processable m (stdout, ()) where
  sh_ = ((\(Script stripNL s) -> fmap (\out -> (fromLBS (stripNL out), ())) (liftIO (readProcessStdout_ (fromString s)))) =<<)

instance (MonadIO m, Outputable stderr) => Processable m ((), stderr) where
  sh_ = ((\(Script stripNL s) -> fmap (\err -> ((), fromLBS (stripNL err))) (liftIO (readProcessStderr_ (fromString s)))) =<<)

instance (MonadIO m) => Processable m ExitCode where
  sh_ = ((\(Script _ s) -> liftIO (runProcess (fromString s))) =<<)

instance (MonadIO m, Outputable stdoutAndStderr) => Processable m (ExitCode, stdoutAndStderr) where
  sh_ = ((\(Script stripNL s) -> fmap (\(exitCode, outErr) -> (exitCode, fromLBS (stripNL outErr))) (liftIO (readProcessInterleaved (fromString s)))) =<<)

instance (MonadIO m, Outputable stdout) => Processable m (ExitCode, stdout, ()) where
  sh_ = ((\(Script stripNL s) -> fmap (\(exitCode, out) -> (exitCode, fromLBS (stripNL out), ())) (liftIO (readProcessStdout (fromString s)))) =<<)

instance (MonadIO m, Outputable stderr) => Processable m (ExitCode, (), stderr) where
  sh_ = ((\(Script stripNL s) -> fmap (\(exitCode, err) -> (exitCode, (), fromLBS (stripNL err))) (liftIO (readProcessStderr (fromString s)))) =<<)

instance (MonadIO m, Outputable stdout, Outputable stderr) => Processable m (ExitCode, stdout, stderr) where
  sh_ = ((\(Script stripNL s) -> fmap (\(exitCode, out, err) -> (exitCode, fromLBS (stripNL out), fromLBS (stripNL err))) (liftIO (readProcess (fromString s)))) =<<)

instance
  (MonadIO m, MonadResource m, Inputable stdin, Outputable stdoutAndStderr, Monoid stdoutAndStderr) =>
  Processable (ConduitT stdin Void m) stdoutAndStderr
  where
  sh_ =
    ( ( \(Script strip s) -> do
          let stripC = awaitForever (\i -> maybe (yield (strip i)) (\_ -> yield i) =<< peekC)
          bracketP
            ( startProcess
                ( fromString s
                    & setStdin createSinkClose
                    & setStdout createSource
                    & setStderr createSource
                )
            )
            stopProcess
            ( \p -> do
                awaitForever (yield . LB.toStrict . toLBS) .| transPipe liftIO (getStdin p)
                outErr <-
                  ( mapOutput fromLBS $
                      mapOutput LB.fromStrict (transPipe liftIO (getStdout p >> getStderr p))
                        .| stripC
                    )
                    .| foldC
                checkExitCode p
                pure outErr
            )
      )
        =<<
    )

instance
  (MonadIO m, MonadResource m, Inputable stdin, Outputable stdout, Outputable stderr, Monoid stdout, Monoid stderr) =>
  Processable (ConduitT stdin Void m) (stdout, stderr)
  where
  sh_ =
    ( ( \(Script strip s) -> do
          let stripC = awaitForever (\i -> maybe (yield (strip i)) (\_ -> yield i) =<< peekC)
          bracketP
            ( startProcess
                ( fromString s
                    & setStdin createSinkClose
                    & setStdout createSource
                    & setStderr createSource
                )
            )
            stopProcess
            ( \p -> do
                awaitForever (yield . LB.toStrict . toLBS) .| transPipe liftIO (getStdin p)
                outErr <-
                  (,)
                    <$> ( ( mapOutput fromLBS $
                              mapOutput LB.fromStrict (transPipe liftIO (getStdout p))
                                .| stripC
                          )
                            .| foldC
                        )
                    <*> ( ( mapOutput fromLBS $
                              mapOutput LB.fromStrict (transPipe liftIO (getStderr p))
                                .| stripC
                          )
                            .| foldC
                        )
                checkExitCode p
                pure outErr
            )
      )
        =<<
    )

instance
  (Monad m, MonadIO m, MonadResource m, Inputable stdin) =>
  Processable (ConduitT stdin Void m) ()
  where
  sh_ =
    ( ( \(Script _ s) -> do
          bracketP
            ( startProcess
                ( fromString s
                    & setStdin createSinkClose
                    & setStdout inherit
                    & setStderr inherit
                )
            )
            stopProcess
            ( \p -> do
                awaitForever (yield . LB.toStrict . toLBS) .| transPipe liftIO (getStdin p)
                checkExitCode p
            )
      )
        =<<
    )

instance
  (MonadIO m, MonadResource m, Inputable stdin, Outputable stdout, Monoid stdout) =>
  Processable (ConduitT stdin Void m) (stdout, ())
  where
  sh_ =
    ( ( \(Script strip s) -> do
          let stripC = awaitForever (\i -> maybe (yield (strip i)) (\_ -> yield i) =<< peekC)
          bracketP
            ( startProcess
                ( fromString s
                    & setStdin createSinkClose
                    & setStdout createSource
                    & setStderr inherit
                )
            )
            stopProcess
            ( \p -> do
                awaitForever (yield . LB.toStrict . toLBS) .| transPipe liftIO (getStdin p)
                out <-
                  ( mapOutput fromLBS $
                      mapOutput LB.fromStrict (transPipe liftIO (getStdout p))
                        .| stripC
                    )
                    .| foldC
                checkExitCode p
                pure (out, ())
            )
      )
        =<<
    )

instance
  (MonadIO m, MonadResource m, Inputable stdin, Outputable stderr, Monoid stderr) =>
  Processable (ConduitT stdin Void m) ((), stderr)
  where
  sh_ =
    ( ( \(Script strip s) -> do
          let stripC = awaitForever (\i -> maybe (yield (strip i)) (\_ -> yield i) =<< peekC)
          bracketP
            ( startProcess
                ( fromString s
                    & setStdin createSinkClose
                    & setStdout inherit
                    & setStderr createSource
                )
            )
            stopProcess
            ( \p -> do
                awaitForever (yield . LB.toStrict . toLBS) .| transPipe liftIO (getStdin p)
                err <-
                  ( mapOutput fromLBS $
                      mapOutput LB.fromStrict (transPipe liftIO (getStderr p))
                        .| stripC
                    )
                    .| foldC
                checkExitCode p
                pure ((), err)
            )
      )
        =<<
    )

instance
  (MonadIO m, MonadResource m, Inputable stdin, Outputable stdoutAndStderr) =>
  Processable (ConduitT stdin stdoutAndStderr m) ()
  where
  sh_ =
    ( ( \(Script strip s) -> do
          let stripC = awaitForever (\i -> maybe (yield (strip i)) (\_ -> yield i) =<< peekC)
          bracketP
            ( startProcess
                ( fromString s
                    & setStdin createSinkClose
                    & setStdout createSource
                    & setStderr createSource
                )
            )
            stopProcess
            ( \p -> do
                awaitForever (yield . LB.toStrict . toLBS) .| transPipe liftIO (getStdin p)
                mapOutput fromLBS $
                  mapOutput LB.fromStrict (transPipe liftIO (getStdout p >> getStderr p))
                    .| stripC
                checkExitCode p
            )
      )
        =<<
    )

instance
  (MonadIO m, MonadResource m, Inputable stdin, Outputable stdout, Outputable stderr) =>
  Processable (ConduitT stdin (Either stderr stdout) m) ()
  where
  sh_ =
    ( ( \(Script strip s) -> do
          let stripC = awaitForever (\i -> maybe (yield (strip i)) (\_ -> yield i) =<< peekC)
          bracketP
            ( startProcess
                ( fromString s
                    & setStdin createSinkClose
                    & setStdout createSource
                    & setStderr createSource
                )
            )
            stopProcess
            ( \p -> do
                awaitForever (yield . LB.toStrict . toLBS) .| transPipe liftIO (getStdin p)
                transPipe liftIO $ do
                  mapOutput Right (mapOutput fromLBS (mapOutput LB.fromStrict (getStdout p) .| stripC))
                  mapOutput Left (mapOutput fromLBS (mapOutput LB.fromStrict (getStderr p) .| stripC))
                checkExitCode p
            )
      )
        =<<
    )

class Outputable a where
  fromLBS :: LB.ByteString -> a

instance Outputable String where
  fromLBS = LB.toString

instance Outputable B.ByteString where
  fromLBS = LB.toStrict

instance Outputable LB.ByteString where
  fromLBS = id

instance Outputable T.Text where
  fromLBS = T.decodeUtf8 . fromLBS

instance Outputable LT.Text where
  fromLBS = LT.decodeUtf8

class Inputable a where
  toLBS :: a -> LB.ByteString

instance Inputable () where
  toLBS _ = LB.pack ""

instance Inputable String where
  toLBS = LB.fromString

instance Inputable B.ByteString where
  toLBS = LB.fromStrict

instance Inputable LB.ByteString where
  toLBS = id

instance Inputable T.Text where
  toLBS = LT.encodeUtf8 . LT.fromStrict

instance Inputable LT.Text where
  toLBS = LT.encodeUtf8

data DecodeException = DecodeException String
  deriving (Show)

instance Exception DecodeException

instance (FromJSON a) => Outputable a where
  fromLBS = either (throw . DecodeException) id . eitherDecode

class Quotable a where
  toString :: a -> String
  default toString :: (Show a) => a -> String
  toString = show

instance Quotable String where
  toString = id

instance Quotable Int

instance Quotable B.ByteString where
  toString = B.toString

instance Quotable LB.ByteString where
  toString = LB.toString

instance Quotable T.Text where
  toString = T.unpack

instance Quotable LT.Text where
  toString = LT.unpack

squote :: String -> String
squote s = "'" <> quote' s <> "'"
  where
    quote' [] = []
    quote' ('\'' : cs) = '\'' : '\\' : '\'' : '\'' : quote' cs
    quote' (c : cs) = c : quote' cs

dquote :: String -> String
dquote s = "\"" <> quote' s <> "\""
  where
    quote' [] = []
    quote' ('\\' : c : cs) = c : quote' cs
    quote' ('"' : cs) = '\\' : '"' : quote' cs
    quote' (c : cs) = c : quote' cs

data Expr a
  = Lit String
  | Var Bool a

sh :: QuasiQuoter
sh = QuasiQuoter quoteExp undefined undefined undefined
  where
    quoteExp :: String -> Q Exp
    quoteExp =
      either (fail . errorBundlePretty) makeExp'
        . parse (parser <* eof) ""

    parser :: ParsecT Void String Identity ([Expr (Q Exp)], Bool)
    parser =
      (,)
        <$> ( many . choice . map try $
                [ do
                    -- TODO Splice arbitrary Haskell expressions
                    Var False . varE . mkName
                      <$> ( string "#{{"
                              *> takeWhile1P Nothing (/= '}')
                              <* string "}}"
                          ),
                  do
                    Var True . varE . mkName
                      <$> ( string "'#{{"
                              *> takeWhile1P Nothing (/= '}')
                              <* string "}}'"
                          ),
                  do
                    Var False . appE [|pure|] . varE . mkName
                      <$> ( string "#{"
                              *> takeWhile1P Nothing (/= '}')
                              <* string "}"
                          ),
                  do
                    Var True . appE [|pure|] . varE . mkName
                      <$> ( string "'#{"
                              *> takeWhile1P Nothing (/= '}')
                              <* string "}'"
                          ),
                  do
                    Lit <$> takeWhile1P Nothing (\c -> all (c /=) ['\'', '#', '\\']),
                  do
                    Lit . (: []) <$> satisfy ((||) <$> (== '\'') <*> (== '#')),
                  do
                    Lit . (: []) <$> satisfy (== '\\') <* lookAhead (satisfy (const True))
                ]
            )
        <*> (isNothing <$> optional (string "\\"))

    makeExp' :: ([Expr (Q Exp)], Bool) -> Q Exp
    makeExp' (exprs, False) = [|sh_ (Script id <$> $(makeExp exprs))|]
    makeExp' (exprs, True) = [|sh_ (Script stripTrailingNLs <$> $(makeExp exprs))|]

    makeExp :: [Expr (Q Exp)] -> Q Exp
    makeExp exprs = do
      exprs' :: [(Expr Name)] <-
        sequence
          [ case expr of
              Lit s -> pure (Lit s)
              Var q _ -> Var q <$> newName "arg"
            | expr <- exprs
          ]
      doE $
        [ BindS <$> (varP nam) <*> [|toString <$> $exp|]
          | (Var _ exp, Var _ nam) <- zip exprs exprs'
        ]
          ++ [ NoBindS <$> [|pure $(foldr (\a b -> [|$a ++ $b|]) [|""|] (map toExp exprs'))|]
             ]

    toExp :: Expr Name -> Q Exp
    toExp (Lit s) = [|s|]
    toExp (Var True nam) = [|squote $(varE nam)|]
    toExp (Var False nam) = [|dquote $(varE nam)|]

stripTrailingNLs :: LB.ByteString -> LB.ByteString
stripTrailingNLs s = maybe s (stripTrailingNLs) (LB.stripSuffix (LB.pack "\n") s)