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

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

import qualified Rebase.Data.Text as Text
         ;import Rebase.Data.Text (pack, unpack)

import Development.Shake          hiding (getEnv)
import Development.Shake.Command  ()
import Development.Shake.FilePath

import Options.Applicative as Opt

import ConfigFile

noParent :: BaseImageSpecification -> Bool
noParent (EmptyImageOfBytes _)     = True
noParent (ParentImageConfigFile _) = False

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_ "btrfstune -f -S0" [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) "mkdir -p var/cache/debconf"
      cmd_ (Cwd mountpoint) "btrfs subvolume create var/cache/apt/archives"
  -- TODO: catch errors and umount, rmdir mountpoint
  {- 1. debconf -}
  forM_ debconfConfig $
    readFile' >=> liftIO . appendFile (mountpoint </> "var/cache/debconf/config.dat")
  {- 2. dpkg installs -}
  let (debs, packageNames) = partitionPackages $ unpack . coerce <$> toList packages
  when (noParent initialImage || not (null packageNames)) $
    -- When there is no parent, selfstrap should install packages marked
    -- "Required" even if no packages are specified. When there is a parent,
    -- assume that this has already happened.
    cmd_ "selfstrap --skip-update" (("--unpack" `consWhen` unpackOnly) ["-t", mountpoint]) packageNames
  forM_ debs $ \pkg -> do cmd_ "dpkg -i --root" [mountpoint, pkg]
  {- 3. binaries -}
  forM_ (unpack <$> binaries) $ \b -> do
    p <- fromMaybe (fail $ "not found in ${PATH}: " ++ b) <$> liftIO (pathLocate b)
    cmd_ "cp -L" [p] [mountpoint </> "usr/local/bin" </> b]
  {- 4. custom setup commands -}
  forM_ chrootCommands $ \c -> do
    cmd_ "chroot" [mountpoint] "/bin/sh -c" [unpack c]
  -- cleanup
  cmd_ "umount" [mountpoint]
  cmd_ "rmdir" [mountpoint]
  cmd_ "btrfstune -S1" [out]
  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

data Options = Options { optCommand :: Command }
data Command = Build BuildOpts
data BuildOpts = BuildOpts { optTarget :: String }

buildOpts :: Parser Command
buildOpts = Build . BuildOpts <$> argument str idm

-- TODO: Fail early on:
-- 1. not running as root
-- 2. no "selfstrap" in PATH
main :: IO ()
main = customExecParser (prefs showHelpOnEmpty) (info opts desc) >>= run
  where
    opts :: Parser Options
    opts = Options <$> parseCommand
    desc = fullDesc <> progDesc detailed <> header "fsmgr - Debian rootfs image manager"
    detailed = unwords ["This program generates btrfs filesystem images using 'selfstrap'",
                        "which, much like 'debootstrap', creates new installations of Debian."]
    parseCommand :: Parser Command
    parseCommand = subparser $ Opt.command "build" (info' buildOpts buildDesc)
    buildDesc = ["build the image specified by the YAML config file"]
    info' o d = info (helper <*> o) (progDesc $ unwords d)

run :: Options -> IO ()
run (Options (Build (BuildOpts target)))= shakeBuildOneImage (stripSuffix ".yaml" target)

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

head1 :: String -> String
head1 = lines >>> \case
  []  -> ""
  x:_ -> x

pathLocate :: String -> IO (Maybe FilePath)
pathLocate c = (getEnv "SUDO_USER" >>=) $ fmap (validatePath . head1 . fromStdout) <$> \case
  "" -> cmd "which" [c]
  u  -> do
    path <- fromStdout <$> cmd "su -" [u] "-c" ["printf %s \"$PATH\""]
    cmd (AddEnv "PATH" path) "which" [c]
  where
    validatePath "" = Nothing
    validatePath x  = Just x

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
    need ["_build" </> out]
    -- WithStderr False needed for `cp` to interact with the tty
    cmd_ (WithStderr False) "cp --reflink=always -i" ["_build" </> out, out]