{-# language RecordWildCards #-}
{-# language ApplicativeDo #-}
{-# language LambdaCase #-}
{-# language OverloadedStrings #-}

{-|
Module      : Main
Description : argonix
Copyright   : (c) Valentin Reis, 2018
License     : MIT
Maintainer  : fre@freux.fr -}

module Main where

import           Data.Default
import           Turtle                  hiding ( header
                                                , text
                                                , arguments
                                                )
import           Options.Applicative
import           Data.Maybe
import           Prelude                 hiding ( FilePath )

import           Data.Text                      ( pack
                                                , unpack
                                                )
import           System.Console.ANSI
import           System.Console.ANSI.Types      ( Color )
import           Turtle.Shell
import           Control.Foldl
import           System.Posix.Process
import           System.IO               hiding ( FilePath )

-- | Miscellaneous printing utilities
colorShell :: Color -> Shell () -> Shell ()
colorShell color she = setC color *> she *> setC White
  where setC c = liftIO $ setSGR [SetColor Foreground Dull c]
verboseShell :: Bool -> Text -> Shell Line -> Shell ExitCode
verboseShell v c i = when v (printCommand c) >> shell c i
printInfo :: Text -> Shell ()
printCommand :: Text -> Shell ()
printError :: Text -> Shell ()
printWarning :: Text -> Shell ()
printSuccess :: Text -> Shell ()
printTest :: Text -> Shell ()
dieRed :: Text -> Shell ()

printInfo = printf ("Info: " % s % "\n")
printCommand = printf ("Running: " % s % "\n")
printWarning = colorShell Yellow . printf ("Warning: " % s % "\n")
printError = colorShell Red . printf ("Error: " % s % "\n")
printSuccess = colorShell Green . printf ("Success: " % s % "\n")
printTest = colorShell Green . printf ("RUNNING TEST: " % s % "\n")
dieRed strw =
  colorShell Red (printf ("Failure: " % s) strw) >> exit (ExitFailure 1)

checkFsAttributes :: FilePath -> Shell ()
checkFsAttributes workingDirectory = do
  let dir = case toText workingDirectory of
        Left  di -> di
        Right di -> di
  let findmnt = inproc "findmnt" ["-T", dir, "-o", "OPTIONS"] empty
  b <- liftIO
    $ Turtle.Shell.fold (grep (has "nosuid") findmnt) Control.Foldl.length
  when (b > 0) $ die $ format
    ("The output directory, " % fp % ", must not mounted with \"nosuid\"")
    workingDirectory

main :: IO ()
main = do
  hSetBuffering System.IO.stdout NoBuffering
  join $ execParser (info (opts <**> helper) idm)
 where
  opts :: Parser (IO ())
  opts = hsubparser
    (  command
        "build"
        (info (wrap NixBuild <$> targetParser <*> commonParser)
              (progDesc "Run an argo-compatible nix-build.")
        )
    <> command
         "shell"
         (info (wrap NixShell <$> targetParser <*> commonParser)
               (progDesc "Enter an argo-compatible nix-shell")
         )
    <> command
         "remote-shell"
         (info
           (remotely <$> targetParser <*> commonParser <*> remoteParser)
           (progDesc
             "Enter an argo-compatible nix-shell on a remote machine with nix enabled"
           )
         )
    <> help "Type of operation to run."
    )

data ArgsCommon = ArgsCommon
  { argopkgs   :: Text
   , verbosity :: Verbosity
   , run       :: Maybe Text
   , overrides :: [(String, String)]
   , grafting  :: Grafting
   , sandboxing :: Sandboxing
  } deriving (Show)
data Verbosity = Verbose | Normal deriving (Show, Eq)
data Sandboxing = Sandbox | NoSandbox deriving (Show, Eq)
data Grafting = Libnrm | NoGraft deriving (Show, Eq)

data ArgsRemote = ArgsRemote
  {  targetMachine
   , retreive
   , retreiveAs    :: Maybe Text
  } deriving (Show)

instance Default ArgsCommon where
  def = ArgsCommon
    { verbosity  = Normal,
      argopkgs   = "<argopkgs>",
      run        = Nothing,
      overrides  = [],
      grafting   = NoGraft,
      sandboxing = NoSandbox
    }

instance Default ArgsRemote where
  def = ArgsRemote
    { targetMachine = Nothing,
      retreive      = Nothing,
      retreiveAs    = Nothing
    }

targetParser :: Parser String
targetParser =
  strArgument (metavar "TARGET" <> showDefault <> help "The build target.")

remoteParser :: Parser ArgsRemote
remoteParser = do
  retreive <- optional $ strOption
    (long "retreive" <> metavar "RELATIVE_PATH" <> help
      "File/Folder to retreive from the remote machine."
    )
  retreiveAs <- optional $ strOption
    (long "retreiveAs" <> metavar "LOCAL_PATH" <> help
      "File/Folder to save retreived data as (locally)."
    )
  targetMachine <- optional $ strOption
    (long "target_machine" <> metavar "USER@HOST" <> help
      "Target machine. defaults to localhost via forking."
    )
  pure ArgsRemote {..}

commonParser :: Parser ArgsCommon
commonParser = do
  argopkgs <- strOption
    (  long "argopkgs"
    <> metavar "ARGOPKGS"
    <> showDefault
    <> value (argopkgs def)
    <> help "Nix expression that produces the argopkgs source path."
    )
  verbosity <- flag
    Normal
    Verbose
    (long "verbose" <> short 'v' <> help "Enable verbose mode")
  sandboxing <- flag
    NoSandbox
    Sandbox
    (long "sandboxing" <> short 's' <> help "Enable nix sandboxing.")
  grafting <- flag
    NoGraft
    Libnrm
    (long "grafting" <> short 'g' <> help
      "Enable libnrm grafting to avoid rebuilding applications."
    )
  overrides <- catMaybes <$> targets
  run       <- optional $ strOption
    (long "run" <> metavar "COMMAND" <> help
      "Command to run the environment instead of an interactive shell"
    )
  pure ArgsCommon {..}
 where
  targets :: Parser [Maybe (String, String)]
  targets = traverse
    optSrc
    [ "aml"
    , "libnrm"
    , "nrm"
    , "containers"
    , "argotk"
    , "amg"
    , "lammps"
    , "qmcpack"
    , "stream"
    , "openmc"
    , "experiments"
    ]
  optSrc :: String -> Parser (Maybe (String, String))
  optSrc longform = do
    parsed <- optional $ strOption
      (long longform <> metavar "PATH" <> help
        (longform <> " source folder override.")
      )
    pure $ mapT longform parsed
  mapT :: String -> Maybe String -> Maybe (String, String)
  mapT longform thePath = case thePath of
    Nothing -> Nothing
    Just p  -> Just (longform, p)

sudoRemoveFile :: (Text -> Shell ()) -> Text -> FilePath -> Shell ()
sudoRemoveFile printer desc filePath = do
  foundSocket <- testfile filePath
  when foundSocket $ go False
  printInfo $ format ("OK: " % s % " " % fp) desc filePath
 where
  go useSudo = do
    printer $ format ("found file " % s % " at " % fp % ".. ") desc filePath
    shell
        (format ((if useSudo then "sudo " else "") % "rm -rf " % fp) filePath)
        Turtle.empty
      >>= \case
            ExitSuccess -> colorShell Green $ printf " Successfully removed.\n"
            ExitFailure _ -> if useSudo
              then printer $ format
                ("Failed to remove file " % s % ", even with sudo.")
                desc
              else do
                printer $ format
                  ("Failed to remove file " % s % ". Trying sudo..\n")
                  desc
                go True

cleanSocket :: FilePath -> Shell ()
cleanSocket = sudoRemoveFile printWarning "socket"

setupSystem :: ArgsCommon -> Shell ()
setupSystem sa = do
  doVerbose $ printInfo "Setting the nix-build environment up."
  doVerbose $ printInfo "Cleaning sockets."
  Prelude.mapM_
    cleanSocket
    [ "/tmp/nrm-downstream-in"
    , "/tmp/nrm-downstream-event"
    , "/tmp/nrm-upstream-in"
    , "/tmp/nrm-upstream-event"
    ]
  doVerbose $ printInfo "Setting up a cache directory:"
  cachedir <- single
    $ inproc "mktemp" ["-d", "/tmp/deletable-nix-cache-XXXX"] empty
  export "XDG_CACHE_HOME" $ lineToText cachedir
  doVerbose $ printInfo $ lineToText cachedir <> " exported to XDG_CACHE_HOME"
  vshell "sudo rm -rf result" empty >>= \case
    ExitSuccess -> printInfo "removed ./result"
    ExitFailure n ->
      die ("Failed to remove ./result " <> repr n)
  doVerbose $ printInfo "running nix-build for the containers attribute."
  doVerbose $ printCommand $ "nix-build " <> pack
    (unwords (nixArguments "containers" sa))
  nodeos_config <- single
    $ inproc "nix-build" (fmap pack (nixArguments "containers" sa)++["--no-out-link"]) empty
  doVerbose $ printInfo "Checking filesystem attributes on /tmp"
  checkFsAttributes "/tmp"
  vshell "sudo rm -rf /tmp/argo_nodeos_config" empty
  vshell
    (format ("cp " % s % "/bin/argo_nodeos_config /tmp/argo_nodeos_config")
            (lineToText nodeos_config)
    )
    empty
  vshell "sudo chown root:root /tmp/argo_nodeos_config" empty >>= \case
    ExitSuccess -> return ()
    ExitFailure n ->
      die ("Failed to set argo_nodeos_config permissions " <> repr n)
  vshell "sudo chmod u+sw /tmp/argo_nodeos_config" empty >>= \case
    ExitSuccess   -> return ()
    ExitFailure n -> die ("Setting suid bit failed with exit code " <> repr n)
  vshell "sudo /tmp/argo_nodeos_config --clean_config=kill_content:true" empty
  void $ printInfo "Done setting the environment for nix-build up."
 where
  vshell    = verboseShell verbose
  verbose   = verbosity sa == Verbose
  doVerbose = when verbose

nixArguments :: String -> ArgsCommon -> [String]
nixArguments target ArgsCommon {..} =
  [unpack argopkgs, "-A", target]
    ++ concat [ ["--arg", longform <> "-src", p] | (longform, p) <- overrides ]
    ++ (if grafting == Libnrm then ["--arg", "rebuildApps", "true"] else [])
    ++ [ "--show-trace" | verbosity == Verbose ]

data NixCommand = NixBuild | NixShell
toCommand :: IsString p => NixCommand -> p
toCommand NixBuild = "nix-build"
toCommand NixShell = "nix-shell"

wrap :: NixCommand -> String -> ArgsCommon -> IO ()
wrap nixCommand target sa@ArgsCommon {..} = sh $ do
  when (verbosity == Verbose) $ liftIO $ print sa
  _ <- setupSystem sa
  printCommand $ pack nixc <> " " <> pack (unwords arglist)
  liftIO (executeFile nixc True arglist Nothing)
 where
  nixc = toCommand nixCommand
-- Sources of impurity for this build are: "/tmp/ /etc/argo/ /var/run/
-- /var/lock/. Moreover, sandboxing is disabled, in particular because of:
-- /tmp/nrm-* sockets, /etc/argo, /var/run/, /var/lock/ which all need read
-- access.  until these components are patched to allow for alternative paths,
-- no sandbox use is possible. If sandbox paths were read only we could add:
   {-, "--option"-}
   {-, "build-use-sandbox"-}
   {-, "--option"-}
   {-, "extra-sandbox-paths"-}
   {-, "/tmp/ /etc/argo/ /var/run/ /var/lock/"-}
  arglist =
    nixArguments target sa
      ++ [ "--pure"
         , "--allow-new-privileges"
         , "-K"
         , "--option"
         , "build-use-sandbox"
         ]
      ++ [if sandboxing == Sandbox then "true" else "false"]


remotely :: String -> ArgsCommon -> ArgsRemote -> IO ()
remotely _ _ _ = putStrLn "unsupported in this version"

{-nixshell :: IO ()-}
{-nixshell = undefined-}
  {-ArgsCommon {..} <- liftIO $ execParser opts-}

  {--- building nixArguments (pure stuff) and shellArguments (impure stuff)-}
  {-let nixArguments = ["-A", if enableApps then "expe" else "test"] ++ concat-}
        {-[ ["--arg", longform <> "-src", encodeString p]-}
        {-| (longform, p) <- overrides-}
        {-]-}
  {-sudo <- which "sudo" >>= \case-}
    {-(Just sudo) -> printf ("Found sudo at " % fp % "\n") sudo >> return sudo-}
    {-Nothing     -> die "sudo not in $PATH."-}
  {-export "SUDO" $ pack $ encodeString sudo-}
  {-let shellArguments =-}
        {-[unpack argopkgs]-}
          {-++ [ "--keep"-}
             {-, "SUDO"-}
             {-, "--pure"-}
             {-, "--allow-new-privileges"-}
             {-, "--option"-}
             {-, "build-extra-sandbox-paths"-}
             {-, encodeString (directory sudo)-}
             {-]-}
          {-++ nixArguments-}
          {-++ (case run of-}
               {-Just cmd -> ["--run", unpack ("\"exec " <> cmd <> "\"")]-}
               {-Nothing  -> []-}
             {-)-}

  {-cachedir <- single $ inproc-}
    {-"mktemp"-}
    {-["-d", "--suffix=nixcache", "/tmp/deletable-nix-cache-XXXX"]-}
    {-empty-}
  {-export "XDG_CACHE_HOME" $ lineToText cachedir-}
  {-case targetMachine of-}
    {-Nothing -> do-}
      {-printf s "Running nix-shell with the following arguments: \n"-}
      {-liftIO $ print shellArguments-}
      {-liftIO $ executeFile "nix-shell" True shellArguments Nothing-}
    {-Just host -> do-}
      {-export "NIX_SSHOPTS" "source .profile; source .bash_profile;"-}
      {-printf s "argonix: creating a derivation in the local store:\n"-}

      {-drv <- single $ inproc-}
        {-"nix-instantiate"-}
        {-(["--quiet", argopkgs, "-A", "test"] ++ fmap pack nixArguments)-}
        {-empty-}

      {-printf (s % "\n") $ lineToText drv-}

      {-if remoteBuild-}
        {-then do-}
          {-printf-}
            {-s-}
            {-"argonix: copying the derivation's closure (without build output).\n"-}
          {-proc "nix-copy-closure" ["--to", host, lineToText drv] empty-}
        {-else do-}
          {-printf s "argonix: building the derivation's output:\n"-}
          {-proc "nix-store" ["--realize", lineToText drv, "--quiet"] empty-}

          {-printf-}
            {-s-}
            {-"argonix: copying the derivation's closure (with build output).\n"-}
          {-proc "nix-copy-closure"-}
               {-["--include-outputs", "--to", host, lineToText drv]-}
               {-empty-}

      {-case run of-}
        {-Just cmd -> do-}
          {-_ <- proc-}
            {-"ssh"-}
            {-[ "-t"-}
            {-, host-}
            {-, "source .profile; source .bash_profile; SUDO=$(which sudo) exec nix-shell --keep SUDO --pure --allow-new-privileges --option build-extra-sandbox-paths $(which sudo | xargs dirname) "-}
            {-<> lineToText drv-}
            {-<> " --run "-}
            {-<> "\""-}
            {-<> cmd-}
            {-<> "\""-}
            {-]-}
            {-empty-}
          {-Prelude.mapM_-}
            {-(\r -> proc "scp"-}
                        {-["-r", host <> ":" <> r, fromMaybe "." retreiveAs]-}
                        {-empty-}
            {-)-}
            {-retreive-}
        {-Nothing -> liftIO $ executeFile-}
          {-"ssh"-}
          {-True-}
          {-[ "-t"-}
          {-, unpack host-}
          {-, unpack-}
            {-("source .profile; source .bash_profile; SUDO=$(which sudo) exec nix-shell --keep SUDO --pure --allow-new-privileges --option build-extra-sandbox-paths $(which sudo | xargs dirname) "-}
            {-<> lineToText drv-}
            {-)-}
          {-]-}
          {-Nothing-}
 {-where-}
  {-opts = info-}
    {-(commonParser <**> helper)-}
    {-(  fullDesc-}
    {-<> progDesc "Argo environment provisioning/deployment/execution."-}
    {-<> header "argonix"-}
    {-)-}