{-| Data types that describe all B9 relevant elements of virtual machine disk
images.-}
module B9.DiskImages where

import           B9.QCUtil
import           GHC.Generics (Generic)
import           Control.Parallel.Strategies
import           Data.Binary
import           Data.Data
import           Data.Hashable
import           Data.Maybe
import           Data.Semigroup as Sem
import           System.FilePath
import           Test.QuickCheck
import qualified Text.PrettyPrint.Boxes as Boxes
import           Text.Printf

-- * Data types for disk image description, e.g. 'ImageTarget',
-- 'ImageDestination', 'Image', 'MountPoint', 'SharedImage'

-- | Build target for disk images; the destination, format and size of the image
-- to generate, as well as how to create or obtain the image before a
-- 'B9.Vm.VmScript' is executed with the image mounted at a 'MountPoint'.
data ImageTarget = ImageTarget
                     ImageDestination
                     ImageSource
                     MountPoint
                     deriving (Read, Show, Typeable, Data, Eq,Generic)

instance Hashable ImageTarget
instance Binary ImageTarget
instance NFData ImageTarget

-- | A mount point or 'NotMounted'
data MountPoint = MountPoint FilePath | NotMounted
                     deriving (Show, Read, Typeable, Data, Eq,Generic)

instance Hashable MountPoint
instance Binary MountPoint
instance NFData MountPoint

-- | The destination of an image.
data ImageDestination = Share String ImageType ImageResize
                      -- ^ Create the image and some meta data so that other
                      -- builds can use them as 'ImageSource's via 'From'.
                      | LiveInstallerImage String FilePath ImageResize
                      -- ^ __DEPRECATED__ Export a raw image that can directly
                      -- be booted.
                      | LocalFile Image ImageResize
                      -- ^ Write an image file to the path in the first
                      -- argument., possible resizing it,
                      | Transient
                      -- ^ Do not export the image. Usefule if the main
                      -- objective of the b9 build is not an image file, but
                      -- rather some artifact produced by executing by a
                      -- containerize build.
                      deriving (Read, Show, Typeable, Data,Eq,Generic)

instance Hashable ImageDestination
instance Binary ImageDestination
instance NFData ImageDestination

-- | Specification of how the image to build is obtained.
data ImageSource = EmptyImage String FileSystem ImageType ImageSize
                  -- ^ Create an empty image file having a file system label
                  -- (first parameter), a file system type (e.g. 'Ext4') and an
                  -- 'ImageSize'
                 | CopyOnWrite Image
                  -- ^ __DEPRECATED__
                 | SourceImage Image Partition ImageResize
                  -- ^ Clone an existing image file; if the image file contains
                  -- partitions, select the partition to use, b9 will extract
                  -- that partition by reading the offset of the partition from
                  -- the partition table and extract it using @dd@.
                 | From String ImageResize
                  -- ^ Use an image previously shared by via 'Share'.
                 deriving (Show,Read,Typeable,Data,Eq,Generic)

instance Hashable ImageSource
instance Binary ImageSource
instance NFData ImageSource

-- | The partition to extract.
data Partition = NoPT -- ^ There is no partition table on the image
               | Partition Int -- ^ Extract partition @n@ @n@ must be in @0..3@
  deriving (Eq, Show, Read, Typeable, Data,Generic)

instance Hashable Partition
instance Binary Partition
instance NFData Partition

-- | A vm disk image file consisting of a path to the image file, and the type
-- and file system.
data Image = Image FilePath ImageType FileSystem
           deriving (Eq, Show, Read, Typeable, Data,Generic)

instance Hashable Image
instance Binary Image
instance NFData Image

-- | An image type defines the actual /file format/ of a file containing file
-- systems. These are like /virtual harddrives/
data ImageType = Raw | QCow2 | Vmdk
               deriving (Eq,Read,Typeable,Data,Show,Generic)

instance Hashable ImageType
instance Binary ImageType
instance NFData ImageType

-- | The file systems that b9 can use and convert.
data FileSystem = NoFileSystem | Ext4 | Ext4_64 | ISO9660 | VFAT
                deriving (Eq,Show,Read,Typeable,Data,Generic)

instance Hashable FileSystem
instance Binary FileSystem
instance NFData FileSystem

