{-|
Module      : THSH.Fn
Description : Fn funclets are Haskell functions that talk with other funclets including the main shell script.
Copyright   : (c) Miao ZhiCheng, 2024
License     : MIT
Maintainer  : zhicheng.miao@gmail.com
Stability   : experimental
Portability : POSIX
-}
module THSH.Fn
  ( FnFunction (..)
  , ContentFn (..), stringContentFn, stringContentIOFn, textContentFn, textContentIOFn
  , LineReadFn (..), lineReadFn
  , Fn, fn ) where

import           Control.Concurrent      (forkIO)
import           Control.Concurrent.MVar (newEmptyMVar, putMVar, takeMVar)
import           Control.Exception       (bracket)
import           System.Exit             (ExitCode (..))
import           System.IO               (BufferMode (NoBuffering), Handle, hClose, hGetContents, hGetLine, hIsEOF,
                                          hPutStr, hSetBuffering)
-- process
import           System.Process          (createPipe)
-- text
import qualified Data.Text               as T
import qualified Data.Text.IO
--
import           THSH.Funclet            (Funclet (..))


-- | A 'FnFunction' is a function that, given a set of handles to communicate with it, it returns an exit code.
class FnFunction f where
  runFn :: f -> (Handle, Handle, Handle) -> IO ExitCode

-- | The new type wrapper of any "FnFunction" instance.
newtype Fn f = MkFn f

-- | Marker for the thsh quasi-quote to recognize a 'FnFunction' code block.
fn :: FnFunction f => f -> Fn f
fn :: forall f. FnFunction f => f -> Fn f
fn = f -> Fn f
forall f. f -> Fn f
MkFn

-- | A 'FnFunction' that converts the entire input content to another as 'String'.
data ContentFn m s = MkContentFn (s -> m s) (Handle -> m s) (Handle -> s -> m ())

instance FnFunction (ContentFn IO s) where
  runFn :: ContentFn IO s -> (Handle, Handle, Handle) -> IO ExitCode
runFn (MkContentFn s -> IO s
f Handle -> IO s
r Handle -> s -> IO ()
w) (Handle
hIn, Handle
hOut, Handle
_) = do
    content <- Handle -> IO s
r Handle
hIn
    w hOut =<< f content
    pure ExitSuccess

-- | 'ContentFn' for the 'String' type.
stringContentFn :: (String -> String) -> ContentFn IO String
stringContentFn :: (String -> String) -> ContentFn IO String
stringContentFn String -> String
f = (String -> IO String)
-> (Handle -> IO String)
-> (Handle -> String -> IO ())
-> ContentFn IO String
forall (m :: * -> *) s.
(s -> m s)
-> (Handle -> m s) -> (Handle -> s -> m ()) -> ContentFn m s
MkContentFn (String -> IO String
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (String -> IO String) -> (String -> String) -> String -> IO String
forall b c a. (b -> c) -> (a -> b) -> a -> c
. String -> String
f) Handle -> IO String
hGetContents Handle -> String -> IO ()
hPutStr

-- | IO variant of 'stringContentFn'.
stringContentIOFn :: (String -> IO String) -> ContentFn IO String
stringContentIOFn :: (String -> IO String) -> ContentFn IO String
stringContentIOFn String -> IO String
f = (String -> IO String)
-> (Handle -> IO String)
-> (Handle -> String -> IO ())
-> ContentFn IO String
forall (m :: * -> *) s.
(s -> m s)
-> (Handle -> m s) -> (Handle -> s -> m ()) -> ContentFn m s
MkContentFn String -> IO String
f Handle -> IO String
hGetContents Handle -> String -> IO ()
hPutStr

