{-# LANGUAGE ConstraintKinds #-}
{-# LANGUAGE DataKinds #-}
{-# LANGUAGE GADTs #-}
{-# LANGUAGE RankNTypes #-}
{-# LANGUAGE TypeApplications #-}
{-# LANGUAGE TypeOperators #-}
{-# OPTIONS_GHC -fno-warn-missing-signatures #-}
{-# OPTIONS_GHC -fno-warn-incomplete-uni-patterns #-}

{-|

Install [Kata Containers](https://katacontainers.io) on a Kubernetes cluster.

-}

module Test.Sandwich.Contexts.Kubernetes.KataContainers (
  -- * Introduce Kata Containers
  introduceKataContainers

  -- * Bracket-style versions
  , withKataContainers
  , withKataContainers'

  -- * Types
  , KataContainersOptions(..)
  , SourceCheckout(..)
  , defaultKataContainersOptions

  , kataContainers
  , KataContainersContext(..)
  , HasKataContainersContext
  ) where

import Control.Lens
import Control.Monad
import Control.Monad.IO.Unlift
import Data.Aeson (FromJSON)
import Data.String.Interpolate
import qualified Data.Text as T
import qualified Data.Text.IO as T
import qualified Data.Yaml as Yaml
import Kubernetes.OpenAPI.Model as Kubernetes
import Kubernetes.OpenAPI.ModelLens as Kubernetes
import Relude hiding (withFile)
import Safe
import System.Exit
import System.FilePath
import Test.Sandwich
import Test.Sandwich.Contexts.Files
import Test.Sandwich.Contexts.Kubernetes.FindImages
import Test.Sandwich.Contexts.Kubernetes.Images
import Test.Sandwich.Contexts.Kubernetes.Kubectl
import Test.Sandwich.Contexts.Kubernetes.Types
import Test.Sandwich.Contexts.Nix
import Test.Sandwich.Waits
import UnliftIO.Process


data KataContainersContext = KataContainersContext {
  kataContainersOptions :: KataContainersOptions
  } deriving (Show)

data SourceCheckout =
  SourceCheckoutFilePath FilePath
  | SourceCheckoutNixDerivation Text
  deriving (Show, Eq)

data KataContainersOptions = KataContainersOptions {
  kataContainersSourceCheckout :: SourceCheckout
  -- | If set, this will overwrite the image in the DaemonSet in @kata-deploy.yaml@ and will set the 'ImagePullPolicy'
  -- to 'IfNotPresent'.
  -- This is useful because it's currently (8\/15\/2024) set to @quay.io\/kata-containers\/kata-deploy:latest@,
  -- with @imagePullPolicy: Always@. This is not reproducible and also doesn't allow us to cache images.
  , kataContainersKataDeployImage :: Maybe Text
  -- | Whether to pull the image using Docker and load it onto the cluster using 'loadImageIfNecessary''.
  , kataContainersPreloadImages :: Bool
  -- | Whether to label the node(s) with @katacontainers.io/kata-runtime=true@, since this seems not to happen
  -- automatically with kata-deploy.
  , kataContainersLabelNode :: Bool
  } deriving (Show)
defaultKataContainersOptions :: KataContainersOptions
defaultKataContainersOptions = KataContainersOptions {
  kataContainersSourceCheckout = SourceCheckoutNixDerivation kataContainersDerivation
  , kataContainersKataDeployImage = Just kataContainersDeployImage
  , kataContainersPreloadImages = True
  , kataContainersLabelNode = True
  }

kataContainers :: Label "kataContainers" KataContainersContext
kataContainers = Label
type HasKataContainersContext context = HasLabel context "kataContainers" KataContainersContext

type ContextWithKataContainers context =
  LabelValue "kataContainers" KataContainersContext
  :> LabelValue "file-kubectl" (EnvironmentFile "kubectl")
  :> context

-- | Install Kata Containers on the cluster and introduce a 'KataContainersContext'.
introduceKataContainers :: (
  Typeable context, KubernetesClusterBasicWithoutReader context m, HasNixContext context
  )
  -- | Options
  => KataContainersOptions
  -> SpecFree (ContextWithKataContainers context) m ()
  -> SpecFree context m ()
introduceKataContainers options = introduceBinaryViaNixPackage @"kubectl" "kubectl" . introduceWith "introduce KataContainers" kataContainers (void . withKataContainers options)

-- | Bracket-style version of 'introduceKataContainers'.
withKataContainers :: forall context m a. (
  HasCallStack, Typeable context, MonadFail m, KubectlBasic context m
  )
  -- | Options
  => KataContainersOptions
  -> (KataContainersContext -> m a)
  -> m a
withKataContainers options action = do
  kcc <- getContext kubernetesCluster
  kubectlBinary <- askFile @"kubectl"
  withKataContainers' kcc kubectlBinary options action

-- | Same as 'withKataContainers', but allows you to pass in the 'KubernetesClusterContext' and @kubectl@ binary path.
withKataContainers' :: forall context m a. (
  HasCallStack, Typeable context, MonadFail m, KubernetesBasic context m
  )
  => KubernetesClusterContext
  -- | Path to @kubectl@ binary
  -> FilePath
  -> KataContainersOptions
  -> (KataContainersContext -> m a)
  -> m a
withKataContainers' kcc@(KubernetesClusterContext {..}) kubectlBinary options@(KataContainersOptions {..}) action = do
  -- Preflight checks
  case kubernetesClusterType of
    KubernetesClusterKind {} -> expectationFailure [i|Can't install Kata Containers on Kind at present.|]
    KubernetesClusterMinikube {..} -> do
      output <- readCreateProcessWithLogging (proc kubernetesClusterTypeMinikubeBinary [
                                                 "--profile", toString kubernetesClusterTypeMinikubeProfileName
                                                 , "ssh", [i|egrep -c 'vmx|svm' /proc/cpuinfo|]
                                                 ]) ""
      case readMay output of
        Just (0 :: Int) -> expectationFailure [i|Preflight check: didn't find "vmx" or "svm" in /proc/cpuinfo. Please make sure virtualization support is enabled.|]
        Just _ -> return ()
        Nothing -> expectationFailure [i|Preflight check: couldn't parse output of minikube ssh "egrep -c 'vmx|svm' /proc/cpuinfo"|]

  -- Get Kata source dir
  kataRoot <- case kataContainersSourceCheckout of
    SourceCheckoutFilePath x -> pure x
    SourceCheckoutNixDerivation d -> getContextMaybe nixContext >>= \case
      Nothing -> expectationFailure [i|Wanted to build Kata Containers source checkout via derivation, but no Nix context was provided.|]
      Just nc -> buildNixCallPackageDerivation' nc d

  info [i|kataRoot: #{kataRoot}|]

  env <- getKubectlEnvironment kcc

  -- Now follow the instructions from
  -- https://github.com/kata-containers/kata-containers/blob/main/docs/install/minikube-installation-guide.md#installing-kata-containers

  -- Read the RBAC and DaemonSet configs
  rbacContents <- liftIO $ T.readFile $ kataRoot </> "tools/packaging/kata-deploy/kata-rbac/base/kata-rbac.yaml"
  deploymentContents' <- liftIO $ T.readFile $ kataRoot </> "tools/packaging/kata-deploy/kata-deploy/base/kata-deploy.yaml"
  let deploymentContents = case kataContainersKataDeployImage of
        Nothing -> deploymentContents'
        Just deployImage -> deploymentContents'
                          & setDaemonSetImage deployImage

  -- Preload any images
  when kataContainersPreloadImages $ do
    let images = findAllImages (rbacContents <> "\n---\n" <> deploymentContents)

    forM_ images $ \image -> do
      info [i|Preloading image: #{image}|]
      loadImageIfNecessary' kcc (ImageLoadSpecDocker image IfNotPresent)

  -- Install kata-deploy
  debug [i|Applying kata-rbac.yaml|]
  createProcessWithLoggingAndStdin ((proc kubectlBinary ["apply", "-f", "-"]) { env = Just env }) (toString rbacContents)
    >>= waitForProcess >>= (`shouldBe` ExitSuccess)
  debug [i|Applying kata-deploy.yaml|]
  createProcessWithLoggingAndStdin ((proc kubectlBinary ["apply", "-f", "-"]) { env = Just env }) (toString deploymentContents)
    >>= waitForProcess >>= (`shouldBe` ExitSuccess)

  debug [i|Waiting for kata-deploy pod to exist|]
  podName <- waitUntil 600 $ do
    pods <- (T.words . toText) <$> readCreateProcessWithLogging ((
      (proc kubectlBinary ["-n", "kube-system"
                          , "get", "pods", "-o=name"]) { env = Just env }
      ) { env = Just env }) ""

    case headMay [t | t <- pods, "pod/kata-deploy" `T.isPrefixOf` t] of
      Just x -> pure x
      Nothing -> expectationFailure [i|Couldn't find kata-deploy pod in: #{pods}|]

  info [i|Got podName: #{podName}|]

  -- Wait until the kata-deploy pod starts sleeping
  waitUntil 600 $ do
    (exitCode, sout, serr) <- readCreateProcessWithExitCode (
      (proc kubectlBinary ["-n", "kube-system"
                          , "exec", toString podName
                          , "--"
                          , "ps", "-ef"
                          ])
      { env = Just env }
      ) ""
    case exitCode of
      ExitSuccess -> return ()
      ExitFailure n -> expectationFailure [i|Command failed with code #{n}. Stderr: #{serr}|]

    toText sout `textShouldContain` "sleep infinity"

  -- Now install the runtime classes
  debug [i|Applying kata-runtimeClasses.yaml|]
  runtimeClassesContents <- liftIO $ T.readFile $ kataRoot </> "tools/packaging/kata-deploy/runtimeclasses/kata-runtimeClasses.yaml"
  createProcessWithLoggingAndStdin ((proc kubectlBinary ["apply", "-f", "-"]) { env = Just env }) (toString runtimeClassesContents)
    >>= waitForProcess >>= (`shouldBe` ExitSuccess)

  -- Finally, label the node(s)
  debug [i|Labeling nodes with katacontainers.io/kata-runtime=true|]
  when kataContainersLabelNode $ do
    createProcessWithLoggingAndStdin ((proc kubectlBinary ["label", "nodes", "--all", "--overwrite", "katacontainers.io/kata-runtime=true"]) { env = Just env }) (toString deploymentContents)
      >>= waitForProcess >>= (`shouldBe` ExitSuccess)

  action $ KataContainersContext options

kataContainersDeployImage :: Text
kataContainersDeployImage = "quay.io/kata-containers/kata-deploy:3.19.1"

-- | Checkout of the @kata-containers@ repo. Currently at release 3.19.1.
kataContainersDerivation :: Text
kataContainersDerivation = [__i|{fetchFromGitHub}:

                                fetchFromGitHub {
                                  owner = "kata-containers";
                                  repo = "kata-containers";
                                  rev = "acae4480ac84701d7354e679714cc9d084b37f44";
                                  sha256 = "sha256-h9Jsto2l1NhQEwIQoecT/D+yt/QbGoqqH/l6NNzJOwk=";
                                }
                               |]

setDaemonSetImage :: Text -> Text -> Text
setDaemonSetImage image = mconcat . fmap setDaemonSetImage' . T.splitOn "---\n"
  where
    setDaemonSetImage' :: Text -> Text
    setDaemonSetImage' (decode -> Right x@(V1DaemonSet {v1DaemonSetKind=(Just "DaemonSet")})) = x
      & set (v1DaemonSetSpecL . _Just . v1DaemonSetSpecTemplateL . v1PodTemplateSpecSpecL . _Just . v1PodSpecContainersL . ix 0 . v1ContainerImageL) (Just image)
      & set (v1DaemonSetSpecL . _Just . v1DaemonSetSpecTemplateL . v1PodTemplateSpecSpecL . _Just . v1PodSpecContainersL . ix 0 . v1ContainerImagePullPolicyL) (Just "IfNotPresent")
      & Yaml.encode
      & decodeUtf8
    setDaemonSetImage' t = t

decode :: FromJSON a => Text -> Either Yaml.ParseException a
decode = Yaml.decodeEither' . encodeUtf8