-- | A data type for image file or file system size; instead of passing 'Int's
-- around this also captures a size unit so that the 'Int' can be kept small
data ImageSize = ImageSize Int SizeUnit
                 deriving (Eq, Show, Read, Typeable, Data, Generic)

instance Hashable ImageSize
instance Binary ImageSize
instance NFData ImageSize

-- | Enumeration of size multipliers. The exact semantics may vary depending on
-- what external tools look at these. E.g. the size unit is convert to a size
-- parameter of the @qemu-img@ command line tool.
data SizeUnit = B | KB | MB | GB
              deriving (Eq, Show, Read, Ord, Typeable, Data, Generic)

instance Hashable SizeUnit
instance Binary SizeUnit
instance NFData SizeUnit

-- | How to resize an image file.
data ImageResize = ResizeImage ImageSize
                   -- ^ Resize the image __but not the file system__. Note that
                   -- a file system contained in the image file might be
                   -- corrupted by this operation. To not only resize the image
                   -- file but also the fil system contained in it, use
                   -- 'Resize'.
                 | Resize ImageSize
                   -- ^ Resize an image and the contained file system.
                 | ShrinkToMinimum
                   -- ^ Resize an image and the contained file system to the
                   -- smallest size to fit the contents of the file system.
                 | KeepSize
                   -- ^ Do not change the image size.
                   deriving (Eq, Show, Read, Typeable, Data, Generic)

instance Hashable ImageResize
instance Binary ImageResize
instance NFData ImageResize

-- | A type alias that indicates that something of type @a@ is mount at a
-- 'MountPoint'
type Mounted a = (a, MountPoint)

-- * Shared Images

-- | 'SharedImage' holds all data necessary to describe an __instance__ of a shared
--    image identified by a 'SharedImageName'. Shared images are stored in
--    'B9.Repository's.
data SharedImage =
    SharedImage SharedImageName
                SharedImageDate
                SharedImageBuildId
                ImageType
                FileSystem
    deriving (Eq,Read,Show,Typeable,Data,Generic)

instance Hashable SharedImage
instance Binary SharedImage
instance NFData SharedImage

-- | The name of the image is the de-facto identifier for push, pull, 'From' and
--   'Share'.  B9 always selects the newest version the shared image identified
--   by that name when using a shared image as an 'ImageSource'. This is a
--   wrapper around a string that identifies a 'SharedImage'
newtype SharedImageName = SharedImageName String deriving (Eq,Ord,Read,Show,Typeable,Data,Hashable,Binary,NFData)

-- | Get the String representation of a 'SharedImageName'.
fromSharedImageName :: SharedImageName -> String
fromSharedImageName (SharedImageName b) = b

-- | The exact time that build job __started__.
--   This is a wrapper around a string contains the build date of a
--   'SharedImage'; this is purely additional convenience and typesafety
newtype SharedImageDate = SharedImageDate String deriving (Eq,Ord,Read,Show,Typeable,Data,Hashable,Binary,NFData)

-- | Every B9 build running in a 'B9Monad'
--   contains a random unique id that is generated once per build (no matter how
--   many artifacts are created in that build) This field contains the build id
--   of the build that created the shared image instance.  This is A wrapper
--   around a string contains the build id of a 'SharedImage'; this is purely
--   additional convenience and typesafety
newtype SharedImageBuildId = SharedImageBuildId String deriving (Eq,Ord,Read,Show,Typeable,Data,Hashable,Binary,NFData)

-- | Get the String representation of a 'SharedImageBuildId'.
fromSharedImageBuildId :: SharedImageBuildId -> String
fromSharedImageBuildId (SharedImageBuildId b) = b

