{-|
Module      : THSH.Script
Description : Script funclet create a system process for the shell script.
Copyright   : (c) Miao ZhiCheng, 2024
License     : MIT
Maintainer  : zhicheng.miao@gmail.com
Stability   : experimental
Portability : POSIX
-}
{-# LANGUAGE QuasiQuotes #-}
module THSH.Script
  ( Script (..)
  , genFuncletPipeCode
  , sh
  ) where

-- base module
import           Control.Concurrent         (forkIO, newEmptyMVar, putMVar, takeMVar)
import           Control.Exception          (bracket, catch)
import           Control.Monad              (unless, void)
import           System.Exit                (ExitCode (..))
import           System.IO                  (BufferMode (NoBuffering), IOMode (ReadMode, ReadWriteMode, WriteMode),
                                             hClose, hGetLine, hPutStr, hPutStrLn, hSetBuffering, openBinaryFile,
                                             stderr, withFile)
-- filepath module
import           System.FilePath            ((</>))
-- process module
import           System.Process             (CreateProcess (std_err, std_in, std_out), StdStream (CreatePipe),
                                             callCommand, createProcess, shell)
-- temporary module
import           System.IO.Temp             (withSystemTempDirectory)
-- PyF module
import           PyF                        (fmt, str)
--
import           THSH.Funclet               (AnyFunclet, Funclet (..))
import           THSH.Internal.ProcessUtils (binaryCat, pollProcessExitCode)


data Script = MkScript { Script -> String
source   :: String
                       , Script -> [AnyFunclet]
funclets :: [AnyFunclet]
                       }

instance Funclet Script where
  runFunclet :: Script -> (ExitCode -> IO ()) -> IO (Handle, Handle, Handle)
runFunclet (MkScript { String
source :: Script -> String
source :: String
source, [AnyFunclet]
funclets :: Script -> [AnyFunclet]
funclets :: [AnyFunclet]
funclets }) ExitCode -> IO ()
cb = do
    MVar (Handle, Handle, Handle)
handles <- IO (MVar (Handle, Handle, Handle))
forall a. IO (MVar a)
newEmptyMVar
    ThreadId
_ <- IO () -> IO ThreadId
forkIO (IO () -> IO ThreadId) -> IO () -> IO ThreadId
forall a b. (a -> b) -> a -> b
$ String -> (String -> IO ()) -> IO ()
forall (m :: * -> *) a.
(MonadIO m, MonadMask m) =>
String -> (String -> m a) -> m a
withSystemTempDirectory String
"thsh-script.d" ((String -> IO ()) -> IO ()) -> (String -> IO ()) -> IO ()
forall a b. (a -> b) -> a -> b
$ \ String
dir -> do
          let initCodePath :: String
initCodePath = String
dir String -> String -> String
</> String
"init.sh"
              srcPath :: String
srcPath      = String
dir String -> String -> String
</> String
"source.sh"
              ctlFifo :: String
ctlFifo      = String
dir String -> String -> String
</> String
"cr.fifo"

          -- write init code
          String -> IOMode -> (Handle -> IO ()) -> IO ()
forall r. String -> IOMode -> (Handle -> IO r) -> IO r
withFile String
initCodePath IOMode
WriteMode ((Handle -> IO ()) -> IO ()) -> (Handle -> IO ()) -> IO ()
forall a b. (a -> b) -> a -> b
$ \Handle
fh -> Handle -> String -> IO ()
hPutStr Handle
fh (Int -> String
gen_init_code ([AnyFunclet] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [AnyFunclet]
funclets))

          -- call init code
          String -> IO ()
callCommand (String
"sh " String -> String -> String
forall a. Semigroup a => a -> a -> a
<> String
initCodePath)

          -- write source file
          String -> IOMode -> (Handle -> IO ()) -> IO ()
forall r. String -> IOMode -> (Handle -> IO r) -> IO r
withFile String
srcPath IOMode
WriteMode ((Handle -> IO ()) -> IO ()) -> (Handle -> IO ()) -> IO ()
forall a b. (a -> b) -> a -> b
$ \Handle
fh -> do
            Handle -> String -> IO ()
hPutStr Handle
fh String
gen_stub_code
            Handle -> String -> IO ()
hPutStr Handle
fh String
source

          -- create the main process
          (Just Handle
hInW, Just Handle
hOutR, Just Handle
hErrR, ProcessHandle
mainProc) <- CreateProcess
-> IO (Maybe Handle, Maybe Handle, Maybe Handle, ProcessHandle)
createProcess (CreateProcess
 -> IO (Maybe Handle, Maybe Handle, Maybe Handle, ProcessHandle))
-> CreateProcess
-> IO (Maybe Handle, Maybe Handle, Maybe Handle, ProcessHandle)
forall a b. (a -> b) -> a -> b
$
            (String -> CreateProcess
shell (String
"sh " String -> String -> String
forall a. Semigroup a => a -> a -> a
<> String
srcPath)) { std_in  = CreatePipe
                                       , std_out = CreatePipe
                                       , std_err = CreatePipe
                                       }
          MVar (Handle, Handle, Handle) -> (Handle, Handle, Handle) -> IO ()
forall a. MVar a -> a -> IO ()
putMVar MVar (Handle, Handle, Handle)
handles (Handle
hInW, Handle
hOutR, Handle
hErrR)

          -- create control thread
          Bool -> IO () -> IO ()
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
unless ([AnyFunclet] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [AnyFunclet]
funclets Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
0) (IO () -> IO ()) -> IO () -> IO ()
forall a b. (a -> b) -> a -> b
$ IO ThreadId -> IO ()
forall (f :: * -> *) a. Functor f => f a -> f ()
void (IO ThreadId -> IO ()) -> (IO () -> IO ThreadId) -> IO () -> IO ()
forall b c a. (b -> c) -> (a -> b) -> a -> c
. IO () -> IO ThreadId
forkIO (IO () -> IO ()) -> IO () -> IO ()
forall a b. (a -> b) -> a -> b
$ String -> IOMode -> (Handle -> IO ()) -> IO ()
forall r. String -> IOMode -> (Handle -> IO r) -> IO r
withFile String
ctlFifo IOMode
ReadMode ((Handle -> IO ()) -> IO ()) -> (Handle -> IO ()) -> IO ()
forall a b. (a -> b) -> a -> b
$ \ Handle
ch -> let
            go :: IO ()
go = do
              IO String -> (IOError -> IO String) -> IO String
forall e a. Exception e => IO a -> (e -> IO a) -> IO a
catch (Handle -> IO String
hGetLine Handle
ch) (\ (IOError
e :: IOError) -> Handle -> String -> IO ()
hPutStrLn Handle
stderr (IOError -> String
forall a. Show a => a -> String
show IOError
e) IO () -> IO String -> IO String
forall a b. IO a -> IO b -> IO b
forall (m :: * -> *) a b. Monad m => m a -> m b -> m b
>> String -> IO String
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure []) IO String -> (String -> IO ()) -> IO ()
forall a b. IO a -> (a -> IO b) -> IO b
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= \String
cmd -> do
                case String
cmd of
                  []        -> () -> IO ()
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure () -- likely end of file
                  Char
's':Char
' ':String
i -> AnyFunclet -> String -> IO ()
forall f. Funclet f => f -> String -> IO ()
start_funclet_proc ([AnyFunclet]
funclets [AnyFunclet] -> Int -> AnyFunclet
forall a. HasCallStack => [a] -> Int -> a
!! (String -> Int
forall a. Read a => String -> a
read String
i :: Int)) (String
dir String -> String -> String
</> String
i) IO () -> IO () -> IO ()
forall a b. IO a -> IO b -> IO b
forall (m :: * -> *) a b. Monad m => m a -> m b -> m b
>> () -> IO ()
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure ()
                  String
_         -> Handle -> String -> IO ()
hPutStrLn Handle
stderr (String -> IO ()) -> String -> IO ()
forall a b. (a -> b) -> a -> b
$ String
"[thsh-script] unknown control command: " String -> String -> String
forall a. Semigroup a => a -> a -> a
<> String
cmd
                if String
cmd String -> String -> Bool
forall a. Eq a => a -> a -> Bool
/= String
"" then IO ()
go else () -> IO ()
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure ()
            in IO ()
go

          -- wait for the main sub process to finish
          ExitCode
ec <- ProcessHandle -> IO ExitCode
pollProcessExitCode ProcessHandle
mainProc

          ExitCode -> IO ()
cb ExitCode
ec

    (Handle
hInW, Handle
hOutR, Handle
hErrR) <- MVar (Handle, Handle, Handle) -> IO (Handle, Handle, Handle)
forall a. MVar a -> IO a
takeMVar MVar (Handle, Handle, Handle)
handles
    (Handle, Handle, Handle) -> IO (Handle, Handle, Handle)
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Handle
hInW, Handle
hOutR, Handle
hErrR)

