{-# 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"
    , "numabench"
    , "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"

socketList :: [FilePath]
socketList =
  [ "/tmp/nrm-downstream-in"
  , "/tmp/nrm-downstream-event"
  , "/tmp/nrm-upstream-in"
  , "/tmp/nrm-upstream-event"
  ]

setupSystem :: ArgsCommon -> Shell ()
setupSystem sa = do
  doVerbose $ printInfo "Setting the nix-build environment up."
  doVerbose $ printInfo "Cleaning sockets."
  Prelude.mapM_ cleanSocket socketList
  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)
  vshell "sudo rm -rf nixtmpdir" empty >>= \case
    ExitSuccess   -> printInfo "removed ./nixtmpdir"
    ExitFailure n -> die ("Failed to remove ./nixtmpdir " <> 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
  liftIO $ createAndSetTMPDIR "nixtmpdir"
  void $ printInfo "Done setting the environment for nix-build up."
 where
  vshell    = verboseShell verbose
  verbose   = verbosity sa == Verbose
  doVerbose = when verbose
  createAndSetTMPDIR name = do
    {-localpath <- (</> fromText name) <$> pwd-}
    let path = "/tmp" </> fromText name
    testpath path >>= flip when (rmtree path >> mkdir path)
    case toText path of
      Right p -> sh $ do
        shell ("chmod 777 " <> p) empty >>= \case
          ExitSuccess   -> printf ("Successfully chmod" % s %"\n") p
          ExitFailure _ -> die "failed at chmod"
        shell ("ln -s " <> p <> name) empty >>= \case
          ExitSuccess   -> printf ("Successfully ln -s" % s %"\n") p
          ExitFailure _ -> die "failed at ln -s"
        export "TMPDIR" p
        printInfo $ format ("TMPDIR exported to " % s) p
      Left _ -> die "Path error when setting TMPDIR"

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

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

-- 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/"-}
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)
  void $ shell "echo \"TMPDIR=$TMPDIR\"" empty
  void $ shell "echo \"XDG_CACHE_HOME=$XDG_CACHE_HOME\"" empty
  liftIO (executeFile nixc True arglist Nothing)
 where
  nixc = toCommand nixCommand
  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"