-- | Shared images are orderd by name, build date and build id
instance Ord SharedImage where
  compare (SharedImage n d b _ _) (SharedImage n' d' b' _ _) =
    compare n n' Sem.<> compare d d' Sem.<> compare b b'

-- * Constroctor and accessors for 'Image' 'ImageTarget' 'ImageSource'
-- 'ImageDestination' and 'SharedImage'

-- | Return the name of the file corresponding to an 'Image'
imageFileName :: Image -> FilePath
imageFileName (Image f _ _) = f

-- | Return the 'ImageType' of an 'Image'
imageImageType :: Image -> ImageType
imageImageType (Image _ t _) = t

-- | Return the files generated for a 'LocalFile' or a 'LiveInstallerImage'; 'SharedImage' and 'Transient'
-- are treated like they have no ouput files because the output files are manged
-- by B9.
getImageDestinationOutputFiles :: ImageTarget -> [FilePath]
getImageDestinationOutputFiles (ImageTarget d _ _) =
    case d of
        LiveInstallerImage liName liPath _ ->
            let path = liPath </> "machines" </> liName </> "disks" </> "raw"
            in [path </> "0.raw", path </> "0.size", path </> "VERSION"]
        LocalFile (Image lfPath _ _) _ -> [lfPath]
        _ -> []

-- | Return the name of a shared image, if the 'ImageDestination' is a 'Share'
--   destination
imageDestinationSharedImageName :: ImageDestination -> Maybe SharedImageName
imageDestinationSharedImageName (Share n _ _) = Just (SharedImageName n)
imageDestinationSharedImageName _ = Nothing

-- | Return the name of a shared source image, if the 'ImageSource' is a 'From'
--   source
imageSourceSharedImageName :: ImageSource -> Maybe SharedImageName
imageSourceSharedImageName (From n _) = Just (SharedImageName n)
imageSourceSharedImageName _ = Nothing

-- | Get the 'ImageDestination' of an 'ImageTarget'
itImageDestination :: ImageTarget -> ImageDestination
itImageDestination (ImageTarget d _ _) = d

-- | Get the 'ImageSource' of an 'ImageTarget'
itImageSource :: ImageTarget -> ImageSource
itImageSource (ImageTarget _ s _) = s

-- | Get the 'MountPoint' of an 'ImageTarget'
itImageMountPoint :: ImageTarget -> MountPoint
itImageMountPoint (ImageTarget _ _ m) = m


-- | Return true if a 'Partition' parameter is actually refering to a partition,
-- false if it is 'NoPT'
isPartitioned :: Partition -> Bool
isPartitioned p
  | p == NoPT = False
  | otherwise = True

-- | Return the 'Partition' index or throw a runtime error if aplied to 'NoPT'
getPartition :: Partition -> Int
getPartition (Partition p) = p
getPartition NoPT = error "No partitions!"

-- | Return the file name extension of an image file with a specific image
-- format.
imageFileExtension :: ImageType -> String
imageFileExtension Raw = "raw"
imageFileExtension QCow2 = "qcow2"
imageFileExtension Vmdk = "vmdk"

-- | Change the image file format and also rename the image file name to
-- have the appropriate file name extension. See 'imageFileExtension' and
-- 'replaceExtension'
changeImageFormat :: ImageType -> Image -> Image
changeImageFormat fmt' (Image img _ fs) = Image img' fmt' fs
  where img' = replaceExtension img (imageFileExtension fmt')

changeImageDirectory :: FilePath -> Image -> Image
changeImageDirectory dir (Image img fmt fs) = Image img' fmt fs
  where img' = dir </> takeFileName img

-- * Constructors and accessors for 'ImageSource's
getImageSourceImageType :: ImageSource -> Maybe ImageType
getImageSourceImageType (EmptyImage _ _ t _) = Just t
getImageSourceImageType (CopyOnWrite i) = Just $ imageImageType i
getImageSourceImageType (SourceImage i _ _) = Just $ imageImageType i
getImageSourceImageType (From _ _) = Nothing

-- * Constructors and accessors for 'SharedImage's

-- | Return the name of a shared image.
sharedImageName :: SharedImage -> SharedImageName
sharedImageName (SharedImage n _ _ _ _) = n

-- | Return the build date of a shared image.
sharedImageDate :: SharedImage -> SharedImageDate
sharedImageDate (SharedImage _ n _ _ _) = n

-- | Return the build id of a shared image.
sharedImageBuildId :: SharedImage -> SharedImageBuildId
sharedImageBuildId (SharedImage _ _ n _ _) = n