-- | 'ContentFn' for the 'Data.Text' type from the text package.
textContentFn :: (T.Text -> T.Text) -> ContentFn IO T.Text
textContentFn :: (Text -> Text) -> ContentFn IO Text
textContentFn Text -> Text
f = (Text -> IO Text)
-> (Handle -> IO Text)
-> (Handle -> Text -> IO ())
-> ContentFn IO Text
forall (m :: * -> *) s.
(s -> m s)
-> (Handle -> m s) -> (Handle -> s -> m ()) -> ContentFn m s
MkContentFn (Text -> IO Text
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Text -> IO Text) -> (Text -> Text) -> Text -> IO Text
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Text -> Text
f) Handle -> IO Text
Data.Text.IO.hGetContents Handle -> Text -> IO ()
Data.Text.IO.hPutStr

-- | IO variant of 'textContentFn'.
textContentIOFn :: (T.Text -> IO T.Text) -> ContentFn IO T.Text
textContentIOFn :: (Text -> IO Text) -> ContentFn IO Text
textContentIOFn Text -> IO Text
f = (Text -> IO Text)
-> (Handle -> IO Text)
-> (Handle -> Text -> IO ())
-> ContentFn IO Text
forall (m :: * -> *) s.
(s -> m s)
-> (Handle -> m s) -> (Handle -> s -> m ()) -> ContentFn m s
MkContentFn Text -> IO Text
f Handle -> IO Text
Data.Text.IO.hGetContents Handle -> Text -> IO ()
Data.Text.IO.hPutStr

-- | A 'FnFunction' that reads line by line via 'Read' instances of @a@ and accumulates context @b@.
data LineReadFn m a b = Read a
                      => MkLineReadFn
                         (a -> b -> m (b, Maybe String)) -- ^ read an element; accumulate context; and maybe an output
                         (b -> m (Maybe String))         -- ^ final output
                         b                               -- ^ initial context

-- | Idiomatic wrapper for the `MkLineReadFn`
lineReadFn :: forall a b.
              Read a
           => (a -> b -> (b, Maybe String))
           -> (b -> Maybe String)
           -> b
           -> LineReadFn IO a b
lineReadFn :: forall a b.
Read a =>
(a -> b -> (b, Maybe String))
-> (b -> Maybe String) -> b -> LineReadFn IO a b
lineReadFn a -> b -> (b, Maybe String)
f b -> Maybe String
fin b
b0 = (a -> b -> IO (b, Maybe String))
-> (b -> IO (Maybe String)) -> b -> LineReadFn IO a b
forall (m :: * -> *) a b.
Read a =>
(a -> b -> m (b, Maybe String))
-> (b -> m (Maybe String)) -> b -> LineReadFn m a b
MkLineReadFn (((b, Maybe String) -> IO (b, Maybe String)
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure ((b, Maybe String) -> IO (b, Maybe String))
-> (b -> (b, Maybe String)) -> b -> IO (b, Maybe String)
forall b c a. (b -> c) -> (a -> b) -> a -> c
.) ((b -> (b, Maybe String)) -> b -> IO (b, Maybe String))
-> (a -> b -> (b, Maybe String)) -> a -> b -> IO (b, Maybe String)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. a -> b -> (b, Maybe String)
f) (Maybe String -> IO (Maybe String)
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Maybe String -> IO (Maybe String))
-> (b -> Maybe String) -> b -> IO (Maybe String)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. b -> Maybe String
fin) b
b0

instance FnFunction (LineReadFn IO a b) where
  runFn :: LineReadFn IO a b -> (Handle, Handle, Handle) -> IO ExitCode
runFn (MkLineReadFn a -> b -> IO (b, Maybe String)
f b -> IO (Maybe String)
fin b
b0) (Handle
hIn, Handle
hOut, Handle
_) = do
    let go :: b -> [IO (Maybe a)] -> IO ()
