content: {-# LANGUAGE ConstraintKinds #-}
  {-# LANGUAGE DataKinds #-}
  {-# LANGUAGE GADTs #-}
  {-# LANGUAGE TypeOperators #-}
  
  {-|
  
  Functions for managing images on a Kubernetes cluster.
  
  -}
  
  module Test.Sandwich.Contexts.Kubernetes.Images (
    -- * Introduce a set of images
    introduceImages
  
    -- * Query images
    , getLoadedImages
    , getLoadedImages'
  
    , clusterContainsImage
    , clusterContainsImage'
  
    -- * Load images
    , loadImage
    , loadImage'
  
    -- * Load images if not present
    , loadImageIfNecessary
    , loadImageIfNecessary'
  
    -- * Retry helpers
    , withImageLoadRetry
    , withImageLoadRetry'
  
    -- * Util
    , findAllImages
    , findAllImages'
  
    -- * Types
    , kubernetesClusterImages
    , HasKubernetesClusterImagesContext
  
    , ImageLoadSpec(..)
    , ImagePullPolicy(..)
    ) where
  
  import Control.Monad.Catch (Handler(..), MonadMask)
  import Control.Monad.IO.Unlift
  import Control.Monad.Logger
  import Control.Retry
  import Data.String.Interpolate
  import Data.Text as T
  import Relude
  import Test.Sandwich
  import Test.Sandwich.Contexts.Kubernetes.FindImages
  import qualified Test.Sandwich.Contexts.Kubernetes.KindCluster as Kind
  import qualified Test.Sandwich.Contexts.Kubernetes.MinikubeCluster.Images as Minikube
  import Test.Sandwich.Contexts.Kubernetes.Types
  import Test.Sandwich.Contexts.Kubernetes.Util.Images
  
  
  -- | Get the images loaded onto the cluster.
  getLoadedImages :: (
    HasCallStack, KubernetesClusterBasic context m
    )
    -- | List of image names
    => m (Set Text)
  getLoadedImages = getContext kubernetesCluster >>= getLoadedImages'
  
  -- | Same as 'getLoadedImages', but allows you to pass in the 'KubernetesClusterContext'.
  getLoadedImages' :: (
    HasCallStack, KubernetesBasic context m
    )
    -- | Cluster context
    => KubernetesClusterContext
    -- | List of image names
    -> m (Set Text)
  getLoadedImages' kcc@(KubernetesClusterContext {kubernetesClusterType, kubernetesClusterName}) = do
    timeAction [i|Getting loaded images|] $ do
      case kubernetesClusterType of
        (KubernetesClusterKind {..}) ->
          Kind.getLoadedImagesKind kcc kubernetesClusterTypeKindClusterDriver kubernetesClusterTypeKindBinary Nothing
          -- Kind.loadImage kindBinary kindClusterName image env
        (KubernetesClusterMinikube {..}) ->
          -- Note: don't pass minikubeFlags here. These are pretty much intended for "minikube start" only.
          -- TODO: clarify the documentation and possibly add an extra field where extra options can be passed
          -- to "minikube image" commands.
          Minikube.getLoadedImagesMinikube kubernetesClusterTypeMinikubeBinary kubernetesClusterName []
  
  -- | Test if a cluster has a given image loaded.
  clusterContainsImage :: (
    HasCallStack, KubernetesClusterBasic context m
    )
    -- | Image
    => Text
    -> m Bool
  clusterContainsImage image = do
    kcc <- getContext kubernetesCluster
    clusterContainsImage' kcc image
  
  -- | Same as 'clusterContainsImage', but allows you to pass in the 'KubernetesClusterContext'.
  clusterContainsImage' :: (
    HasCallStack, MonadUnliftIO m, MonadLogger m
    )
    -- | Cluster context
    => KubernetesClusterContext
    -- | Image
    -> Text
    -> m Bool
  clusterContainsImage' kcc@(KubernetesClusterContext {kubernetesClusterType, kubernetesClusterName}) image = do
    case kubernetesClusterType of
      KubernetesClusterKind {..} ->
        Kind.clusterContainsImageKind kcc kubernetesClusterTypeKindClusterDriver kubernetesClusterTypeKindBinary kubernetesClusterTypeKindClusterEnvironment image
      KubernetesClusterMinikube {..} ->
        Minikube.clusterContainsImageMinikube kubernetesClusterTypeMinikubeBinary kubernetesClusterName [] image
  
  -- | Same as 'loadImage', but first checks if the given image is already present on the cluster.
  loadImageIfNecessary :: (
    HasCallStack, MonadFail m, KubernetesClusterBasic context m
    )
    -- | Image load spec
    => ImageLoadSpec
    -> m ()
  loadImageIfNecessary image = do
    kcc <- getContext kubernetesCluster
    loadImageIfNecessary' kcc image
  
  -- | Same as 'loadImage', but allows you to pass in the 'KubernetesClusterContext'.
  loadImageIfNecessary' :: (
    HasCallStack, MonadFail m, KubernetesBasic context m
    )
    -- | Cluster context
    => KubernetesClusterContext
    -- | Image load spec
    -> ImageLoadSpec
    -- | The transformed image name
    -> m ()
  loadImageIfNecessary' kcc imageLoadSpec = do
    unlessM (imageLoadSpecToImageName imageLoadSpec >>= clusterContainsImage' kcc) $
      void $ loadImage' kcc imageLoadSpec
  
  -- | Load an image into a Kubernetes cluster. This will load the image onto the cluster
  -- and return the modified image name (i.e. the name by which the cluster knows the image).
  loadImage :: (
    HasCallStack, MonadFail m, KubernetesClusterBasic context m
    )
    -- | Image load spec
    => ImageLoadSpec
    -- | The loaded image name
    -> m Text
  loadImage image = do
    kcc <- getContext kubernetesCluster
    loadImage' kcc image
  
  -- | Same as 'loadImage', but allows you to pass in the 'KubernetesClusterContext'.
  loadImage' :: (
    HasCallStack, MonadFail m, KubernetesBasic context m
    )
    -- | Cluster context
    => KubernetesClusterContext
    -- | Image load spec
    -> ImageLoadSpec
    -- | The loaded image name
    -> m Text
  loadImage' (KubernetesClusterContext {kubernetesClusterType, kubernetesClusterName}) imageLoadSpec = do
    debug [i|Loading container image '#{imageLoadSpec}'|]
    timeAction [i|Loading container image '#{imageLoadSpec}'|] $ do
      case kubernetesClusterType of
        (KubernetesClusterKind {..}) ->
          Kind.loadImageKind kubernetesClusterTypeKindBinary kubernetesClusterTypeKindClusterName imageLoadSpec kubernetesClusterTypeKindClusterEnvironment
        (KubernetesClusterMinikube {..}) ->
          -- Don't pass minikubeFlags; see comment above.
          Minikube.loadImageMinikube kubernetesClusterTypeMinikubeBinary kubernetesClusterName [] imageLoadSpec
  
          -- Because of the possible silent failure in "minikube image load", confirm that this
          -- image made it onto the cluster.
          -- At the moment this approach doesn't work, because if you do
          -- "minikube image load busybox:1.36.1-musl"
          -- followed by
          -- "minikube image ls",
          -- the result contains "docker.io/library/busybox:1.36.1-musl".
          -- Where did the docker.io/library/ come from? Need to understand this before we can
          -- check this properly.
          --
          -- image' <- Minikube.loadImage minikubeBinary kubernetesClusterName minikubeFlags image
          -- loadedImages <- Set.toList <$> getLoadedImages' kcc
          -- loadedImages `shouldContain` [image']
          -- return image'
  
  -- | A combinator to wrap around your 'loadImage' or 'loadImageIfNecessary' calls to retry
  -- on failure. Image loads sometimes fail on Minikube (version 1.33.0 at time of writing).
  withImageLoadRetry :: (MonadLoggerIO m, MonadMask m) => ImageLoadSpec -> m a -> m a
  withImageLoadRetry = withImageLoadRetry' (exponentialBackoff 50000 <> limitRetries 5)
  
  -- | Same as 'withImageLoadRetry', but allows you to specify the retry policy.
  withImageLoadRetry' :: (MonadLoggerIO m, MonadMask m) => RetryPolicyM m -> ImageLoadSpec -> m a -> m a
  withImageLoadRetry' policy ils action =
    recovering policy [\_status -> Handler (\(e :: FailureReason) -> do
                                               warn [i|#{ils}: retrying load due to exception: #{e}|]
                                               return True)] $ \_ ->
      action
  
  
  -- | Introduce a list of images into a Kubernetes cluster.
  -- Stores the list of transformed image names under the "kubernetesClusterImages" label.
  introduceImages :: (
    HasCallStack, KubernetesClusterBasicWithoutReader context m
    )
    -- | Images to load
    => [ImageLoadSpec]
    -> SpecFree (LabelValue "kubernetesClusterImages" [Text] :> context) m ()
    -> SpecFree context m ()
  introduceImages images = introduceWith "Introduce cluster images" kubernetesClusterImages $ \action ->
    forM images (\x -> loadImage x) >>= (void . action)
  