-- | This module contains the logic to interact with the python wrapper
module Haxible.Connection (Connections (..), TaskCall (..), withConnections, cleanVar) where

import Control.Exception (bracket)
import Data.Aeson (eitherDecodeStrict, encode)
import Data.Aeson.KeyMap qualified
import Data.ByteString (hGetLine, toStrict)
import Data.ByteString.Char8 (hPutStrLn)
import Data.Hashable (Hashable)
import Data.Pool qualified
import Data.Text qualified as Text
import Data.Text.Encoding (decodeUtf8)
import Data.Typeable (Typeable)
import GHC.IO.Handle (hFlush)
import Haxible.Prelude
import Say
import System.Console.ANSI (Color (..), ColorIntensity (Dull), ConsoleLayer (Foreground), hSupportsANSIColor, setSGRCode)
import System.Console.ANSI.Codes (SGR (..))
import System.IO (Handle, hClose, stdout)
import System.Process (Pid, getPid)
import System.Process.Typed

data TaskCall = TaskCall
  { -- | The playbook attributes, such as `hosts` or `become`.
    TaskCall -> Vars
playAttrs :: Vars,
    -- | The module name for debug purpose, it is more convenient to access than reading it from the moduleObject.
    TaskCall -> Text
module_ :: Text,
    -- | The task object, e.g `{"file": {"path": "/etc/zuul"}}`.
    TaskCall -> Value
moduleObject :: Value,
    -- | Extra task attributes, such as "when"
    TaskCall -> Vars
taskAttrs :: Vars,
    -- | Extra task vars, e.g. role defaults
    TaskCall -> Vars
taskVars :: Vars
  }
  deriving (TaskCall -> TaskCall -> Bool
(TaskCall -> TaskCall -> Bool)
-> (TaskCall -> TaskCall -> Bool) -> Eq TaskCall
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
/= :: TaskCall -> TaskCall -> Bool
$c/= :: TaskCall -> TaskCall -> Bool
== :: TaskCall -> TaskCall -> Bool
$c== :: TaskCall -> TaskCall -> Bool
Eq, Int -> TaskCall -> ShowS
[TaskCall] -> ShowS
TaskCall -> String
(Int -> TaskCall -> ShowS)
-> (TaskCall -> String) -> ([TaskCall] -> ShowS) -> Show TaskCall
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
showList :: [TaskCall] -> ShowS
$cshowList :: [TaskCall] -> ShowS
show :: TaskCall -> String
$cshow :: TaskCall -> String
showsPrec :: Int -> TaskCall -> ShowS
$cshowsPrec :: Int -> TaskCall -> ShowS
Show, Typeable, (forall x. TaskCall -> Rep TaskCall x)
-> (forall x. Rep TaskCall x -> TaskCall) -> Generic TaskCall
forall x. Rep TaskCall x -> TaskCall
forall x. TaskCall -> Rep TaskCall x
forall a.
(forall x. a -> Rep a x) -> (forall x. Rep a x -> a) -> Generic a
$cto :: forall x. Rep TaskCall x -> TaskCall
$cfrom :: forall x. TaskCall -> Rep TaskCall x
Generic, Eq TaskCall
Eq TaskCall
-> (Int -> TaskCall -> Int)
-> (TaskCall -> Int)
-> Hashable TaskCall
Int -> TaskCall -> Int
TaskCall -> Int
forall a. Eq a -> (Int -> a -> Int) -> (a -> Int) -> Hashable a
hash :: TaskCall -> Int
$chash :: TaskCall -> Int
hashWithSalt :: Int -> TaskCall -> Int
$chashWithSalt :: Int -> TaskCall -> Int
Hashable)

-- | A connection run converts a TaskCall into a (result code, result value)
newtype Connections = Connections {Connections -> TaskCall -> IO (Int, Value)
run :: TaskCall -> IO (Int, Value)}

-- | Add horizontal line separator
-- >>> addSep 10 "TASK"
-- "TASK ****"
addSep :: Int -> Text -> Text
addSep :: Int -> Text -> Text
addSep Int
width Text
x = Text
x Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<> Text
" " Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<> Text
sep
  where
    sep :: Text
sep = Int -> Text -> Text
Text.replicate (Int
width Int -> Int -> Int
forall a. Num a => a -> a -> a
- Text -> Int
Text.length Text
x Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
2) Text
"*"

-- | Format process id.
-- >>> formatPid 42
-- "<42>"
formatPid :: Pid -> Text
formatPid :: Pid -> Text
formatPid Pid
pid = Text
"<" Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<> String -> Text
forall source target. From source target => source -> target
from (Pid -> String
forall a. Show a => a -> String
show Pid
pid) Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<> Text
">"

