{-# LANGUAGE ViewPatterns          #-}
{-# OPTIONS_GHC -fno-warn-type-defaults #-}
{-# LANGUAGE DuplicateRecordFields #-}
{-# 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 qualified Options.Applicative as Opt
         ;import Options.Applicative hiding (action, command)

import ConfigFile
import System.Directory     (getCurrentDirectory, createDirectoryIfMissing)
import System.Posix.Process (getProcessID)
import System.Posix.Types   (CUid (..))
import System.Posix.User    (getEffectiveUserID)

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

buildRoot :: DiskImageConfig -> FilePath -> Action ()
buildRoot config@DiskImageConfig{..} finalOut = do
  let out        = finalOut <.> "tmp"
      mountpoint = finalOut <.> "mnt"
  cmd_ "sh -xc" ["! mountpoint \"$0\" || umount \"$0\" ", mountpoint]

  let (abortion :: IO ()) = ignoreErrors' $ do
        cmd_ "umount" [mountpoint]
        cmd_ "rmdir" [mountpoint]
        cmd_ "rm -f" [out]

  handle' abortion $ buildInitialImage config mountpoint out

  handle' abortion $ do
    {- 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
    when (not $ null debs) $ do
      cmd_ (AddEnv "DEBIAN_FRONTEND" "noninteractive")
           ["dpkg"] [if unpackOnly then "--unpack" else "--install"] debs
    {- 2.5. install apt package cache -}
    -- TODO
    {- 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]
    {- 3.5 skel -}
    forM_ (unpack <$> skelFiles) $ \f -> do
      homeDir <- getHomeDir
      target <- absPath mountpoint <&> (</> "etc/skel")
      cmd_ "mkdir -p" [target]
      cmd_ (Cwd homeDir) "cp -r --preserve=mode,timestamps -L --parents -t" [target] [f]
    {- 4. custom setup commands -}
    forM_ chrootCommands $ \c -> do
      cmd_ "chroot" [mountpoint] "/bin/sh -c" [unpack c]
    {- 5. create a backup snapshot -}
    cmd_ "umount" [mountpoint]
    cmd_ "mount -t btrfs -o subvol=/" [out, mountpoint]
    ignoreErrors $ cmd_ "btrfs subvolume delete" [mountpoint </> "root~orig"]
    cmd_ "btrfs subvolume snapshot -r" [mountpoint </> "root", mountpoint </> "root~orig"]

    -- cleanup
    cmd_ "umount" [mountpoint]
    cmd_ "rmdir" [mountpoint]
    cmd_ "btrfstune -S1" [out]
    cmd_ "mv" [out, finalOut]

absPath :: MonadIO m => FilePath -> m FilePath
absPath f@('/':_) = return f
absPath f = do
  wd <- liftIO getCurrentDirectory
  return $ wd </> f

getHomeDir :: Action FilePath
getHomeDir = do
  Stdout homeDir <- cmd ["sh", "-c", "getent passwd \"$SUDO_USER\"|(IFS=: read _ _ _ _ _ d _; printf %s \"$d\")"]
  return homeDir

buildInitialImage :: DiskImageConfig -> FilePath -> FilePath -> Action ()
buildInitialImage DiskImageConfig{..} mountpoint out = do
  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"

ignoreErrors' :: IO () -> IO ()
ignoreErrors' = flip catch (\(SomeException _) -> return ())

ignoreErrors :: Action () -> Action ()
ignoreErrors = flip actionCatch (\(SomeException _) -> return ())

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]

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 | Chroot ChrootOpts
data BuildOpts = BuildOpts { optTarget :: String }
data ChrootOpts = ChrootOpts { optTarget :: String, optChrootCommand :: [String] }

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

chrootOpts :: Parser Command
chrootOpts = fmap Chroot $ ChrootOpts <$> argument str idm <*> many (argument str idm)

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)
                             <> Opt.command "chroot" (info' chrootOpts chrootDesc)
    buildDesc  = ["build the image specified by the YAML config file"]
    chrootDesc = ["chroot into the mounted image.  ",
                  "Iff the chroot exits with a success value, filesystem changes will persist"]

    info' o d = info (helper <*> o) (progDesc $ unwords d)

run :: Options -> IO ()
run (Options (Build (BuildOpts target)))        = earlyFail >> shakeBuildOneImage target
run (Options (Chroot (ChrootOpts target args))) = earlyFail >> chrootImage target args

useCGroups :: Bool
useCGroups = True               -- TODO: make command-line option

handle' :: IO b -> Action a -> Action a
handle' = flip actionOnException

chrootImage :: FilePath -> [String] -> IO ()
chrootImage target args =
  shake shakeOptions {shakeFiles = "_build"} $ do
    shakeRules
    action $ do
      pid <- show <$> liftIO getProcessID
      let inp = target -<.> "btrfs"
      let tmp = inp <.> "tmp" <.> pid
      let mnt = tmp <.> "mnt"
      orderOnly [inp]
      cmd_ "cp --reflink=always" [inp, tmp]
      cmd_ "btrfstune -S0 -f" tmp
      cmd_ "mkdir" [mnt]
      cmd_ "mount -t btrfs" [tmp, mnt]
      let (umount :: IO ()) = do
            cmd_ "umount" [mnt]
            cmd_ "rm -f" [tmp]
      handle' umount $
        if useCGroups
          then liftIO $ cgroupChroot ("fsmgr" <.> takeFileName target) mnt args
          else cmd_ (WithStderr False) "chroot" (mnt : args) -- TODO: cgroup
      cmd_ "umount" [mnt]
      cmd_ "sync"
      cmd_ "btrfstune -S1" [tmp]
      cmd_ "mv" [tmp, inp]

shakeBuildOneImage :: FilePath -> IO ()
shakeBuildOneImage target =
  shake shakeOptions {shakeFiles = "_build", shakeVerbosity = Loud} $ 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

infix 1 ~%>
(~%>) :: FilePattern -> (FilePath -> Action ()) -> Rules ()
filepat ~%> act = phonys f
  where
    f file | (filepat ?== file) = Just (act file)
    f _    = Nothing

cgroupChroot :: String -> FilePath -> [String] -> IO ()
cgroupChroot groupName mnt [] = cgroupChroot groupName mnt ["/bin/bash"]
cgroupChroot groupName mnt args = do
  let cgdir = "/sys/fs/cgroup/pids" </> groupName
  createDirectoryIfMissing False cgdir
  cmd_ (Cwd mnt) (WithStderr False)
    "unshare --ipc --uts --cgroup --mount --pid --fork chroot ."
    "sh -exc" ["mount -t proc proc /proc; mount -t devpts devpts /dev/pts; exec \"$@\""]
    "sh" args

earlyFail :: IO ()
earlyFail = do
  CUid euid <- liftIO getEffectiveUserID
  when (euid /= 0) $ fail "you are not root"
  Stdout () <- cmd (Traced []) "which selfstrap"
  return ()

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]