genFuncletPipeCode :: Int -> String
genFuncletPipeCode :: Int -> String
genFuncletPipeCode Int
i = String
"__pipeFunclet " String -> String -> String
forall a. Semigroup a => a -> a -> a
<> (Int -> String
forall a. Show a => a -> String
show Int
i)

sh :: Script -> Script
sh :: Script -> Script
sh = Script -> Script
forall a. a -> a
id

{- INTERNAL FUNCTIONS -}

start_funclet_proc :: Funclet f => f -> FilePath -> IO ()
start_funclet_proc :: forall f. Funclet f => f -> String -> IO ()
start_funclet_proc f
f String
procDir = do
  -- use mVar to communicate errno
  MVar ExitCode
ecVar <- IO (MVar ExitCode)
forall a. IO (MVar a)
newEmptyMVar
  (Handle
fh0, Handle
fh1, Handle
fh2) <- f -> (ExitCode -> IO ()) -> IO (Handle, Handle, Handle)
forall f.
Funclet f =>
f -> (ExitCode -> IO ()) -> IO (Handle, Handle, Handle)
runFunclet f
f (MVar ExitCode -> ExitCode -> IO ()
forall a. MVar a -> a -> IO ()
putMVar MVar ExitCode
ecVar)
  -- piping data between the main process and the funclet process
  IO ThreadId -> IO ()
