{-# LANGUAGE DataKinds #-}
{-# LANGUAGE TypeApplications #-}

module Test.Sandwich.Contexts.Kubernetes.KubectlLogs (
  withKubectlLogs
  , KubectlLogsContext (..)
  ) where

import Control.Monad
import Control.Monad.IO.Unlift
import Control.Monad.Logger
import Data.String.Interpolate
import qualified Data.Text as T
import Relude hiding (withFile)
import System.FilePath
import Test.Sandwich
import Test.Sandwich.Contexts.Files
import Test.Sandwich.Util.Process (gracefullyStopProcess)
import UnliftIO.Exception
import UnliftIO.IO (withFile)
import UnliftIO.Process


-- * Types

data KubectlLogsContext = KubectlLogsContext {
  KubectlLogsContext -> ProcessHandle
kubectlProcessHandle :: ProcessHandle
  }

-- * Implementation

-- | Run a @kubectl logs@ process, placing the logs in a file in the current test node directory.
--
-- Note that this will stop working if the pod you're talking to goes away (even if you do it against a service).
-- If this happens, a rerun of the command is needed to resume log forwarding.
withKubectlLogs :: (
  MonadLogger m, MonadFail m, MonadUnliftIO m
  , HasBaseContextMonad ctx m, HasFile ctx "kubectl"
  )
  -- | Kubeconfig file
  => FilePath
  -- | Namespace
  -> Text
  -- | Log target (pod, service, etc.)
  -> Text
  -- | Specific container to get logs from
  -> Maybe Text
  -- | Whether to interrupt the process to shut it down while cleaning up
  -> Bool
  -- | Callback receiving the 'KubectlLogsContext'
  -> (KubectlLogsContext -> m a)
  -> m a
withKubectlLogs :: forall (m :: * -> *) ctx a.
(MonadLogger m, MonadFail m, MonadUnliftIO m,
 HasBaseContextMonad ctx m, HasFile ctx "kubectl") =>
FilePath
-> Text
-> Text
-> Maybe Text
-> Bool
-> (KubectlLogsContext -> m a)
-> m a
withKubectlLogs FilePath
kubeConfigFile Text
namespace Text
target Maybe Text
maybeContainer Bool
interruptWhenDone KubectlLogsContext -> m a
action = do
  kubectlBinary <- forall (a :: Symbol) context (m :: * -> *).
(MonadReader context m, HasFile context a) =>
m FilePath
askFile @"kubectl"

  let args = [FilePath
"logs", Text -> FilePath
forall a. ToString a => a -> FilePath
toString Text
target
             , FilePath
"--namespace", Text -> FilePath
forall a. ToString a => a -> FilePath
toString Text
namespace
             , FilePath
"--kubeconfig", FilePath
kubeConfigFile]
             [FilePath] -> [FilePath] -> [FilePath]
forall a. Semigroup a => a -> a -> a
<> ([FilePath] -> (Text -> [FilePath]) -> Maybe Text -> [FilePath]
forall b a. b -> (a -> b) -> Maybe a -> b
maybe [] (\Text
x -> [FilePath
"--container", Text -> FilePath
forall a. ToString a => a -> FilePath
toString Text
x]) Maybe Text
maybeContainer)

  Just dir <- getCurrentFolder
  let logPath = FilePath
dir FilePath -> FilePath -> FilePath
</> Text -> FilePath
forall a. ToString a => a -> FilePath
toString (HasCallStack => Text -> Text -> Text -> Text
Text -> Text -> Text -> Text
T.replace Text
"/" Text
"_" Text
target) FilePath -> FilePath -> FilePath
<.> FilePath
"log"

  debug [i|Running kubectl #{unwords $ fmap toText args} --> #{logPath}|]

  withFile logPath WriteMode $ \Handle
h -> do
    Handle -> BufferMode -> m ()
forall (m :: * -> *). MonadIO m => Handle -> BufferMode -> m ()
hSetBuffering Handle
h BufferMode
LineBuffering

    m (Maybe Handle, Maybe Handle, Maybe Handle, ProcessHandle)
-> ((Maybe Handle, Maybe Handle, Maybe Handle, ProcessHandle)
    -> m ())
-> ((Maybe Handle, Maybe Handle, Maybe Handle, ProcessHandle)
    -> m a)
-> m a
forall (m :: * -> *) a b c.
MonadUnliftIO m =>
m a -> (a -> m b) -> (a -> m c) -> m c
bracket (CreateProcess
-> m (Maybe Handle, Maybe Handle, Maybe Handle, ProcessHandle)
forall (m :: * -> *).
MonadIO m =>
CreateProcess
-> m (Maybe Handle, Maybe Handle, Maybe Handle, ProcessHandle)
createProcess ((FilePath -> [FilePath] -> CreateProcess
proc FilePath
kubectlBinary [FilePath]
args) {
                               std_out = UseHandle h
                               , std_err = UseHandle h
                               , create_group = True
                               }))
            (\(Maybe Handle
_, Maybe Handle
_, Maybe Handle
_, ProcessHandle
ps) -> if
                | Bool
interruptWhenDone -> m () -> m ()
forall (f :: * -> *) a. Functor f => f a -> f ()
void (m () -> m ()) -> m () -> m ()
forall a b. (a -> b) -> a -> b
$ ProcessHandle -> Int -> m ()
forall (m :: * -> *).
(MonadIO m, MonadLogger m) =>
ProcessHandle -> Int -> m ()
gracefullyStopProcess ProcessHandle
ps Int
30_000_000
                | Bool
otherwise -> m ExitCode -> m ()
forall (f :: * -> *) a. Functor f => f a -> f ()
void (m ExitCode -> m ()) -> m ExitCode -> m ()
forall a b. (a -> b) -> a -> b
$ ProcessHandle -> m ExitCode
forall (m :: * -> *). MonadIO m => ProcessHandle -> m ExitCode
waitForProcess ProcessHandle
ps
            )
            (\(Maybe Handle
_, Maybe Handle
_, Maybe Handle
_, ProcessHandle
ps) -> do
                KubectlLogsContext -> m a
action (KubectlLogsContext -> m a) -> KubectlLogsContext -> m a
forall a b. (a -> b) -> a -> b
$ ProcessHandle -> KubectlLogsContext
KubectlLogsContext ProcessHandle
ps
            )