formatTask :: Pid -> TaskCall -> Text
formatTask :: Pid -> TaskCall -> Text
formatTask Pid
pid TaskCall
tc = Text
"TASK [" Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<> Text
name Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<> Text
"] " Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<> Pid -> Text
formatPid Pid
pid
  where
    name :: Text
name = Text -> Maybe Text -> Text
forall a. a -> Maybe a -> a
fromMaybe TaskCall
tc.module_ (Getting (First Text) Value Text -> Value -> Maybe Text
forall s (m :: * -> *) a.
MonadReader s m =>
Getting (First a) s a -> m (Maybe a)
preview (Text -> Traversal' Value Value
forall t. AsValue t => Text -> Traversal' t Value
key Text
"name" ((Value -> Const (First Text) Value)
 -> Value -> Const (First Text) Value)
-> Getting (First Text) Value Text
-> Getting (First Text) Value Text
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Getting (First Text) Value Text
forall t. AsPrimitive t => Prism' t Text
_String) TaskCall
tc.moduleObject)

formatResult :: Bool -> Pid -> (Int, Value) -> String
formatResult :: Bool -> Pid -> (Int, Value) -> String
formatResult Bool
withColor Pid
pid (Int
code, Value
val) = String
pre String -> ShowS
forall a. Semigroup a => a -> a -> a
<> Text -> String
forall source target. From source target => source -> target
from Text
txt String -> ShowS
forall a. Semigroup a => a -> a -> a
<> String
post
  where
    (String
pre, String
post)
      | Bool
withColor = ([SGR] -> String
setSGRCode [ConsoleLayer -> ColorIntensity -> Color -> SGR
SetColor ConsoleLayer
Foreground ColorIntensity
Dull Color
resColor], [SGR] -> String
setSGRCode [SGR
Reset])
      | Bool
otherwise = (String
"", String
"")
    txt :: Text
txt = Text
res Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<> Text
": [" Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<> Text
host Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<> Text
"] " Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<> Pid -> Text
formatPid Pid
pid Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<> Text
" => " Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<> Text
jsonDump Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<> Text
"\n"
    jsonDump :: Text
jsonDump = ByteString -> Text
decodeUtf8 (ByteString -> Text) -> (Value -> ByteString) -> Value -> Text
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ByteString -> ByteString
forall source target. From source target => source -> target
from (ByteString -> ByteString)
-> (Value -> ByteString) -> Value -> ByteString
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Value -> ByteString
forall a. ToJSON a => a -> ByteString
encode (Value -> ByteString) -> (Value -> Value) -> Value -> ByteString
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Value -> Value
cleanVar (Value -> Text) -> Value -> Text
forall a b. (a -> b) -> a -> b
$ Value
val
    host :: Text
host = Text -> Maybe Text -> Text
forall a. a -> Maybe a -> a
fromMaybe Text
"unknown?!" (Getting (First Text) Value Text -> Value -> Maybe Text
forall s (m :: * -> *) a.
MonadReader s m =>
Getting (First a) s a -> m (Maybe a)
preview (Text -> Traversal' Value Value
forall t. AsValue t => Text -> Traversal' t Value
key Text
"__haxible_play" ((Value -> Const (First Text) Value)
 -> Value -> Const (First Text) Value)
-> Getting (First Text) Value Text
-> Getting (First Text) Value Text
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Text -> Traversal' Value Value
forall t. AsValue t => Text -> Traversal' t Value
key Text
"hosts" ((Value -> Const (First Text) Value)
 -> Value -> Const (First Text) Value)
-> Getting (First Text) Value Text
-> Getting (First Text) Value Text
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Getting (First Text) Value Text
forall t. AsPrimitive t => Prism' t Text
_String) Value
val)
    resColor :: Color
resColor
      | Int
code Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
/= Int
0 = Color
Red
      | Bool
otherwise = Color
Green
    res :: Text
res
      | Int
code Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
0 = Text
"ok"
      | Bool
otherwise = String -> Text
forall source target. From source target => source -> target
from (Int -> String
forall a. Show a => a -> String
show Int
code)

cleanVar :: Value -> Value
cleanVar :: Value -> Value
cleanVar = \case
  Object Object
obj -> Object -> Value
Object (Object -> Value) -> Object -> Value
forall a b. (a -> b) -> a -> b
$ (Key -> Value -> Bool) -> Object -> Object
forall v. (Key -> v -> Bool) -> KeyMap v -> KeyMap v
Data.Aeson.KeyMap.filterWithKey (\Key
k Value
_ -> Key
k Key -> [Key] -> Bool
forall (t :: * -> *) a. (Foldable t, Eq a) => a -> t a -> Bool
`notElem` [Key]
addedKey) Object
obj
  Value
