{-# 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
                           ,pullRemoteRepos
                           ,pullLatestImage
                           ,) where

import Control.Exception
import Control.Monad
import Control.Monad.IO.Class
import System.Directory
import System.FilePath
import Text.Printf (printf)
import Data.Maybe
import Data.Monoid
import Data.Function
import Control.Applicative
import Text.Show.Pretty (ppShow)
import Data.List
import Data.Data
import Data.Generics.Schemes
import Data.Generics.Aliases

import B9.B9Monad
import B9.Repository
import B9.RepositoryIO
import B9.DiskImages
import qualified B9.PartitionTable as P
import B9.ConfigUtils
import B9.Content.StringTemplate


-- | Replace $... variables inside an 'ImageTarget'
substImageTarget :: [(String,String)] -> ImageTarget -> ImageTarget
substImageTarget env p = everywhere gsubst p
  where gsubst :: forall a. 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'. Note however that this source will
-- may not exist as is the case for 'EmptyImage'.
resolveImageSource :: ImageSource -> B9 Image
resolveImageSource src = do
  case src of
   (EmptyImage fsLabel fsType imgType _size) ->
     let img = Image fsLabel imgType fsType
     in return (changeImageFormat imgType img)
   (SourceImage srcImg _part _resize) ->
     liftIO (ensureAbsoluteImageDirExists srcImg)
   (CopyOnWrite backingImg) ->
     liftIO (ensureAbsoluteImageDirExists backingImg)
   (From name _resize) -> do
     latestImage <- getLatestImageByName name
     liftIO (ensureAbsoluteImageDirExists latestImage)

-- | 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) -> do
     sharedImg <- getLatestImageByName name
     preferredDestImageTypes (SourceImage sharedImg NoPT resize)

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]

ensureAbsoluteImageDirExists :: Image -> IO Image
ensureAbsoluteImageDirExists img@(Image path _ _) = do
  let dir = takeDirectory path
  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 filesyste. 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) -> do
     sharedImg <- getLatestImageByName name
     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))

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
  dbgL (printf "Creating empty raw image '%s' with size %s" imgFile
                (toQemuSizeOptVal imgSize))
  cmd (printf "qemu-img create -f %s '%s' '%s'"
              (imageFileExtension imgFmt)
              imgFile
              (toQemuSizeOptVal imgSize))
  case (imgFmt, imgFs) of
    (Raw, Ext4) -> do
      let fsCmd = "mkfs.ext4"
      dbgL (printf "Creating file system %s" (show imgFs))
      cmd (printf "%s -L '%s' -q '%s'" fsCmd fsLabel imgFile)
    (it, fs) -> do
      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 Ext4) = 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 Ext4) = 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 = convert False

-- | 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 = 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'"
                (imageFileExtension fmtIn) (imageFileExtension fmtOut)
                imgIn imgOut)
    when doMove $ do
      dbgL (printf "Removing '%s'" imgIn)
      liftIO (removeFile imgIn)

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.
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))
  tell (dir </> sharedImageFileName sharedImg) sharedImg
  dbgL (printf "CREATED SHARED IMAGE IN CAHCE '%s'" (ppShow sharedImg))
  return sharedImg


-- | Publish the latest version of a shared image identified by name to the
-- selected repository from the cache.
pushSharedImageLatestVersion :: SharedImageName -> B9 ()
pushSharedImageLatestVersion (SharedImageName imgName) = do
  sharedImage <- getLatestSharedImageByNameFromCache imgName
  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 Bool
pullLatestImage (SharedImageName name) = do
  repos <- getSelectedRepos
  let repoPredicate Cache = False
      repoPredicate (Remote repoId) = repoId `elem` repoIds
      repoIds = map remoteRepoRepoId repos
      hasName sharedImage = name == siName sharedImage
  candidates <- lookupSharedImages repoPredicate hasName
  let (Remote repoId, image) = head (reverse candidates)
  if null candidates
     then do errorL (printf "No shared image named '%s'\
                            \ on these remote repositories: '%s'"
                            name
                            (ppShow repoIds))
             return False
     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'"
                           name
                           repoId)
             return True


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

-- | Return the latest version of a shared image named 'name' from the local cache.
getLatestSharedImageByNameFromCache :: String -> B9 SharedImage
getLatestSharedImageByNameFromCache name = do
  imgs <- lookupSharedImages (== Cache) ((== name) . siName)
  case reverse imgs of
    (Cache, sharedImage):_rest ->
      return sharedImage
    _ ->
      error (printf "No image(s) named '%s' found." name)

-- | 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     -- user has not selected a repo
                    return       -- user has selected a repo, return it as
                                 -- singleton list
                    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)