{-|
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           Data.Function              ((&))
import           Data.Maybe                 (fromJust)
import           System.Exit                (ExitCode (..))
import           System.IO                  (BufferMode (NoBuffering), Handle,
                                             IOMode (ReadMode, ReadWriteMode, WriteMode), hClose, hGetLine, hPutStr,
                                             hPutStrLn, hSetBuffering, openBinaryFile, stderr, withFile)
-- filepath module
import           System.FilePath            ((</>))
-- process module
import           System.Process             (CreateProcess (..), StdStream (CreatePipe, UseHandle), 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 contains shell source code and a list of other funclets it depends on.
data Script = MkScript { Script -> String
source   :: String
                       , Script -> [AnyFunclet]
funclets :: [AnyFunclet]
                       }

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

-- | The 'Script' instance of 'Funclet'.
instance Funclet Script where
  runFunclet :: Script -> (ExitCode -> IO ()) -> IO (Handle, Handle, Handle)
runFunclet Script
script ExitCode -> IO ()
cb = Script
-> (ExitCode -> IO ())
-> Maybe (Handle, Handle, Handle)
-> IO (Maybe (Handle, Handle, Handle))
run_script_funclet Script
script ExitCode -> IO ()
cb Maybe (Handle, Handle, Handle)
forall a. Maybe a
Nothing IO (Maybe (Handle, Handle, Handle))
-> (Maybe (Handle, Handle, Handle) -> IO (Handle, Handle, Handle))
-> IO (Handle, Handle, Handle)
forall a b. IO a -> (a -> IO b) -> IO b
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= (Handle, Handle, Handle) -> IO (Handle, Handle, Handle)
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure ((Handle, Handle, Handle) -> IO (Handle, Handle, Handle))
-> (Maybe (Handle, Handle, Handle) -> (Handle, Handle, Handle))
-> Maybe (Handle, Handle, Handle)
-> IO (Handle, Handle, Handle)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Maybe (Handle, Handle, Handle) -> (Handle, Handle, Handle)
forall a. HasCallStack => Maybe a -> a
fromJust

  runFuncletWithHandles :: Script -> (ExitCode -> IO ()) -> (Handle, Handle, Handle) -> IO ()
runFuncletWithHandles Script
script ExitCode -> IO ()
cb (Handle, Handle, Handle)
providedHandles = IO (Maybe (Handle, Handle, Handle)) -> IO ()
forall (f :: * -> *) a. Functor f => f a -> f ()
void (Script
-> (ExitCode -> IO ())
-> Maybe (Handle, Handle, Handle)
-> IO (Maybe (Handle, Handle, Handle))
run_script_funclet Script
script ExitCode -> IO ()
cb ((Handle, Handle, Handle) -> Maybe (Handle, Handle, Handle)
forall a. a -> Maybe a
Just (Handle, Handle, Handle)
providedHandles))

-- | 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 -}

run_script_funclet :: Script -> (ExitCode -> IO ()) -> Maybe (Handle, Handle, Handle)
                   -> IO (Maybe (Handle, Handle, Handle))
run_script_funclet :: Script
-> (ExitCode -> IO ())
-> Maybe (Handle, Handle, Handle)
-> IO (Maybe (Handle, Handle, Handle))
run_script_funclet (MkScript { String
source :: Script -> String
source :: String
source, [AnyFunclet]
funclets :: Script -> [AnyFunclet]
funclets :: [AnyFunclet]
funclets }) ExitCode -> IO ()
cb Maybe (Handle, Handle, Handle)
providedHandles = do
  mProcHandles <- IO (MVar (Maybe (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 shell script process
        mMainProc <- IO (MVar ProcessHandle)
forall a. IO (MVar a)
newEmptyMVar
        -- TODO: I can't make this work with withCreateProcess
        -- withCreateProcess
        --   (shell ("sh " <> srcPath) & \ procSpec -> case providedHandles of
        --       Just (hInR, hOutW, hErrW) -> procSpec { std_in  = UseHandle hInR
        --                                             , std_out = UseHandle hOutW
        --                                             , std_err = UseHandle hErrW
        --                                             }
        --       Nothing                   -> procSpec { std_in  = CreatePipe
        --                                             , std_out = CreatePipe
        --                                             , std_err = CreatePipe
        --                                             }
        --   )
        --   (\ cases
        --     (Just hInW) (Just hOutR) (Just hErrR) mainProc -> putMVar mProcHandles (Just (hInW, hOutR, hErrR))
        --                                                       >> putMVar mMainProc mainProc
        --     _ _ _                                 mainProc -> putMVar mProcHandles Nothing
        --                                                       >> putMVar mMainProc mainProc
        --   )
        createProcess
          (shell ("sh " <> srcPath) & \ CreateProcess
procSpec -> case Maybe (Handle, Handle, Handle)
providedHandles of
              Just (Handle
hInR, Handle
hOutW, Handle
hErrW) -> CreateProcess
procSpec { std_in  = UseHandle hInR
                                                    , std_out = UseHandle hOutW
                                                    , std_err = UseHandle hErrW
                                                    }
              Maybe (Handle, Handle, Handle)
Nothing                   -> CreateProcess
procSpec { std_in  = CreatePipe
                                                    , std_out = CreatePipe
                                                    , std_err = CreatePipe
                                                    }
          )
          >>= (\ cases
                (Just Handle
hInW, Just Handle
hOutR, Just Handle
hErrR, ProcessHandle
mainProc) -> MVar (Maybe (Handle, Handle, Handle))
-> Maybe (Handle, Handle, Handle) -> IO ()
forall a. MVar a -> a -> IO ()
putMVar MVar (Maybe (Handle, Handle, Handle))
mProcHandles ((Handle, Handle, Handle) -> Maybe (Handle, Handle, Handle)
forall a. a -> Maybe a
Just (Handle
hInW, Handle
hOutR, Handle
hErrR))
                                                                 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 ProcessHandle -> ProcessHandle -> IO ()
forall a. MVar a -> a -> IO ()
putMVar MVar ProcessHandle
mMainProc ProcessHandle
mainProc
                (Maybe Handle
_, Maybe Handle
_, Maybe Handle
_,                           ProcessHandle
mainProc) -> MVar (Maybe (Handle, Handle, Handle))
-> Maybe (Handle, Handle, Handle) -> IO ()
forall a. MVar a -> a -> IO ()
putMVar MVar (Maybe (Handle, Handle, Handle))
mProcHandles Maybe (Handle, Handle, Handle)
forall a. Maybe a
Nothing
                                                                 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 ProcessHandle -> ProcessHandle -> IO ()
forall a. MVar a -> a -> IO ()
putMVar MVar ProcessHandle
mMainProc ProcessHandle
mainProc
          )

        -- 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 (String
"THSH.Script control thread error: " String -> String -> String
forall a. Semigroup a => a -> a -> a
<> 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

        catch
          (takeMVar mMainProc >>= pollProcessExitCode >>= cb)
          (\ (IOError
e :: IOError) -> Handle -> String -> IO ()
hPutStrLn Handle
stderr (String
"THSH.Script funclet thread error: " String -> String -> String
forall a. Semigroup a => a -> a -> a
<> IOError -> String
forall a. Show a => a -> String
show IOError
e)
                               IO () -> IO () -> IO ()
forall a b. IO a -> IO b -> IO b
forall (m :: * -> *) a b. Monad m => m a -> m b -> m b
>> ExitCode -> IO ()
cb (Int -> ExitCode
ExitFailure Int
2))

  takeMVar mProcHandles


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
|]