{-# OPTIONS_GHC -fno-warn-type-defaults #-}
{-# LANGUAGE ExtendedDefaultRules #-}
{-# LANGUAGE NamedFieldPuns       #-}
{-# LANGUAGE NoImplicitPrelude    #-}
{-# LANGUAGE OverloadedStrings    #-}
{-# LANGUAGE RecordWildCards      #-}
{-# LANGUAGE ScopedTypeVariables  #-}

module Main where
import Rebase.Prelude hiding (bool, hash, (<.>))

import Crypto.Hash
import Crypto.Hash.Types.Digest.Read ()
import Data.Yaml.Combinators
import Development.Shake
import Development.Shake.Command     ()
import Development.Shake.FilePath
import qualified Rebase.Data.Set     as Set
import qualified Rebase.Data.Text    as Text
         ;import Rebase.Data.Text    (pack, unpack)
import Rebase.Data.Text.Encoding
import qualified Rebase.Data.Vector  as Vector
{-

Basic idea is to have a fs.yaml that specifies the build procedure. Analogous to
stack.yaml, it should specify everything (every source) directly or indirectly.

We want to make new images from CoW copies of old ones. We want to build these
things incrementally, but still end up with something that will be reproducible
from scratch.

Anyway, what we'll have is a list of packages, which will be unpacked first.
Then a list of debconf values, which will be applied. Then we will have the rest
of the other, slower changes to the image (including dpkg --configure -a). Some
changes can be assumed to produce the same results out of order. Oh right, and
the zeroeth step is to generate the empty (or just initial) filesystem image.
The initial filesystem image, if nonempty, is specified by filename (otherwise,
by size), while the hash of the configuration determines the output filename.

--

'Parent' param should be a config file, not hash. The hash will only be used to
determine whether we need to build it.

-}

newtype DebconfConfig = DebconfConfig Text deriving (Show, Read, Eq, Ord)
newtype Package       = Package Text deriving (Show, Read, Eq, Ord)
data Patch            = Patch deriving (Show, Read)

data BaseImageSpecification
  = EmptyImageOfBytes Int
  | ParentImageConfigFile FilePath
  deriving (Show, Read)

data DiskImageConfig = DiskImageConfig {
  initialImage :: BaseImageSpecification
, packages     :: Set Package
-- , debconfConfig :: DebconfConfig
, unpackOnly   :: Bool
, binaries     :: Vector Text -- :: [Patch]
} deriving (Show, Read)

parsePackageName :: Text -> Package
parsePackageName = Package -- TODO

diskImageConfigParser :: Parser DiskImageConfig
diskImageConfigParser = object $
  DiskImageConfig
  <$> field "parent" ((EmptyImageOfBytes <$> integer) <> (ParentImageConfigFile . unpack <$> string))
  <*> (Set.fromList . toList . fmap parsePackageName <$> defaultField "packages" Vector.empty (array string))
  -- <*> pure (DebconfConfig "")
  <*> defaultField "unpack-only" False bool
  <*> defaultField "binaries" Vector.empty (array string)

diskImageFilename :: DiskImageConfig -> FilePath
diskImageFilename = (++ ".btrfs") . show . sha1 . show
  where
    sha1 :: String -> Digest SHA1
    sha1 = hash . encodeUtf8 . pack

readCfg :: FilePath -> Action DiskImageConfig
readCfg yaml = either error id . parse diskImageConfigParser . encodeUtf8 . pack <$> readFile' yaml

buildRoot :: DiskImageConfig -> FilePath -> Action ()
buildRoot DiskImageConfig{..} finalOut = do
  let out        = finalOut <.> "tmp"
      mountpoint = finalOut <.> "mnt"
  cmd_ "sh -c" ["! mountpoint -q \"$0\" || umount \"$0\" ", mountpoint]
  case initialImage of
    ParentImageConfigFile f -> do
      let parent = "_build" </> f -<.> "btrfs"
      need [parent]
      cmd_ "cp --reflink" [parent, out]
      cmd_ "mkdir -p" [mountpoint]
      cmd_ "mount -t btrfs" [out] mountpoint
    EmptyImageOfBytes n -> do
      cmd_ "truncate -s" [show n] [out]
      cmd_ "mkfs.btrfs" [out]
      cmd_ "mkdir -p" [mountpoint]
      cmd_ "mount -t btrfs" [out] mountpoint

      -- create new default subvolume, and then remount with it
      createDefaultSubvolume mountpoint
      cmd_ "umount" [mountpoint]
      cmd_ "mount -t btrfs" [out, mountpoint]

      cmd_ (Cwd mountpoint) "mkdir -p var/cache/apt"
      cmd_ (Cwd mountpoint) "btrfs subvolume create var/cache/apt/archives"

  -- TODO: catch errors and umount, rmdir mountpoint
  let (debs, packageNames) = partitionPackages $ unpack . coerce <$> toList packages
  cmd_ "selfstrap --skip-update"
    (("--unpack" `consWhen` unpackOnly) ["-t", mountpoint]) packageNames
  forM_ debs $ \pkg -> do
    cmd_ "dpkg -i --root" [mountpoint, pkg]
  cmd_ "mv" [out, finalOut]

partitionPackages :: [String] -> ([String], [String])
partitionPackages = partition (elem '/')

consWhen :: a -> Bool -> [a] -> [a]
a `consWhen` c = if c then (a:) else id

strip :: String -> String
strip = unpack . Text.strip . pack

createDefaultSubvolume :: FilePath -> Action ()
createDefaultSubvolume mountpoint = do
  cmd_ (Cwd mountpoint) "btrfs subvolume create root"
  Stdout subvolIdLine <- cmd (Cwd mountpoint) "sh -c"
    ["btrfs subvolume show root | sed -n -e 's/^[ \t]*Subvolume ID:[ \t]*//p; s/.*is toplevel subvolume/5/p'"]
  let subvolId = strip subvolIdLine
  when (null subvolId) $ fail "could not obtain btrfs Subvolume ID"
  cmd_ "btrfs subvolume set-default" [strip subvolId, mountpoint]

defaultImageName :: String
defaultImageName = "minbase"

stripSuffix :: Text -> String -> String
stripSuffix s t = fromMaybe t $ fmap unpack $ Text.stripSuffix s $ pack t

main :: IO ()
main = do
  -- TODO: Fail early on:
  -- 1. not running as root
  -- 2. no "selfstrap" in PATH
  args <- getArgs
  case args of
    [target] -> shakeBuildOneImage (stripSuffix ".yaml" target)
    []       -> shakeBuildOneImage defaultImageName
    _        -> error "usage"

shakeBuildOneImage :: FilePath -> IO ()
shakeBuildOneImage target =
  shake shakeOptions {shakeFiles = "_build"} $ do
    want [target <.> "btrfs"]
    shakeRules

shakeRules :: Rules ()
shakeRules = do
  "_build/*.yaml.canon" %> \out -> do
    let yaml = dropDirectory1 $ dropExtension out
    need [yaml]
    cfg <- readCfg yaml
    writeFileChanged out (show cfg)
  "_build/*.btrfs" %> \out -> do
    let cfgFile = (out -<.> "yaml.canon")
    need [cfgFile]
    cfg <- read <$> readFile' cfgFile
    buildRoot cfg out

  "*.btrfs" %> \out -> do
    orderOnly ["_build" </> out]
    cmd_ "cp --reflink=always -i" ["_build" </> out, out]