forall (f :: * -> *) a. Functor f => f a -> f ()
void (IO ThreadId -> IO ()) -> (IO () -> IO ThreadId) -> IO () -> IO ()
forall b c a. (b -> c) -> (a -> b) -> a -> c
. IO () -> IO ThreadId
forkIO (IO () -> IO ()) -> IO () -> IO ()
forall a b. (a -> b) -> a -> b
$ IO (Handle, Handle, Handle)
-> ((Handle, Handle, Handle) -> IO ())
-> ((Handle, Handle, Handle) -> IO ())
-> IO ()
forall a b c. IO a -> (a -> IO b) -> (a -> IO c) -> IO c
bracket
    (do
        -- create main process handlers (mh)
        Handle
mh0 <- String -> IOMode -> IO Handle
openBinaryFile (String
procDir String -> String -> String
</> String
"0.fifo") IOMode
ReadMode
        Handle
mh1 <- String -> IOMode -> IO Handle
openBinaryFile (String
procDir String -> String -> String
</> String
"1.fifo") IOMode
ReadWriteMode
        Handle
mh2 <- String -> IOMode -> IO Handle
openBinaryFile (String
procDir String -> String -> String
</> String
"2.fifo") IOMode
ReadWriteMode
        (Handle -> IO ()) -> [Handle] -> IO ()
forall (t :: * -> *) (m :: * -> *) a b.
(Foldable t, Monad m) =>
(a -> m b) -> t a -> m ()
mapM_ (Handle -> BufferMode -> IO ()
`hSetBuffering` BufferMode
NoBuffering) [Handle
mh0, Handle
mh1, Handle
mh2]
        (Handle, Handle, Handle) -> IO (Handle, Handle, Handle)
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Handle
mh0, Handle
mh1, Handle
mh2)
    )
    (\(Handle
mh0, Handle
mh1, Handle
mh2) -> (Handle -> IO ()) -> [Handle] -> IO ()
forall (t :: * -> *) (m :: * -> *) a b.
(Foldable t, Monad m) =>
(a -> m b) -> t a -> m ()
mapM_ Handle -> IO ()
hClose [Handle
mh0, Handle
mh1, Handle
mh2])
    (\(Handle
mh0, Handle
mh1, Handle
mh2) -> do
        -- fork pipes
        [MVar ()]
cs <- (() -> IO (MVar ())) -> [()] -> IO [MVar ()]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
forall (m :: * -> *) a b. Monad m => (a -> m b) -> [a] -> m [b]
mapM (IO (MVar ()) -> () -> IO (MVar ())
forall a b. a -> b -> a
const IO (MVar ())
forall a. IO (MVar a)
newEmptyMVar) [(),(),()]
        (IO () -> IO ThreadId) -> [IO ()] -> IO ()
forall (t :: * -> *) (m :: * -> *) a b.
(Foldable t, Monad m) =>
(a -> m b) -> t a -> m ()
mapM_ IO () -> IO ThreadId
forkIO [ Handle -> Handle -> IO ()
binaryCat Handle
mh0 Handle
fh0 IO () -> IO () -> IO ()
forall a b. IO a -> IO b -> IO b
forall (m :: * -> *) a b. Monad m => m a -> m b -> m b
>> Handle -> IO ()
hClose Handle
fh0 IO () -> IO () -> IO ()
forall a b. IO a -> IO b -> IO b
forall (m :: * -> *) a b. Monad m => m a -> m b -> m b
>> MVar () -> () -> IO ()
forall a. MVar a -> a -> IO ()
putMVar ([MVar ()]
cs [MVar ()] -> Int -> MVar ()
forall a. HasCallStack => [a] -> Int -> a
!! Int
0) ()
                     , Handle -> Handle -> IO ()
binaryCat Handle
fh1 Handle
mh1 IO () -> IO () -> IO ()
forall a b. IO a -> IO b -> IO b
forall (m :: * -> *) a b. Monad m => m a -> m b -> m b
>> MVar () -> () -> IO ()
forall a. MVar a -> a -> IO ()
putMVar ([MVar ()]
cs [MVar ()] -> Int -> MVar ()
forall a. HasCallStack => [a] -> Int -> a
!! Int
1) ()
                     , Handle -> Handle -> IO ()
binaryCat Handle
fh2 Handle
mh2 IO () -> IO () -> IO ()
forall a b. IO a -> IO b -> IO b
forall (m :: * -> *) a b. Monad m => m a -> m b -> m b
>> MVar () -> () -> IO ()
forall a. MVar a -> a -> IO ()
putMVar ([MVar ()]
cs [MVar ()] -> Int -> MVar ()
forall a. HasCallStack => [a] -> Int -> a
!! Int
2) ()
                     ]
        -- wait for all pipes to finish
        (MVar () -> IO ()) -> [MVar ()] -> IO ()
forall (t :: * -> *) (m :: * -> *) a b.
(Foldable t, Monad m) =>
(a -> m b) -> t a -> m ()
mapM_ MVar () -> IO ()
forall a. MVar a -> IO a
takeMVar [MVar ()]
cs
    )
  -- wait for the errno and save it for the main process
  IO ThreadId -> IO ()