-- | Print the contents of the shared image in one line
prettyPrintSharedImages :: [SharedImage] -> String
prettyPrintSharedImages imgs = Boxes.render table
  where
    table = Boxes.hsep 1 Boxes.left cols
      where
        cols = [nameC, dateC, idC]
          where
            nameC = col "Name" ((\(SharedImageName n) -> n) . sharedImageName)
            dateC = col "Date" ((\(SharedImageDate n) -> n) . sharedImageDate)
            idC = col "ID" ((\(SharedImageBuildId n) -> n) . sharedImageBuildId)
            col title accessor =
              (Boxes.text title) Boxes.// (Boxes.vcat Boxes.left cells)
              where
                cells = Boxes.text <$> accessor <$> imgs

-- | Return the disk image of an sharedImage
sharedImageImage :: SharedImage -> Image
sharedImageImage (SharedImage (SharedImageName n) _ (SharedImageBuildId bid) sharedImageType sharedImageFileSystem) =
    Image
        (n ++ "_" ++ bid <.> imageFileExtension sharedImageType)
        sharedImageType
        sharedImageFileSystem

-- | Calculate the path to the text file holding the serialized 'SharedImage'
-- relative to the directory of shared images in a repository.
sharedImageFileName :: SharedImage -> FilePath
sharedImageFileName (SharedImage (SharedImageName n) _ (SharedImageBuildId bid) _ _) =
    n ++ "_" ++ bid <.> sharedImageFileExtension

sharedImagesRootDirectory :: FilePath
sharedImagesRootDirectory = "b9_shared_images"

sharedImageFileExtension :: String
sharedImageFileExtension  = "b9si"

-- | The internal image type to use as best guess when dealing with a 'From'
-- value.
sharedImageDefaultImageType :: ImageType
sharedImageDefaultImageType = QCow2

-- * Constructors for 'ImageTarget's

-- | Use a 'QCow2' image with an 'Ext4' file system
transientCOWImage :: FilePath -> FilePath -> ImageTarget
transientCOWImage fileName mountPoint =
    ImageTarget
        Transient
        (CopyOnWrite (Image fileName QCow2 Ext4))
        (MountPoint mountPoint)

-- | Use a shared image
transientSharedImage :: SharedImageName -> FilePath -> ImageTarget
transientSharedImage (SharedImageName name) mountPoint =
    ImageTarget Transient (From name KeepSize) (MountPoint mountPoint)

-- | Use a shared image
transientLocalImage :: FilePath -> FilePath -> ImageTarget
transientLocalImage name mountPoint =
    ImageTarget Transient (From name KeepSize) (MountPoint mountPoint)

-- | Share a 'QCow2' image with 'Ext4' fs
shareCOWImage :: FilePath -> SharedImageName -> FilePath -> ImageTarget
shareCOWImage srcFilename (SharedImageName destName) mountPoint =
    ImageTarget
        (Share destName QCow2 KeepSize)
        (CopyOnWrite (Image srcFilename QCow2 Ext4))
        (MountPoint mountPoint)

-- | Share an image based on a shared image
shareSharedImage :: SharedImageName
                 -> SharedImageName
                 -> FilePath
                 -> ImageTarget
shareSharedImage (SharedImageName srcName) (SharedImageName destName) mountPoint =
    ImageTarget
        (Share destName QCow2 KeepSize)
        (From srcName KeepSize)
        (MountPoint mountPoint)

-- | Share a 'QCow2' image with 'Ext4' fs
shareLocalImage :: FilePath -> SharedImageName -> FilePath -> ImageTarget
shareLocalImage srcName (SharedImageName destName) mountPoint =
    ImageTarget
        (Share destName QCow2 KeepSize)
        (SourceImage (Image srcName QCow2 Ext4) NoPT KeepSize)
        (MountPoint mountPoint)

-- | Export a 'QCow2' image with 'Ext4' fs
cowToliveInstallerImage :: String
                        -> FilePath
                        -> FilePath
                        -> FilePath
                        -> ImageTarget
cowToliveInstallerImage srcName destName outDir mountPoint =
    ImageTarget
        (LiveInstallerImage destName outDir KeepSize)
        (CopyOnWrite (Image srcName QCow2 Ext4))
        (MountPoint mountPoint)

