content: {-# LANGUAGE ConstraintKinds #-}
  {-# LANGUAGE DataKinds #-}
  {-# LANGUAGE GADTs #-}
  {-# LANGUAGE TypeOperators #-}
  
  module Test.Sandwich.Contexts.Kubernetes.Util.Images (
    dockerPullIfNecessary
    , isDockerImagePresent
  
    , podmanPullIfNecessary
    , isPodmanImagePresent
  
    , readImageName
    , readUncompressedImageName
    , imageLoadSpecToImageName
    ) where
  
  import Control.Monad.IO.Unlift
  import Control.Monad.Logger
  import Data.Aeson as A
  import qualified Data.ByteString.Lazy as BL
  import Data.String.Interpolate
  import qualified Data.Text as T
  import qualified Data.Vector as V
  import Relude
  import Safe
  import System.Exit
  import System.FilePath
  import Test.Sandwich
  import Test.Sandwich.Contexts.Kubernetes.Types
  import Test.Sandwich.Contexts.Kubernetes.Util.Aeson
  import UnliftIO.Directory
  import UnliftIO.Process
  import UnliftIO.Temporary
  
  
  -- * Docker
  
  -- | Pull an image using Docker if it isn't already present.
  -- Returns 'True' if a pull was done.
  dockerPullIfNecessary :: (MonadUnliftIO m, MonadLoggerIO m) => Text -> ImagePullPolicy -> m Bool
  dockerPullIfNecessary = commonPullIfNecessary "docker"
  
  isDockerImagePresent :: (MonadUnliftIO m, MonadLoggerIO m) => Text -> m Bool
  isDockerImagePresent = isImagePresentCommon "docker"
  
  -- * Podman
  
  -- | Pull an image using Docker if it isn't already present.
  -- Returns 'True' if a pull was done.
  podmanPullIfNecessary :: (MonadUnliftIO m, MonadLoggerIO m) => Text -> ImagePullPolicy -> m Bool
  podmanPullIfNecessary = commonPullIfNecessary "podman"
  
  isPodmanImagePresent :: (MonadUnliftIO m, MonadLoggerIO m) => Text -> m Bool
  isPodmanImagePresent = isImagePresentCommon "podman"
  
  -- * Common
  
  commonPullIfNecessary :: (MonadUnliftIO m, MonadLoggerIO m) => String -> Text -> ImagePullPolicy -> m Bool
  commonPullIfNecessary binary image pullPolicy = isImagePresentCommon binary image >>= \case
    True ->
      if | pullPolicy == Always -> doPull
         | otherwise -> return False
    False ->
      if | pullPolicy == Never -> expectationFailure [i|Docker pull policy was "Never" but image wasn't present: '#{image}'|]
         | otherwise -> doPull
    where
      doPull = do
        createProcessWithLogging (proc binary ["pull", toString image])
          >>= waitForProcess >>= (`shouldBe` ExitSuccess)
        return True
  
  isImagePresentCommon :: (MonadUnliftIO m, MonadLoggerIO m) => String -> Text -> m Bool
  isImagePresentCommon binary image = do
    createProcessWithLogging (proc binary ["inspect", "--type=image", toString image]) >>= waitForProcess >>= \case
      ExitSuccess -> return True
      ExitFailure _ -> return False
  
  -- * Image name reading
  
  readImageName :: (HasCallStack, MonadUnliftIO m, MonadLogger m) => FilePath -> m Text
  readImageName path = doesDirectoryExist path >>= \case
    True -> readUncompressedImageName path
    False -> case takeExtension path of
      ".tar" -> extractFromTarball
      ".gz" -> extractFromTarball
      _ -> expectationFailure [i|readImageName: unexpected extension in #{path}. Wanted .tar, .tar.gz, or uncompressed directory.|]
    where
      extractFromTarball = do
        files <- readCreateProcessWithLogging (proc "tar" ["tf", path]) ""
        manifestFileName <- case headMay [t | t <- T.words (toText files), "manifest.json" `T.isInfixOf` t] of
          Just f -> pure $ toString $ T.strip f
          Nothing -> expectationFailure [i|readImageName: couldn't find manifest file in #{path}|]
  
        withSystemTempDirectory "manifest.json" $ \dir -> do
          _ <- readCreateProcessWithLogging ((proc "tar" ["xvf", path, manifestFileName]) { cwd = Just dir }) ""
          liftIO (BL.readFile (dir </> "manifest.json")) >>= getImageNameFromManifestJson path
  
  readUncompressedImageName :: (HasCallStack, MonadIO m) => FilePath -> m Text
  readUncompressedImageName path = liftIO (BL.readFile (path </> "manifest.json")) >>= getImageNameFromManifestJson path
  
  getImageNameFromManifestJson :: (HasCallStack, MonadIO m) => FilePath -> LByteString -> m Text
  getImageNameFromManifestJson path contents = do
    case A.eitherDecode contents of
      Left err -> expectationFailure [i|Couldn't decode manifest.json: #{err}|]
      Right (A.Array entries) -> case concatMap getRepoTags entries of
        (x:_) -> pure x
        [] -> expectationFailure [i|Didn't find a repo tag for image at #{path}|]
      Right x -> expectationFailure [i|Unexpected manifest.json format: #{x}|]
  
    where
      getRepoTags :: A.Value -> [Text]
      getRepoTags (A.Object (aesonLookup "RepoTags" -> Just (A.Array repoItems))) = [t | A.String t <- V.toList repoItems]
      getRepoTags _ = []
  
  imageLoadSpecToImageName :: (MonadUnliftIO m, MonadLogger m) => ImageLoadSpec -> m Text
  imageLoadSpecToImageName (ImageLoadSpecTarball image) = readImageName image
  imageLoadSpecToImageName (ImageLoadSpecDocker image _) = pure image
  imageLoadSpecToImageName (ImageLoadSpecPodman image _) = pure image
  