{-|
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)


-- | A script is a funclet that has its source code and a list of other funclets it depends on.
data Script = MkScript { Script -> String
source   :: String
                       , Script -> [AnyFunclet]
funclets :: [AnyFunclet]
                       }

-- | Marker for the thsh quasi-quote to recognize a 'Script'.
sh :: Script -> Script
sh :: Script -> Script
sh = Script -> Script
forall a. a -> a
id

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
    handles <- IO (MVar (Handle, Handle, Handle))
forall a. IO (MVar a)
newEmptyMVar
    _ <- forkIO $ withSystemTempDirectory "thsh-script.d" $ \ 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 hInW, Just hOutR, Just hErrR, 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
                                       }
          putMVar handles (hInW, hOutR, hErrR)

          -- create control thread
          unless (length funclets == 0) $ void . forkIO $ withFile ctlFifo ReadMode $ \ 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
          ec <- pollProcessExitCode mainProc

          cb ec

    (hInW, hOutR, hErrR) <- takeMVar handles
    pure (hInW, hOutR, hErrR)

-- | The piping code snippet that should substitute the funclet occurrences during quasi quoting.
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)

{- 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
  ecVar <- IO (MVar ExitCode)
forall a. IO (MVar a)
newEmptyMVar
  (fh0, fh1, fh2) <- runFunclet f (putMVar ecVar)
  -- piping data between the main process and the funclet process
  void . forkIO $ bracket
    (do
        -- create main process handlers (mh)
        mh0 <- openBinaryFile (procDir </> "0.fifo") ReadMode
        mh1 <- openBinaryFile (procDir </> "1.fifo") ReadWriteMode
        mh2 <- openBinaryFile (procDir </> "2.fifo") ReadWriteMode
        mapM_ (`hSetBuffering` NoBuffering) [mh0, mh1, mh2]
        pure (mh0, mh1, 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
        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) [(),(),()]
        mapM_ forkIO [ binaryCat mh0 fh0 >> hClose fh0 >> putMVar (cs !! 0) ()
                     , binaryCat fh1 mh1 >> putMVar (cs !! 1) ()
                     , binaryCat fh2 mh2 >> putMVar (cs !! 2) ()
                     ]
        -- wait for all pipes to finish
        mapM_ takeMVar cs
    )
  -- wait for the errno and save it for the main process
  void . forkIO $ takeMVar ecVar >>= \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
|]