x -> Value
x
  where
    -- TODO: keep in sync with the wrapper and the data source
    addedKey :: [Key]
addedKey =
      [Key
"__haxible_play", Key
"__haxible_start", Key
"__haxible_end", Key
"__haxible_module"]
        [Key] -> [Key] -> [Key]
forall a. Semigroup a => a -> a -> a
<> [Key
"__haxible_multi_hosts"]
        [Key] -> [Key] -> [Key]
forall a. Semigroup a => a -> a -> a
<> [Key
"_ansible_no_log", Key
"_ansible_verbose_always"]

-- | Creates the Python interpreters.
withConnections :: Int -> FilePath -> (Connections -> IO ()) -> IO ()
withConnections :: Int -> String -> (Connections -> IO ()) -> IO ()
withConnections Int
count String
inventory Connections -> IO ()
callback =
  IO (Pool (Process Handle Handle ()))
-> (Pool (Process Handle Handle ()) -> IO ())
-> (Pool (Process Handle Handle ()) -> IO ())
-> IO ()
forall a b c. IO a -> (a -> IO b) -> (a -> IO c) -> IO c
bracket (PoolConfig (Process Handle Handle ())
-> IO (Pool (Process Handle Handle ()))
forall a. PoolConfig a -> IO (Pool a)
Data.Pool.newPool PoolConfig (Process Handle Handle ())
poolConfig) Pool (Process Handle Handle ()) -> IO ()
forall a. Pool a -> IO ()
Data.Pool.destroyAllResources Pool (Process Handle Handle ()) -> IO ()
go
  where
    go :: Pool (Process Handle Handle ()) -> IO ()
go Pool (Process Handle Handle ())
pool = do
      Int
termWidth <- Int -> (String -> Int) -> Maybe String -> Int
forall b a. b -> (a -> b) -> Maybe a -> b
maybe Int
80 String -> Int
forall a. Read a => String -> a
read (Maybe String -> Int) -> IO (Maybe String) -> IO Int
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> String -> IO (Maybe String)
lookupEnv String
"COLUMNS"
      Bool
withColor <- Handle -> IO Bool
hSupportsANSIColor Handle
stdout
      Text -> IO ()
forall (m :: * -> *). MonadIO m => Text -> m ()
say (Int -> Text -> Text
addSep Int
termWidth Text
"PLAY [concurrent]" Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<> Text
"\n")

      let runTask :: TaskCall -> Process Handle Handle () -> IO (Int, Value)
          runTask :: TaskCall -> Process Handle Handle () -> IO (Int, Value)
runTask TaskCall
taskCall Process Handle Handle ()
p = do
            let callParams :: [Value]
callParams = [Vars -> Value
mkObj TaskCall
taskCall.playAttrs, TaskCall
taskCall.moduleObject, Vars -> Value
mkObj TaskCall
taskCall.taskAttrs, Vars -> Value
mkObj TaskCall
taskCall.taskVars]
            Pid
pid <- Pid -> Maybe Pid -> Pid
forall a. a -> Maybe a -> a
fromMaybe (String -> Pid
forall a. HasCallStack => String -> a
error String
"no pid?!") (Maybe Pid -> Pid) -> IO (Maybe Pid) -> IO Pid
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> ProcessHandle -> IO (Maybe Pid)
getPid (Process Handle Handle () -> ProcessHandle
forall stdin stdout stderr.
Process stdin stdout stderr -> ProcessHandle
unsafeProcessHandle Process Handle Handle ()
p)
            Text -> IO ()
forall (m :: * -> *). MonadIO m => Text -> m ()
say (Int -> Text -> Text
addSep Int
termWidth (Pid -> TaskCall -> Text
formatTask Pid
pid TaskCall
taskCall))
            Handle -> ByteString -> IO ()
hPutStrLn (Process Handle Handle () -> Handle
forall stdin stdout stderr. Process stdin stdout stderr -> stdin
getStdin Process Handle Handle ()
p) (ByteString -> ByteString
toStrict (ByteString -> ByteString) -> ByteString -> ByteString
forall a b. (a -> b) -> a -> b
$ [Value] -> ByteString
forall a. ToJSON a => a -> ByteString
encode [Value]
callParams)
            Handle -> IO ()
hFlush (Process Handle Handle () -> Handle
forall stdin stdout stderr. Process stdin stdout stderr -> stdin
getStdin Process Handle Handle ()
p)
            ByteString
output <- Handle -> IO ByteString
hGetLine (Process Handle Handle () -> Handle
forall stdin stdout stderr. Process stdin stdout stderr -> stdout
getStdout Process Handle Handle ()
p)
            case ByteString -> Either String (Int, Value)
