{-# LANGUAGE ScopedTypeVariables #-}

{-| Effectful functions that create and convert disk image files. -}
module B9.DiskImageBuilder
  ( materializeImageSource
  , substImageTarget
  , preferredDestImageTypes
  , preferredSourceImageTypes
  , resolveImageSource
  , createDestinationImage
  , resizeImage
  , importImage
  , exportImage
  , exportAndRemoveImage
  , convertImage
  , shareImage
  , ensureAbsoluteImageDirExists
  , pushSharedImageLatestVersion
  , lookupSharedImages
  , getSharedImages
  , getSharedImagesCacheDir
  , getSelectedRepos
  , pullRemoteRepos
  , pullLatestImage
  ) where

import B9.B9Config
import B9.B9Monad
import B9.Artifact.Content.StringTemplate
import B9.DiskImages
import qualified B9.PartitionTable as P
import B9.Repository
import B9.RepositoryIO
import Control.Exception
import Control.Lens ((^.))
import Control.Monad
import Control.Monad.IO.Class
import Data.Data
import Data.Function
import Data.Generics.Aliases
import Data.Generics.Schemes
import Data.List
import Data.Maybe
import System.Directory
import System.FilePath
import System.IO.B9Extras (consult, ensureDir, prettyPrintToFile)
import System.IO.Error (isDoesNotExistError)
import Text.Printf (printf)
import Text.Show.Pretty (ppShow)

-- -- | Convert relative file paths of images, sources and mounted host directories
-- -- to absolute paths relative to '_buildDirRoot'.
-- makeImagePathsAbsoluteToBuildDirRoot :: ImageTarget -> B9 ImageTarget
-- makeImagePathsAbsoluteToBuildDirRoot img =
--   getConfig >>= maybe (return img) (return . go) . _buildDirRoot
--   where
--     go rootDir = everywhere mkAbs img
--       where mkAbs = mkT
-- | Replace $... variables inside an 'ImageTarget'
substImageTarget :: Environment -> ImageTarget -> ImageTarget
substImageTarget env = everywhere gsubst
  where
    gsubst :: Data a => a -> a
    gsubst = mkT substMountPoint `extT` substImage `extT` substImageSource `extT` substDiskTarget
    substMountPoint NotMounted = NotMounted
    substMountPoint (MountPoint x) = MountPoint (sub x)
    substImage (Image fp t fs) = Image (sub fp) t fs
    substImageSource (From n s) = From (sub n) s
    substImageSource (EmptyImage l f t s) = EmptyImage (sub l) f t s
    substImageSource s = s
    substDiskTarget (Share n t s) = Share (sub n) t s
    substDiskTarget (LiveInstallerImage name outDir resize) = LiveInstallerImage (sub name) (sub outDir) resize
    substDiskTarget s = s
    sub = subst env

-- | Resolve an ImageSource to an 'Image'. The ImageSource might
-- not exist, as is the case for 'EmptyImage'.
resolveImageSource :: ImageSource -> B9 Image
resolveImageSource src =
  case src of
    (EmptyImage fsLabel fsType imgType _size) ->
      let img = Image fsLabel imgType fsType
       in return (changeImageFormat imgType img)
    (SourceImage srcImg _part _resize) -> ensureAbsoluteImageDirExists srcImg
    (CopyOnWrite backingImg) -> ensureAbsoluteImageDirExists backingImg
    (From name _resize) ->
      getLatestImageByName (SharedImageName name) >>=
      maybe (errorExitL (printf "Nothing found for %s." (show (SharedImageName name)))) ensureAbsoluteImageDirExists

-- | Return all valid image types sorted by preference.
preferredDestImageTypes :: ImageSource -> B9 [ImageType]
preferredDestImageTypes src =
  case src of
    (CopyOnWrite (Image _file fmt _fs)) -> return [fmt]
    (EmptyImage _label NoFileSystem fmt _size) -> return (nub [fmt, Raw, QCow2, Vmdk])
    (EmptyImage _label _fs _fmt _size) -> return [Raw]
    (SourceImage _img (Partition _) _resize) -> return [Raw]
    (SourceImage (Image _file fmt _fs) _pt resize) ->
      return (nub [fmt, Raw, QCow2, Vmdk] `intersect` allowedImageTypesForResize resize)
    (From name resize) ->
      getLatestImageByName (SharedImageName name) >>=
      maybe
        (errorExitL (printf "Nothing found for %s." (show (SharedImageName name))))
        (\sharedImg -> preferredDestImageTypes (SourceImage sharedImg NoPT resize))

-- | Return all supported source 'ImageType's compatible to a 'ImageDestinaion'
-- in the preferred order.
preferredSourceImageTypes :: ImageDestination -> [ImageType]
preferredSourceImageTypes dest =
  case dest of
    (Share _ fmt resize) -> nub [fmt, Raw, QCow2, Vmdk] `intersect` allowedImageTypesForResize resize
    (LocalFile (Image _ fmt _) resize) -> nub [fmt, Raw, QCow2, Vmdk] `intersect` allowedImageTypesForResize resize
    Transient -> [Raw, QCow2, Vmdk]
    (LiveInstallerImage _name _repo _imgResize) -> [Raw]

allowedImageTypesForResize :: ImageResize -> [ImageType]
allowedImageTypesForResize r =
  case r of
    Resize _ -> [Raw]
    ShrinkToMinimum -> [Raw]
    _ -> [Raw, QCow2, Vmdk]

-- | Create the parent directories for the file that contains the 'Image'.
-- If the path to the image file is relative, prepend '_buildDirRoot' from
-- the 'B9Config'.
ensureAbsoluteImageDirExists :: Image -> B9 Image
ensureAbsoluteImageDirExists img@(Image path _ _) = do
  b9cfg <- getConfig
  let dir =
        let dirRel = takeDirectory path
         in if isRelative dirRel
              then let prefix = fromMaybe "." (b9cfg ^. buildDirRoot)
                    in prefix </> dirRel
              else dirRel
  liftIO $ do
    createDirectoryIfMissing True dir
    dirAbs <- canonicalizePath dir
    return $ changeImageDirectory dirAbs img

-- | Create an image from an image source. The destination image must have a
-- compatible image type and filesystem. The directory of the image MUST be
-- present and the image file itself MUST NOT alredy exist.
materializeImageSource :: ImageSource -> Image -> B9 ()
materializeImageSource src dest =
  case src of
    (EmptyImage fsLabel fsType _imgType size) ->
      let (Image _ imgType _) = dest
       in createEmptyImage fsLabel fsType imgType size dest
    (SourceImage srcImg part resize) -> createImageFromImage srcImg part resize dest
    (CopyOnWrite backingImg) -> createCOWImage backingImg dest
    (From name resize) ->
      getLatestImageByName (SharedImageName name) >>=
      maybe
        (errorExitL (printf "Nothing found for %s." (show (SharedImageName name))))
        (\sharedImg -> materializeImageSource (SourceImage sharedImg NoPT resize) dest)

createImageFromImage :: Image -> Partition -> ImageResize -> Image -> B9 ()
createImageFromImage src part size out = do
  importImage src out
  extractPartition part out
  resizeImage size out
  where
    extractPartition :: Partition -> Image -> B9 ()
    extractPartition NoPT _ = return ()
    extractPartition (Partition partIndex) (Image outFile Raw _) = do
      (start, len, blockSize) <- liftIO (P.getPartition partIndex outFile)
      let tmpFile = outFile <.> "extracted"
      dbgL (printf "Extracting partition %i from '%s'" partIndex outFile)
      cmd (printf "dd if='%s' of='%s' bs=%i skip=%i count=%i &> /dev/null" outFile tmpFile blockSize start len)
      cmd (printf "mv '%s' '%s'" tmpFile outFile)
    extractPartition (Partition partIndex) (Image outFile fmt _) =
      error
        (printf "Extract partition %i from image '%s': Invalid format %s" partIndex outFile (imageFileExtension fmt))

-- | Convert some 'Image', e.g. a temporary image used during the build phase
-- to the final destination.
createDestinationImage :: Image -> ImageDestination -> B9 ()
createDestinationImage buildImg dest =
  case dest of
    (Share name imgType imgResize) -> do
      resizeImage imgResize buildImg
      let shareableImg = changeImageFormat imgType buildImg
      exportAndRemoveImage buildImg shareableImg
      void (shareImage shareableImg (SharedImageName name))
    (LocalFile destImg imgResize) -> do
      resizeImage imgResize buildImg
      exportAndRemoveImage buildImg destImg
    (LiveInstallerImage name repo imgResize) -> do
      resizeImage imgResize buildImg
      let destImg = Image destFile Raw buildImgFs
          (Image _ _ buildImgFs) = buildImg
          destFile = repo </> "machines" </> name </> "disks" </> "raw" </> "0.raw"
          sizeFile = repo </> "machines" </> name </> "disks" </> "raw" </> "0.size"
          versFile = repo </> "machines" </> name </> "disks" </> "raw" </> "VERSION"
      exportAndRemoveImage buildImg destImg
      cmd
        (printf
           "echo $(qemu-img info -f raw '%s' | gawk -e '/virtual size/ {print $4}' | tr -d '(') > '%s'"
           destFile
           sizeFile)
      buildDate <- getBuildDate
      buildId <- getBuildId
      liftIO (writeFile versFile (buildId ++ "-" ++ buildDate))
    Transient -> return ()

createEmptyImage :: String -> FileSystem -> ImageType -> ImageSize -> Image -> B9 ()
createEmptyImage fsLabel fsType imgType imgSize dest@(Image _ imgType' fsType')
  | fsType /= fsType' =
    error
      (printf
         "Conflicting createEmptyImage parameters. Requested is file system %s but the destination image has %s."
         (show fsType)
         (show fsType'))
  | imgType /= imgType' =
    error
      (printf
         "Conflicting createEmptyImage parameters. Requested is image type %s but the destination image has type %s."
         (show imgType)
         (show imgType'))
  | otherwise = do
    let (Image imgFile imgFmt imgFs) = dest
        qemuImgOpts = conversionOptions imgFmt
    dbgL
      (printf "Creating empty raw image '%s' with size %s and options %s" imgFile (toQemuSizeOptVal imgSize) qemuImgOpts)
    cmd
      (printf
         "qemu-img create -f %s %s '%s' '%s'"
         (imageFileExtension imgFmt)
         qemuImgOpts
         imgFile
         (toQemuSizeOptVal imgSize))
    case (imgFmt, imgFs) of
      (Raw, Ext4_64) -> do
        let fsCmd = "mkfs.ext4"
        dbgL (printf "Creating file system %s" (show imgFs))
        cmd (printf "%s -F -L '%s' -O 64bit -q '%s'" fsCmd fsLabel imgFile)
      (Raw, Ext4) -> do
        let fsCmd = "mkfs.ext4"
        dbgL (printf "Creating file system %s" (show imgFs))
        cmd (printf "%s -F -L '%s' -O ^64bit -q '%s'" fsCmd fsLabel imgFile)
      (it, fs) -> error (printf "Cannot create file system %s in image type %s" (show fs) (show it))

createCOWImage :: Image -> Image -> B9 ()
createCOWImage (Image backingFile _ _) (Image imgOut imgFmt _) = do
  dbgL (printf "Creating COW image '%s' backed by '%s'" imgOut backingFile)
  cmd (printf "qemu-img create -f %s -o backing_file='%s' '%s'" (imageFileExtension imgFmt) backingFile imgOut)

-- | Resize an image, including the file system inside the image.
resizeImage :: ImageResize -> Image -> B9 ()
resizeImage KeepSize _ = return ()
resizeImage (Resize newSize) (Image img Raw fs)
  | fs == Ext4 || fs == Ext4_64 = do
    let sizeOpt = toQemuSizeOptVal newSize
    dbgL (printf "Resizing ext4 filesystem on raw image to %s" sizeOpt)
    cmd (printf "e2fsck -p '%s'" img)
    cmd (printf "resize2fs -f '%s' %s" img sizeOpt)
resizeImage (ResizeImage newSize) (Image img _ _) = do
  let sizeOpt = toQemuSizeOptVal newSize
  dbgL (printf "Resizing image to %s" sizeOpt)
  cmd (printf "qemu-img resize -q '%s' %s" img sizeOpt)
resizeImage ShrinkToMinimum (Image img Raw fs)
  | fs == Ext4 || fs == Ext4_64 = do
    dbgL "Shrinking image to minimum size"
    cmd (printf "e2fsck -p '%s'" img)
    cmd (printf "resize2fs -f -M '%s'" img)
resizeImage _ img = error (printf "Invalid image type or filesystem, cannot resize image: %s" (show img))

-- | Import a disk image from some external source into the build directory
-- if necessary convert the image.
importImage :: Image -> Image -> B9 ()
importImage imgIn imgOut@(Image imgOutPath _ _) = do
  alreadyThere <- liftIO (doesFileExist imgOutPath)
  unless alreadyThere (convert False imgIn imgOut)

-- | Export a disk image from the build directory; if necessary convert the image.
exportImage :: Image -> Image -> B9 ()
exportImage = convert False

-- | Export a disk image from the build directory; if necessary convert the image.
exportAndRemoveImage :: Image -> Image -> B9 ()
exportAndRemoveImage = convert True

-- | Convert an image in the build directory to another format and return the new image.
convertImage :: Image -> Image -> B9 ()
convertImage imgIn imgOut@(Image imgOutPath _ _) = do
  alreadyThere <- liftIO (doesFileExist imgOutPath)
  unless alreadyThere (convert True imgIn imgOut)

-- | Convert/Copy/Move images
convert :: Bool -> Image -> Image -> B9 ()
convert doMove (Image imgIn fmtIn _) (Image imgOut fmtOut _)
  | imgIn == imgOut = do
    ensureDir imgOut
    dbgL (printf "No need to convert: '%s'" imgIn)
  | doMove && fmtIn == fmtOut = do
    ensureDir imgOut
    dbgL (printf "Moving '%s' to '%s'" imgIn imgOut)
    liftIO (renameFile imgIn imgOut)
  | otherwise = do
    ensureDir imgOut
    dbgL
      (printf "Converting %s to %s: '%s' to '%s'" (imageFileExtension fmtIn) (imageFileExtension fmtOut) imgIn imgOut)
    cmd
      (printf
         "qemu-img convert -q -f %s -O %s %s '%s' '%s'"
         (imageFileExtension fmtIn)
         (imageFileExtension fmtOut)
         (conversionOptions fmtOut)
         imgIn
         imgOut)
    when doMove $ do
      dbgL (printf "Removing '%s'" imgIn)
      liftIO (removeFile imgIn)

conversionOptions :: ImageType -> String
conversionOptions Vmdk = " -o adapter_type=lsilogic "
conversionOptions QCow2 = " -o compat=1.1,lazy_refcounts=on "
conversionOptions _ = " "

toQemuSizeOptVal :: ImageSize -> String
toQemuSizeOptVal (ImageSize amount u) =
  show amount ++
  case u of
    GB -> "G"
    MB -> "M"
    KB -> "K"
    B -> ""

-- | Publish an sharedImage made from an image and image meta data to the
-- configured repository
shareImage :: Image -> SharedImageName -> B9 SharedImage
shareImage buildImg sname@(SharedImageName name) = do
  sharedImage <- createSharedImageInCache buildImg sname
  infoL (printf "SHARED '%s'" name)
  pushToSelectedRepo sharedImage
  return sharedImage

-- | Return a 'SharedImage' with the current build data and build id from the
-- name and disk image.
getSharedImageFromImageInfo :: SharedImageName -> Image -> B9 SharedImage
getSharedImageFromImageInfo name (Image _ imgType imgFS) = do
  buildId <- getBuildId
  date <- getBuildDate
  return (SharedImage name (SharedImageDate date) (SharedImageBuildId buildId) imgType imgFS)

-- | Convert the disk image and serialize the base image data structure.
-- If the 'maxLocalSharedImageRevisions' configuration is set to @Just n@
-- also delete all but the @n - 1@ newest images from the local cache.
createSharedImageInCache :: Image -> SharedImageName -> B9 SharedImage
createSharedImageInCache img sname@(SharedImageName name) = do
  dbgL (printf "CREATING SHARED IMAGE: '%s' '%s'" (ppShow img) name)
  sharedImg <- getSharedImageFromImageInfo sname img
  dir <- getSharedImagesCacheDir
  convertImage img (changeImageDirectory dir (sharedImageImage sharedImg))
  prettyPrintToFile (dir </> sharedImageFileName sharedImg) sharedImg
  dbgL (printf "CREATED SHARED IMAGE IN CACHE '%s'" (ppShow sharedImg))
  cleanOldSharedImageRevisionsFromCache sname
  return sharedImg

-- | Publish the latest version of a shared image identified by name to the
-- selected repository from the cache.
pushSharedImageLatestVersion :: SharedImageName -> B9 ()
pushSharedImageLatestVersion name@(SharedImageName imgName) =
  getLatestSharedImageByNameFromCache name >>=
  maybe
    (errorExitL (printf "Nothing found for %s." (show imgName)))
    (\sharedImage -> do
       dbgL (printf "PUSHING '%s'" (ppShow sharedImage))
       pushToSelectedRepo sharedImage
       infoL (printf "PUSHED '%s'" imgName))

-- | Upload a shared image from the cache to a selected remote repository
pushToSelectedRepo :: SharedImage -> B9 ()
pushToSelectedRepo i = do
  c <- getSharedImagesCacheDir
  r <- getSelectedRemoteRepo
  when (isJust r) $ do
    let (Image imgFile' _imgType _imgFS) = sharedImageImage i
        cachedImgFile = c </> imgFile'
        cachedInfoFile = c </> sharedImageFileName i
        repoImgFile = sharedImagesRootDirectory </> imgFile'
        repoInfoFile = sharedImagesRootDirectory </> sharedImageFileName i
    pushToRepo (fromJust r) cachedImgFile repoImgFile
    pushToRepo (fromJust r) cachedInfoFile repoInfoFile

-- | Pull metadata files from all remote repositories.
pullRemoteRepos :: B9 ()
pullRemoteRepos = do
  repos <- getSelectedRepos
  mapM_ dl repos
  where
    dl = pullGlob sharedImagesRootDirectory (FileExtension sharedImageFileExtension)

-- | Pull the latest version of an image, either from the selected remote
-- repo or from the repo that has the latest version.
pullLatestImage :: SharedImageName -> B9 (Maybe SharedImageBuildId)
pullLatestImage name@(SharedImageName dbgName) = do
  repos <- getSelectedRepos
  let repoPredicate Cache = False
      repoPredicate (Remote repoId) = repoId `elem` repoIds
      repoIds = map remoteRepoRepoId repos
      hasName sharedImage = name == sharedImageName sharedImage
  candidates <- lookupSharedImages repoPredicate hasName
  let (Remote repoId, image) = last candidates
  if null candidates
    then do
      errorL (printf "No shared image named '%s' on these remote repositories: '%s'" dbgName (ppShow repoIds))
      return Nothing
    else do
      dbgL (printf "PULLING SHARED IMAGE: '%s'" (ppShow image))
      cacheDir <- getSharedImagesCacheDir
      let (Image imgFile' _imgType _fs) = sharedImageImage image
          cachedImgFile = cacheDir </> imgFile'
          cachedInfoFile = cacheDir </> sharedImageFileName image
          repoImgFile = sharedImagesRootDirectory </> imgFile'
          repoInfoFile = sharedImagesRootDirectory </> sharedImageFileName image
          repo = fromJust (lookupRemoteRepo repos repoId)
      pullFromRepo repo repoImgFile cachedImgFile
      pullFromRepo repo repoInfoFile cachedInfoFile
      infoL (printf "PULLED '%s' FROM '%s'" dbgName repoId)
      cleanOldSharedImageRevisionsFromCache name
      return (Just (sharedImageBuildId image))

-- | Return the 'Image' of the latest version of a shared image named 'name'
-- from the local cache.
getLatestImageByName :: SharedImageName -> B9 (Maybe Image)
getLatestImageByName name = do
  sharedImage <- getLatestSharedImageByNameFromCache name
  cacheDir <- getSharedImagesCacheDir
  let image = changeImageDirectory cacheDir . sharedImageImage <$> sharedImage
  case image of
    Just i -> dbgL (printf "USING SHARED SOURCE IMAGE '%s'" (show i))
    Nothing -> errorL (printf "SOURCE IMAGE '%s' NOT FOUND" (show name))
  return image

-- | Return the latest version of a shared image named 'name' from the local cache.
getLatestSharedImageByNameFromCache :: SharedImageName -> B9 (Maybe SharedImage)
getLatestSharedImageByNameFromCache name@(SharedImageName dbgName) = do
  imgs <- lookupSharedImages (== Cache) ((== name) . sharedImageName)
  case reverse imgs of
    (Cache, sharedImage):_rest -> return (Just sharedImage)
    _ -> do
      errorL (printf "No image(s) named '%s' found." dbgName)
      return Nothing

-- | Return a list of all existing sharedImages from cached repositories.
getSharedImages :: B9 [(Repository, [SharedImage])]
getSharedImages = do
  reposAndFiles <- repoSearch sharedImagesRootDirectory (FileExtension sharedImageFileExtension)
  mapM (\(repo, files) -> (repo, ) . catMaybes <$> mapM consult' files) reposAndFiles
  where
    consult' f = do
      r <- liftIO (try (consult f))
      case r of
        Left (e :: SomeException) -> do
          dbgL (printf "Failed to load shared image meta-data from '%s': '%s'" (takeFileName f) (show e))
          dbgL (printf "Removing bad meta-data file '%s'" f)
          liftIO (removeFile f)
          return Nothing
        Right c -> return (Just c)

-- | Find shared images and the associated repos from two predicates. The result
-- is the concatenated result of the sorted shared images satisfying 'imgPred'.
lookupSharedImages :: (Repository -> Bool) -> (SharedImage -> Bool) -> B9 [(Repository, SharedImage)]
lookupSharedImages repoPred imgPred = do
  xs <- getSharedImages
  let rs = [(r, s) | (r, ss) <- xs, s <- ss]
      matchingRepo = filter (repoPred . fst) rs
      matchingImg = filter (imgPred . snd) matchingRepo
      sorted = sortBy (compare `on` snd) matchingImg
  return (mconcat (pure <$> sorted))

-- | Return either all remote repos or just the single selected repo.
getSelectedRepos :: B9 [RemoteRepo]
getSelectedRepos = do
  allRepos <- getRemoteRepos
  selectedRepo <- getSelectedRemoteRepo
  let repos = maybe allRepos return selectedRepo -- 'Maybe' a repo
  return repos

-- | Return the path to the sub directory in the cache that contains files of
-- shared images.
getSharedImagesCacheDir :: B9 FilePath
getSharedImagesCacheDir = do
  cacheDir <- localRepoDir <$> getRepoCache
  return (cacheDir </> sharedImagesRootDirectory)

-- | Depending on the 'maxLocalSharedImageRevisions' 'B9Config' settings either
-- do nothing or delete all but the configured number of most recent shared
-- images with the given name from the local cache.
cleanOldSharedImageRevisionsFromCache :: SharedImageName -> B9 ()
cleanOldSharedImageRevisionsFromCache sn = do
  b9Cfg <- getConfig
  forM_ (b9Cfg ^. maxLocalSharedImageRevisions) $ \maxRevisions -> do
    toDelete <- take maxRevisions <$> newestSharedImages
    imgDir <- getSharedImagesCacheDir
    let filesToDelete = (imgDir </>) <$> (infoFiles ++ imgFiles)
        infoFiles = sharedImageFileName <$> toDelete
        imgFiles = imageFileName . sharedImageImage <$> toDelete
    unless (null filesToDelete) $ do
      traceL (printf "DELETING %d OBSOLETE REVISIONS OF: %s" (length filesToDelete) (show sn))
      mapM_ traceL filesToDelete
      mapM_ removeIfExists filesToDelete
  where
    newestSharedImages :: B9 [SharedImage]
    newestSharedImages = reverse . map snd <$> lookupSharedImages (== Cache) ((sn ==) . sharedImageName)
    removeIfExists fileName = liftIO $ removeFile fileName `catch` handleExists
      where
        handleExists e
          | isDoesNotExistError e = return ()
          | otherwise = throwIO e