forall (f :: * -> *) a. Functor f => f a -> f ()
void (IO ThreadId -> IO ()) -> (IO () -> IO ThreadId) -> IO () -> IO ()
forall b c a. (b -> c) -> (a -> b) -> a -> c
. IO () -> IO ThreadId
forkIO (IO () -> IO ()) -> IO () -> IO ()
forall a b. (a -> b) -> a -> b
$ MVar ExitCode -> IO ExitCode
forall a. MVar a -> IO a
takeMVar MVar ExitCode
ecVar IO ExitCode -> (ExitCode -> IO ()) -> IO ()
forall a b. IO a -> (a -> IO b) -> IO b
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= \ExitCode
ec -> do
    case ExitCode
ec of
      ExitCode
ExitSuccess       -> () -> IO ()
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure ()
      ExitFailure Int
errno -> String -> IOMode -> (Handle -> IO ()) -> IO ()
forall r. String -> IOMode -> (Handle -> IO r) -> IO r
withFile (String
procDir String -> String -> String
</> String
"errno") IOMode
WriteMode (Handle -> String -> IO ()
`hPutStr` (Int -> String
forall a. Show a => a -> String
show Int
errno))

gen_init_code :: Int -> String
gen_init_code :: Int -> String
gen_init_code Int
nFunclets = String
[str|\
#set -x
__initFunclets() {
  n="$1"
  d=$(dirname "$0")

  # nothing to initialize when there is no funclet
  [ "$n" == "-1" ] && return

  # create control fifos
  mkfifo "$d"/c{w,r}.fifo
  tail -f "$d"/cw.fifo > "$d"/cr.fifo &

  # create funclet fifos
  seq 0 "$n" | while read i; do
    mkdir -p "$d/$i"
    mkfifo "$d/$i"/{0,1,2}.fifo
  done
}
|] String -> String -> String
forall a. Semigroup a => a -> a -> a
<> [fmt|__initFunclets {show (nFunclets - 1)}
|]

gen_stub_code :: String
gen_stub_code :: String
gen_stub_code = String
[str|\
#+BEGIN_SRC THSH script stub code
__pipeFunclet() (
  # connect to the fifos of nth functlet
  i="$1"
  d="$(dirname "$0")/$i"

  cat <&0 > "$d"/0.fifo & pid0=$!
  cat >&1 < "$d"/1.fifo & pid1=$!
  cat >&2 < "$d"/2.fifo & pid2=$!
  trap 'kill $pid0 $pid1 $pid2' SIGINT

  echo "s $i" > "$(dirname "$0")"/cw.fifo

  wait $pid0 $pid1 $pid2
  : __pipeFunclet ended
)
#+END_SRC
|]