forall a. FromJSON a => ByteString -> Either String a
eitherDecodeStrict ByteString
output of
              Right (Int, Value)
res -> do
                String -> IO ()
forall (m :: * -> *). MonadIO m => String -> m ()
sayString (Bool -> Pid -> (Int, Value) -> String
formatResult Bool
withColor Pid
pid (Int, Value)
res)
                (Int, Value) -> IO (Int, Value)
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Int, Value)
res
              Left String
err -> String -> IO (Int, Value)
forall a. HasCallStack => String -> a
error (String -> IO (Int, Value)) -> String -> IO (Int, Value)
forall a b. (a -> b) -> a -> b
$ ByteString -> String
forall a. Show a => a -> String
show ByteString
output String -> ShowS
forall a. Semigroup a => a -> a -> a
<> String
": " String -> ShowS
forall a. Semigroup a => a -> a -> a
<> String
err

      Connections -> IO ()
callback ((TaskCall -> IO (Int, Value)) -> Connections
Connections ((TaskCall -> IO (Int, Value)) -> Connections)
-> (TaskCall -> IO (Int, Value)) -> Connections
forall a b. (a -> b) -> a -> b
$ Pool (Process Handle Handle ())
-> (Process Handle Handle () -> IO (Int, Value)) -> IO (Int, Value)
forall a r. Pool a -> (a -> IO r) -> IO r
Data.Pool.withResource Pool (Process Handle Handle ())
pool ((Process Handle Handle () -> IO (Int, Value)) -> IO (Int, Value))
-> (TaskCall -> Process Handle Handle () -> IO (Int, Value))
-> TaskCall
-> IO (Int, Value)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. TaskCall -> Process Handle Handle () -> IO (Int, Value)
runTask)

    poolConfig :: Data.Pool.PoolConfig (Process Handle Handle ())
    poolConfig :: PoolConfig (Process Handle Handle ())
poolConfig =
      Data.Pool.PoolConfig
        { createResource :: IO (Process Handle Handle ())
createResource = do
            ProcessConfig Handle Handle () -> IO (Process Handle Handle ())
forall (m :: * -> *) stdin stdout stderr.
MonadIO m =>
ProcessConfig stdin stdout stderr
-> m (Process stdin stdout stderr)
startProcess
              (ProcessConfig Handle Handle () -> IO (Process Handle Handle ()))
-> (ProcessConfig () () () -> ProcessConfig Handle Handle ())
-> ProcessConfig () () ()
-> IO (Process Handle Handle ())
forall b c a. (b -> c) -> (a -> b) -> a -> c
. StreamSpec 'STInput Handle
-> ProcessConfig () Handle () -> ProcessConfig Handle Handle ()
forall stdin stdin0 stdout stderr.
StreamSpec 'STInput stdin
-> ProcessConfig stdin0 stdout stderr
-> ProcessConfig stdin stdout stderr
setStdin StreamSpec 'STInput Handle
forall (anyStreamType :: StreamType).
StreamSpec anyStreamType Handle
createPipe
              (ProcessConfig () Handle () -> ProcessConfig Handle Handle ())
-> (ProcessConfig () () () -> ProcessConfig () Handle ())
-> ProcessConfig () () ()
-> ProcessConfig Handle Handle ()
forall b c a. (b -> c) -> (a -> b) -> a -> c
. StreamSpec 'STOutput Handle
-> ProcessConfig () () () -> ProcessConfig () Handle ()
forall stdout stdin stdout0 stderr.
StreamSpec 'STOutput stdout
-> ProcessConfig stdin stdout0 stderr
-> ProcessConfig stdin stdout stderr
setStdout StreamSpec 'STOutput Handle
forall (anyStreamType :: StreamType).
StreamSpec anyStreamType Handle
createPipe
              (ProcessConfig () () () -> IO (Process Handle Handle ()))
-> ProcessConfig () () () -> IO (Process Handle Handle ())
forall a b. (a -> b) -> a -> b
$ String -> [String] -> ProcessConfig () () ()
proc String
"python" [String
"./app/wrapper.py", String
inventory],
          freeResource :: Process Handle Handle () -> IO ()
freeResource = \Process Handle Handle ()
p -> do
            Handle -> IO ()
hClose (Process Handle Handle () -> Handle
forall stdin stdout stderr. Process stdin stdout stderr -> stdin
getStdin Process Handle Handle ()
p)
            Process Handle Handle () -> IO ()
forall (m :: * -> *) stdin stdout stderr.
MonadIO m =>
Process stdin stdout stderr -> m ()
stopProcess Process Handle Handle ()
p,
          poolCacheTTL :: Double
poolCacheTTL = Double
3600,
          poolMaxResources :: Int
poolMaxResources = Int
count
        }