go b
b (IO (Maybe a)
a:[IO (Maybe a)]
as) = IO (Maybe a)
a IO (Maybe a) -> (Maybe a -> 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
>>= \case
          Just a
a'  -> a -> b -> IO (b, Maybe String)
f a
a' b
b IO (b, Maybe String) -> ((b, Maybe 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
>>= \ (b
b', Maybe String
r) ->
            case Maybe String
r of
              Just String
r' -> Handle -> String -> IO ()
hPutStr Handle
hOut String
r'
              Maybe String
Nothing -> () -> IO ()
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure ()
            IO () -> IO () -> IO ()
forall a b. IO a -> IO b -> IO b
forall (m :: * -> *) a b. Monad m => m a -> m b -> m b
>> b -> [IO (Maybe a)] -> IO ()
go b
b' [IO (Maybe a)]
as
          Maybe a
Nothing -> b -> IO (Maybe String)
fin b
b IO (Maybe String) -> (Maybe 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
>>= \case
            Just String
a' -> Handle -> String -> IO ()
hPutStr Handle
hOut String
a'
            Maybe String
Nothing -> () -> IO ()
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure () -- input lines finished
        go b
_ [IO (Maybe a)]
_ = String -> IO ()
forall a. HasCallStack => String -> a
error String
"impossible"
    -- repeatedly reading lines for @go@ to process, which should end with an infinite list of Nothings.
    b -> [IO (Maybe a)] -> IO ()
go b
b0 ([IO (Maybe a)] -> IO ()) -> [IO (Maybe a)] -> IO ()
forall a b. (a -> b) -> a -> b
$ IO (Maybe a) -> [IO (Maybe a)]
forall a. a -> [a]
repeat (Handle -> IO Bool
hIsEOF Handle
hIn IO Bool -> (Bool -> IO (Maybe a)) -> IO (Maybe a)
forall a b. IO a -> (a -> IO b) -> IO b
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= \ case
                       Bool
False -> Handle -> IO String
hGetLine Handle
hIn IO String -> (String -> IO (Maybe a)) -> IO (Maybe a)
forall a b. IO a -> (a -> IO b) -> IO b
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= Maybe a -> IO (Maybe a)
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Maybe a -> IO (Maybe a))
-> (String -> Maybe a) -> String -> IO (Maybe a)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. a -> Maybe a
forall a. a -> Maybe a
Just (a -> Maybe a) -> (String -> a) -> String -> Maybe a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (String -> a
forall a. Read a => String -> a
read :: String -> a)
                       Bool
True  -> Maybe a -> IO (Maybe a)
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Maybe a
forall a. Maybe a
Nothing)
    ExitCode -> IO ExitCode
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure ExitCode
ExitSuccess

instance FnFunction f => Funclet (Fn f) where
  runFunclet :: Fn f -> (ExitCode -> IO ()) -> IO (Handle, Handle, Handle)
runFunclet (MkFn f
f) ExitCode -> IO ()
cb = do
    handles <- IO (MVar (Handle, Handle, Handle))
forall a. IO (MVar a)
newEmptyMVar
    _ <- forkIO $ bracket
      (do
          (hInR, hInW)   <- createPipe
          (hOutR, hOutW) <- createPipe
          (hErrR, hErrW) <- createPipe
          mapM_ (`hSetBuffering` NoBuffering) [hInR, hInW, hOutR, hOutW, hErrR, hErrW]
          putMVar handles (hInW, hOutR, hErrR)
          pure (hInR, hOutW, hErrW)
      )
      (\(Handle
hInR, Handle
hOutW, Handle
hErrW) -> (Handle -> IO ()) -> [Handle] -> IO ()
forall (t :: * -> *) (m :: * -> *) a b.
(Foldable t, Monad m) =>
(a -> m b) -> t a -> m ()
mapM_ Handle -> IO ()
hClose [Handle
hInR, Handle
hOutW, Handle
hErrW])
      (\(Handle
hInR, Handle
hOutW, Handle
hErrW) -> ExitCode -> IO ()
cb (ExitCode -> IO ()) -> IO ExitCode -> IO ()
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< f -> (Handle, Handle, Handle) -> IO ExitCode
forall f.
FnFunction f =>
f -> (Handle, Handle, Handle) -> IO ExitCode
runFn f
f (Handle
hInR, Handle
hOutW, Handle
hErrW))
    takeMVar handles