-- | Export a 'QCow2' image file with 'Ext4' fs as
--   a local file
cowToLocalImage :: FilePath -> FilePath -> FilePath -> ImageTarget
cowToLocalImage srcName destName mountPoint =
    ImageTarget
        (LocalFile (Image destName QCow2 Ext4) KeepSize)
        (CopyOnWrite (Image srcName QCow2 Ext4))
        (MountPoint mountPoint)

-- | Export a 'QCow2' image file with 'Ext4' fs as
--   a local file
localToLocalImage :: FilePath -> FilePath -> FilePath -> ImageTarget
localToLocalImage srcName destName mountPoint =
    ImageTarget
        (LocalFile (Image destName QCow2 Ext4) KeepSize)
        (SourceImage (Image srcName QCow2 Ext4) NoPT KeepSize)
        (MountPoint mountPoint)

-- | Create a local image file from the contents of the first partition
--   of a local 'QCow2' image.
partition1ToLocalImage :: FilePath -> FilePath -> FilePath -> ImageTarget
partition1ToLocalImage srcName destName mountPoint =
    ImageTarget
        (LocalFile (Image destName QCow2 Ext4) KeepSize)
        (SourceImage (Image srcName QCow2 Ext4) NoPT KeepSize)
        (MountPoint mountPoint)

-- * 'ImageTarget' Transformations

-- | Split any image target into two image targets, one for creating an intermediate shared image and one
-- from the intermediate shared image to the output image.
splitToIntermediateSharedImage :: ImageTarget
                               -> SharedImageName
                               -> (ImageTarget, ImageTarget)
splitToIntermediateSharedImage (ImageTarget dst src mnt) (SharedImageName intermediateName) =
    (imgTargetShared, imgTargetExport)
  where
    imgTargetShared = ImageTarget intermediateTo src mnt
    imgTargetExport = ImageTarget dst intermediateFrom mnt
    intermediateTo =
        Share
            intermediateName
            (fromMaybe
                 sharedImageDefaultImageType
                 (getImageSourceImageType src))
            KeepSize
    intermediateFrom = From intermediateName KeepSize

-- * 'Arbitrary' instances for quickcheck

instance Arbitrary ImageTarget where
    arbitrary =
        ImageTarget <$> smaller arbitrary <*> smaller arbitrary <*>
        smaller arbitrary

instance Arbitrary ImageSource where
    arbitrary =
        oneof
            [ EmptyImage "img-label" <$> smaller arbitrary <*>
              smaller arbitrary <*>
              smaller arbitrary
            , CopyOnWrite <$> smaller arbitrary
            , SourceImage <$> smaller arbitrary <*> smaller arbitrary <*>
              smaller arbitrary
            , From <$> arbitrarySharedImageName <*> smaller arbitrary]

instance Arbitrary ImageDestination where
    arbitrary =
        oneof
            [ Share <$> arbitrarySharedImageName <*> smaller arbitrary <*>
              smaller arbitrary
            , LiveInstallerImage "live-installer" "output-path" <$>
              smaller arbitrary
            , pure Transient]

instance Arbitrary MountPoint where
    arbitrary = elements [MountPoint "/mnt", NotMounted]

instance Arbitrary ImageResize where
    arbitrary =
        oneof
            [ ResizeImage <$> smaller arbitrary
            , Resize <$> smaller arbitrary
            , pure ShrinkToMinimum
            , pure KeepSize]

instance Arbitrary Partition where
    arbitrary = oneof [Partition <$> elements [0, 1, 2], pure NoPT]

instance Arbitrary Image where
    arbitrary =
        Image "img-file-name" <$> smaller arbitrary <*> smaller arbitrary

instance Arbitrary FileSystem where
    arbitrary = elements [Ext4]

instance Arbitrary ImageType where
    arbitrary = elements [Raw, QCow2, Vmdk]

instance Arbitrary ImageSize where
    arbitrary = ImageSize <$> smaller arbitrary <*> smaller arbitrary

instance Arbitrary SizeUnit where
    arbitrary = elements [B, KB, MB, GB]

instance Arbitrary SharedImageName where
    arbitrary = SharedImageName <$> arbitrarySharedImageName

arbitrarySharedImageName :: Gen String
arbitrarySharedImageName =
    elements [printf "arbitrary-shared-img-name-%d" x | x <- [0 :: Int .. 3]]