content: {-# LANGUAGE OverloadedStrings #-}
  {-# LANGUAGE RecordWildCards #-}
  {-# LANGUAGE NamedFieldPuns #-}
  
  module Docker.Cacher where
  
  import qualified Control.Foldl                 as Fold
  import           Control.Monad                  ( guard
                                                  , when
                                                  )
  import           Data.Either                    ( rights )
  import qualified Data.Map.Strict               as Map
  import           Data.Maybe                     ( fromMaybe
                                                  , mapMaybe
                                                  )
  import qualified Data.Text                     as Text
  import           Data.Text                      ( Text )
  import qualified Data.Text.Lazy                as LT
  import           Language.Docker
  import           Language.Docker.Syntax         ( Tag(..) )
  import           Prelude                 hiding ( FilePath )
  import           Turtle
  
  import qualified Docker.Cacher.Inspect
  import           Docker.Cacher.Inspect          ( ImageConfig(..)
                                                  , StageCache(..)
                                                  )
  import           Docker.Cacher.Internal
  import qualified Data.List.NonEmpty
  import qualified Data.Aeson.Text
  import qualified Data.Coerce
  
  {- Glossary:
      - InstructionPos: in the AST for a docker file, each of the lines are described with the type InstructionPos
      - Dockerfile: A list of InstructionPos
      - stage: Each of the FROM instructions in a Dockerfile
      - cache buster: Files inside a docker image that can be compared with files locally under the same path
  -}
  newtype App =
      App Text
      deriving (Show)
  
  newtype Branch =
      Branch Text
      deriving (Show)
  
  newtype BuildOptions =
      BuildOptions Text
      deriving (Show)
  
  newtype CacheLabels =
      CacheLabels [(Text, Text)]
      deriving (Show)
  
  
  -- | Builds the provided Dockerfile. If it is a multi-stage build, check if those stages are already cached
  --   and change the dockerfile to take advantage of that.
  buildFromCache
    :: App
    -> Branch
    -> Maybe Branch
    -> ImageName SourceImage
    -> Maybe BuildOptions
    -> Dockerfile
    -> Shell ()
  buildFromCache app branch fallbackBranch imageName buildOptions ast = do
    -- Inspect the dockerfile and return the stages that got their cache invalidated. We need
    -- them to rewrite the docker file and replace the stages with the ones we have in local cache.
    changedStages <- getChangedStages app branch fallbackBranch ast
  
    -- We replace the busted stages with cached primed ones
    let cachedStages = replaceStages
          (mapMaybe Docker.Cacher.Inspect.alreadyCached changedStages)
          ast
  
    build app imageName buildOptions cachedStages
  
  
  build
    :: App
    -> ImageName SourceImage
    -> Maybe BuildOptions
    -> Dockerfile
    -> Shell ()
  build app imageName buildOptions ast = do
    echo "I'll start building now the main Dockerfile"
  
    -- Build the main docker file which may already have been rewritten to use the
    -- cached stages.
    status <- buildDockerfile app imageName buildOptions ast
  
    case status of
      ExitSuccess ->
        echo
          "I built the main dockerfile without a problem. Now call this same script with the `Cache` mode"
      ExitFailure _ -> die "Boo, I could not build the project"
  
  
  -- | One a dockefile is built, we can extract each of the stages separately and then tag them, so the cache
  -- can be retreived at a later point.
  cacheBuild :: App -> Branch -> Maybe Branch -> Dockerfile -> Shell ()
  cacheBuild app branch fallbackBranch ast = do
    -- Compare the current dockerfile with whatever we have in the cache. If there are any chages,
    -- then we will need to rebuild the cache for each of the changed stages.
    inspectedStages <- getChangedStages app branch fallbackBranch ast
  
    let stagesToBuildFresh = [ stage | NotCached stage <- inspectedStages ]
    let stagesToReBuild =
          [ (uncached, stage)
          | CacheInvalidated uncached stage <- inspectedStages
          ]
    let stagesToBuildFromFallback =
          [ (uncached, cached)
          | FallbackCache uncached (Cached cached _) <- inspectedStages
          ]
  
    when (stagesToBuildFresh /= []) $ do
      echo "--> Let's build the cache for the first time"
      mapM_ (buildAssetStage app) stagesToBuildFresh -- Build cached images for the first time
  
    when (stagesToBuildFromFallback /= []) $ do
      echo "--> Let's build the cache for the first time using a fallback"
      mapM_ (uncurry (reBuildFromFallback app)) stagesToBuildFromFallback
  
    when (stagesToReBuild /= []) $ do
      echo "--> Let's re-build the cache for stages that changed"
      mapM_ (uncurry (reBuildAssetStage app)) stagesToReBuild -- Build each of the stages so they can be reused later
  
  
  -- | Returns a list of stages which needs to either be built separately or that did not have their cached busted
  --   by the introduction of new code.
  getChangedStages
    :: App -> Branch -> Maybe Branch -> Dockerfile -> Shell [StageCache]
  getChangedStages app branch fallbackBranch ast = do
    let
         -- Filter out the main FROM at the end and only
         -- keep the contents of the file before that instruction.
        assetStages = init (getStages ast)
  
        -- For each of the found stages, before the main
        -- FROM instruction, convert them to Stage records
        stages      = mapMaybe (toStage app branch fallbackBranch) assetStages
  
    when (length assetStages > length stages) showAliasesWarning
    fold
      (
          -- Find the stages that we already have in local cache
          Docker.Cacher.Inspect.getAlreadyCached stages
          -- Determine whether or not the cache was invalidated
      >>= uncurry Docker.Cacher.Inspect.shouldBustCache
      )
      Fold.list
   where
    showAliasesWarning = do
      echo "::::::WARNING::::::"
      echo
        "I found some FROM directives in the dockerfile that did not have an `as` alias"
      echo
        "I'm not smart enough to build multi-stage docker files without aliases."
      echo "While this is safe to do, you will get no cache benefits"
      echo ""
      echo
        "Please always write your FROM directives as `FROM image:tag as myalias`"
  
  
  -- | The goal is to create a temporary dockefile in this same folder with the contents
  --   if the stage variable, call docker build with the generated file and tag the image
  --   so we can find it later.
  buildAssetStage :: App -> Stage SourceImage -> Shell ()
  buildAssetStage app Stage {..} = do
    printLn
      ("\n--> Building asset stage " % s % " at line " % d % " for the first time"
      )
      stageName
      stagePos
    let
      fromInstruction = filter isFrom directives
      sourceImage     = ImageName (extractFullName fromInstruction)
      cacheEverything = canCacheDirectives directives
      embeddedFiles =
        if cacheEverything then extractCopiedFiles directives else []
  
      cacheLabels   = buildCacheLabels stageName stageTag embeddedFiles
  
      newDockerfile = toDockerfile $ do
        if cacheEverything then embed directives else embed fromInstruction
        label (Data.Coerce.coerce cacheLabels)
  
    doStageBuild app
                 sourceImage
                 stageImageName
                 buildImageName
                 cacheLabels
                 newDockerfile
   where
    extractFullName (instr : _) = extractFromInstr (instruction instr)
    extractFullName _           = ""
  
    extractFromInstr (From BaseImage { image, tag }) =
      prettyImage image <> prettyTag tag
    extractFromInstr _ = ""
  
    prettyImage (Image Nothing               img) = img
    prettyImage (Image (Just (Registry reg)) img) = reg <> "/" <> img
  
    prettyTag = maybe "" (\(Tag t) -> ":" <> t)
  
  
  -- | The goal is to create a temporary dockefile in this same folder with the contents
  --   if the stage variable, call docker build with the generated file and tag the image
  --   so we can find it later.
  reBuildAssetStage :: App -> Stage SourceImage -> Stage CachedImage -> Shell ()
  reBuildAssetStage app uncached cached = do
    printLn ("\n--> Rebuilding asset stage " % s % " at line " % d)
            (stageName cached)
            (stagePos cached)
    let embeddedFiles = if canCacheDirectives (directives uncached)
          then extractCopiedFiles (directives uncached)
          else []
        cacheLabels = buildCacheLabels (stageName uncached)
                                       (stageTag uncached)
                                       embeddedFiles
    let ImageName t   = stageImageName cached
        newDockerfile = cacheableDockerFile t (directives uncached) cacheLabels
    doStageBuild app
                 (buildImageName cached) -- The source image is the one having the ONBUILD lines
                 (stageImageName cached)
                 (buildImageName cached)
                 cacheLabels
                 newDockerfile
  
  
  reBuildFromFallback :: App -> Stage SourceImage -> Stage CachedImage -> Shell ()
  reBuildFromFallback app uncached cached = do
    let embeddedFiles = extractCopiedFiles (directives uncached)
        cacheLabels =
          buildCacheLabels (stageName uncached) (stageTag uncached) embeddedFiles
    let sourceImage@(ImageName t) = buildImageName cached
        newDockerfile = cacheableDockerFile t (directives uncached) cacheLabels
    doStageBuild app
                 sourceImage
                 (stageImageName uncached)
                 (buildImageName uncached)
                 cacheLabels
                 newDockerfile
  
  
  doStageBuild
    :: App
    -> ImageName source -- ^ This is the image potentially containing the ONBUILD lines, this image needs to exist
    -> ImageName intermediate -- ^ This is the image name to build as intermediate with no ONBUILD
    -> ImageName target -- ^ This is the final image name to build, after appending the ONBUILD lines to intermediate
    -> CacheLabels
    -> Dockerfile
    -> Shell ()
  doStageBuild app sourceImageName intermediateImage targetImageName cacheLabels directives
    = do
      -- Only build the FROM
      status <- buildDockerfile app intermediateImage Nothing directives
  
      -- Break if previous command failed
      guard (status == ExitSuccess)
  
      ImageConfig _ onBuildLines _ <- Docker.Cacher.Inspect.imageConfig
        sourceImageName
  
      -- Append the ONBUILD lines to the new file
      newDockerfile <- createDockerfile intermediateImage cacheLabels onBuildLines
  
      -- Now build it
      finalStatus   <- buildDockerfile app targetImageName Nothing newDockerfile
  
      -- Stop here if previous command failed
      guard (finalStatus == ExitSuccess)
      echo ""
      echo
        "--> I have tagged a cache container that I can use next time to speed builds!"
  
  
  -- | Simply call docker build for the passed arguments
  buildDockerfile
    :: App -> ImageName a -> Maybe BuildOptions -> Dockerfile -> Shell ExitCode
  buildDockerfile (App app) (ImageName imageName) buildOPtions directives = do
    currentDirectory <- pwd
    tmpFile          <- mktempfile currentDirectory "Dockerfile."
    let BuildOptions opts = fromMaybe (BuildOptions "") buildOPtions
    let allBuildOptions =
          [ "build"
            , "--build-arg"
            , "APP_NAME=" <> app
            , "-f"
            , format fp tmpFile
            , "-t"
            , imageName
            , "."
            ]
            <> [opts]
  
    -- Put the Dockerfile contents in the tmp file
    liftIO (writeTextFile tmpFile (LT.toStrict (prettyPrint directives)))
  
    -- Build the generated dockerfile
    shell ("docker " <> Text.intercalate " " allBuildOptions) empty
  
  
  -- | Given a list of instructions, build a dockerfile where the imageName is the FROM for the file and
  --   the list of instructions are wrapped with ONBUILD
  createDockerfile :: ImageName a -> CacheLabels -> [Text] -> Shell Dockerfile
  createDockerfile (ImageName imageName) (CacheLabels cacheLabels) onBuildLines =
    do
      let eitherDirectives = map parseText onBuildLines
          validDirectives  = rights eitherDirectives -- Just in case, filter out bad directives
          file             = toDockerfile $ do
            from $ toImage imageName `tagged` "latest"
            label cacheLabels
            -- Append each of the ONBUILD instructions
            mapM (onBuildRaw . toInstruction) validDirectives
      return file
   where
    toInstruction [InstructionPos inst _ _] = inst
    toInstruction _                         = error "This is not possible"
  
  
  --
  -- | Returns a list of directives grouped by the appeareance of the FROM directive
  --   This will return the group of all stages found in the Dockerfile
  getStages :: Dockerfile -> [Dockerfile]
  getStages ast = filter startsWithFROM (group ast [])
   where
    group [] acc = reverse acc -- End of recursion
    group (directive@(InstructionPos (From _) _ _) : rest) acc =
      group rest ([directive] : acc) -- Append a new group
    group (directive : rest) [] = group rest [[directive]] -- Create a new group
    group (directive : rest) (current : prev) =
      group rest ((current ++ [directive]) : prev) -- Continue the currently open group
  
    --
    -- | Returns true if the first element in the list is a FROM directive
    startsWithFROM (InstructionPos (From _) _ _ : _) = True
    startsWithFROM _ = False
  
  
  -- | Converts a list of instructions into a Stage record
  toStage :: App -> Branch -> Maybe Branch -> Dockerfile -> Maybe (Stage a)
  toStage (App app) branch fallback directives = do
    (stageName, stageTag, stagePos, stageAlias) <- extractInfo directives -- If getStageInfo returns Nothing, skip the rest
    let newImageName (Branch branchName) =
          app <> "__branch__" <> branchName <> "__stage__" <> sanitize stageName
        stageImageName = ImageName (newImageName branch)
        buildImageName = ImageName (newImageName branch <> "-build")
        stageFallbackImage =
          fmap (\br -> ImageName (newImageName br <> "-build")) fallback
    return Stage { .. }
   where
    extractInfo :: Dockerfile -> Maybe (Text, Text, Linenumber, Text)
    extractInfo (InstructionPos { instruction, lineNumber } : _) =
      getStageInfo instruction lineNumber
    extractInfo _ = Nothing
  
    getStageInfo
      :: Instruction Text -> Linenumber -> Maybe (Text, Text, Linenumber, Text)
    getStageInfo (From BaseImage { image, tag = Just (Tag tag), alias = Just (ImageAlias alias) }) pos
      = Just (imageName image, tag, pos, alias)
    getStageInfo (From BaseImage { image, tag = Nothing, alias = Just (ImageAlias alias) }) pos
      = Just (imageName image, "latest", pos, alias)
    getStageInfo _ _ = Nothing
  
    --
    -- | Makes a string safe to use it as a file name
    sanitize = Text.replace "/" "-" . Text.replace ":" "-"
  
  
  -- | Given a list of stages and the AST for a Dockerfile, replace all the FROM instructions
  --   with their corresponding images as described in the Stage record.
  replaceStages :: [Stage CachedImage] -> Dockerfile -> Dockerfile
  replaceStages stages = fmap
    (\InstructionPos {..} ->
      InstructionPos { instruction = replaceStage instruction, .. }
    )
   where
    stagesMap = Map.fromList (map createStagePairs stages)
  
    createStagePairs stage@Stage {..} = (stageAlias, stage)
  
    --
    -- | Find whehter or not we have extracted a stage with the same alias
    --   If we did, then replace the FROM directive with our own version
    replaceStage directive@(From BaseImage { alias = Just (ImageAlias imageAlias) })
      = replaceKnownAlias directive imageAlias
    replaceStage directive = directive
  
    replaceKnownAlias directive imageAlias =
      case Map.lookup imageAlias stagesMap of
        Nothing -> directive
        Just Stage { buildImageName, stageAlias } ->
          let ImageName t = buildImageName
          in  From
                (BaseImage { image    = toImage t
                           , tag      = Just "latest"
                           , alias    = formatAlias stageAlias
                           , digest   = Nothing
                           , platform = Nothing
                           }
                )
  
    formatAlias = Just . fromString . Text.unpack
  
  
  -- | Finds all COPY and ADD instructions in the dockerfile and returns
  -- a concatenated list of all the source paths collected
  extractCopiedFiles :: Dockerfile -> [(SourcePath, TargetPath)]
  extractCopiedFiles = concatMap (extractFiles . instruction)
   where
    extractFiles (Copy CopyArgs { sourcePaths, sourceFlag = NoSource, targetPath })
      = zip (Data.List.NonEmpty.toList sourcePaths) (repeat targetPath)
    extractFiles (Copy CopyArgs { sourceFlag = _ }) = []
    extractFiles (Add AddArgs { sourcePaths, targetPath }) =
      zip (Data.List.NonEmpty.toList sourcePaths) (repeat targetPath)
    extractFiles _ = []
  
  
  buildCacheLabels :: Text -> Text -> [(SourcePath, TargetPath)] -> CacheLabels
  buildCacheLabels imageName imageTag files =
    CacheLabels $ ("cached_image", imageName <> ":" <> imageTag) : case files of
      [] -> []
      _  -> [("cached_files", encodedFiles)]
   where
    encodedFiles = LT.toStrict (Data.Aeson.Text.encodeToLazyText plainTextList)
  
    plainTextList :: [(Text, Text)]
    plainTextList = Data.Coerce.coerce files
  
  
  canCacheDirectives :: Dockerfile -> Bool
  canCacheDirectives df = not (null cacheLabels)
   where
    cacheLabels =
      [ True
      | Label pairs <- map instruction df
      , (key, val)  <- pairs
      , key == "cache_instructions"
      , val == "cache"
      ]
  
  
  cacheableDirectives :: Dockerfile -> Dockerfile
  cacheableDirectives df = if canCacheDirectives df
    then filter (not . isFrom) . filter (not . isOnBuild) $ df
    else []
  
  
  cacheableDockerFile :: Text -> Dockerfile -> CacheLabels -> Dockerfile
  cacheableDockerFile t directives (CacheLabels cacheLabels) = toDockerfile $ do
    -- Use the cached image as base for the new one
    from (toImage t `tagged` "latest")
    -- But we want the contents of the original one
    -- without the ONBUILD
    embed (cacheableDirectives directives)
    label cacheLabels
  
  
  isFrom :: InstructionPos args -> Bool
  isFrom (InstructionPos From{} _ _) = True
  isFrom _                           = False
  
  
  isOnBuild :: InstructionPos args -> Bool
  isOnBuild (InstructionPos OnBuild{} _ _) = True
  isOnBuild _                              = False
  
  
  toImage :: Text -> Image
  toImage = fromString . Text.unpack
  