module Propellor.CmdLine (
        defaultMain,
        processCmdLine,
) where

import System.Environment (getArgs)
import Data.List
import System.Exit
import System.PosixCompat
import Network.Socket

import Propellor.Base
import Propellor.Gpg
import Propellor.Git
import Propellor.Git.VerifiedBranch
import Propellor.Bootstrap
import Propellor.Spin
import Propellor.Types.CmdLine
import qualified Propellor.Property.Docker as Docker
import qualified Propellor.Property.Chroot as Chroot
import qualified Propellor.Shim as Shim
import Utility.FileSystemEncoding

usage :: Handle -> IO ()
usage h = hPutStrLn h $ unlines
        [ "Usage:"
        , "  with no arguments, provision the current host"
        , ""
        , "  --init"
        , "      initialize ~/.propellor"
        , "  hostname"
        , "      provision the current host as if it had the specified hostname"
        , "  --spin targethost [--via relayhost]"
        , "      provision the specified host"
        , "  --build"
        , "      recompile using your current config"
        , "  --add-key keyid"
        , "      add an additional signing key to the private data"
        , "  --rm-key keyid"
        , "      remove a signing key from the private data"
        , "  --list-fields"
        , "      list private data fields"
        , "  --set field context"
        , "      set a private data field"
        , "  --unset field context"
        , "      clear a private data field"
        , "  --unset-unused"
        , "      clear unused fields from the private data"
        , "  --dump field context"
        , "      show the content of a private data field"
        , "  --edit field context"
        , "      edit the content of a private data field"
        , "  --merge"
        , "      combine multiple spins into a single git commit"
        , "  --check"
        , "      double-check that propellor can actually run here"]

usageError :: [String] -> IO a
usageError ps = do
        usage stderr
        error ("(Unexpected: " ++ show ps)

processCmdLine :: IO CmdLine
processCmdLine = go =<< getArgs
  where
        go ("--check":_) = return Check
        go ("--spin":ps) = case reverse ps of
                (r:"--via":hs) -> Spin
                        <$> mapM hostname (reverse hs)
                        <*> pure (Just r)
                _ -> Spin <$> mapM hostname ps <*> pure Nothing
        go ("--build":[]) = return Build
        go ("--add-key":k:[]) = return $ AddKey k
        go ("--rm-key":k:[]) = return $ RmKey k
        go ("--set":f:c:[]) = withprivfield f c Set
        go ("--unset":f:c:[]) = withprivfield f c Unset
        go ("--unset-unused":[]) = return UnsetUnused
        go ("--dump":f:c:[]) = withprivfield f c Dump
        go ("--edit":f:c:[]) = withprivfield f c Edit
        go ("--list-fields":[]) = return ListFields
        go ("--merge":[]) = return Merge
        go ("--help":_) = do
                usage stdout
                exitFailure
        go ("--boot":_:[]) = return $ Update Nothing -- for back-compat
        go ("--serialized":s:[]) = serialized Serialized s
        go ("--continue":s:[]) = serialized Continue s
        go ("--gitpush":fin:fout:_) = return $ GitPush (Prelude.read fin) (Prelude.read fout)
        go ("--run":h:[]) = go [h]
        go (h:[])
                | "--" `isPrefixOf` h = usageError [h]
                | otherwise = Run <$> hostname h
        go [] = do
                s <- takeWhile (/= '\n') <$> readProcess "hostname" ["-f"]
                if null s
                        then errorMessage "Cannot determine hostname! Pass it on the command line."
                        else return $ Run s
        go v = usageError v

        withprivfield s c f = case readish s of
                Just pf -> return $ f pf (Context c)
                Nothing -> errorMessage $ "Unknown privdata field " ++ s

        serialized mk s = case readish s of
                Just cmdline -> return $ mk cmdline
                Nothing -> errorMessage $ "serialization failure (" ++ s ++ ")"

data CanRebuild = CanRebuild | NoRebuild

-- | Runs propellor on hosts, as controlled by command-line options.
defaultMain :: [Host] -> IO ()
defaultMain hostlist = withConcurrentOutput $ do
        useFileSystemEncoding
        setupGpgEnv
        Shim.cleanEnv
        checkDebugMode
        cmdline <- processCmdLine
        debug ["command line: ", show cmdline]
        go CanRebuild cmdline
  where
        go cr (Serialized cmdline) = go cr cmdline
        go _ Check = return ()
        go cr Build = buildFirst Nothing cr Build $ return ()
        go _ (Set field context) = setPrivData field context
        go _ (Unset field context) = unsetPrivData field context
        go _ (UnsetUnused) = unsetPrivDataUnused hostlist
        go _ (Dump field context) = dumpPrivData field context
        go _ (Edit field context) = editPrivData field context
        go _ ListFields = listPrivDataFields hostlist
        go _ (AddKey keyid) = addKey keyid
        go _ (RmKey keyid) = rmKey keyid
        go _ c@(ChrootChain _ _ _ _ _) = Chroot.chain hostlist c
        go _ (DockerChain hn cid) = Docker.chain hostlist hn cid
        go _ (DockerInit hn) = Docker.init hn
        go _ (GitPush fin fout) = gitPushHelper fin fout
        go cr (Relay h) = forceConsole >>
                updateFirst Nothing cr (Update (Just h)) (update (Just h))
        go _ (Update Nothing) = forceConsole >>
                fetchFirst (onlyprocess (update Nothing))
        go _ (Update (Just h)) = update (Just h)
        go _ Merge = mergeSpin
        go cr cmdline@(Spin hs mrelay) = buildFirst Nothing cr cmdline $ do
                unless (isJust mrelay) commitSpin
                forM_ hs $ \hn -> withhost hn $ spin mrelay hn
        go cr cmdline@(Run hn) = ifM ((==) 0 <$> getRealUserID)
                ( updateFirst (findHost hostlist hn) cr cmdline $ runhost hn
                , fetchFirst $ go cr (Spin [hn] Nothing)
                )
        go cr cmdline@(SimpleRun hn) = forceConsole >>
                fetchFirst (buildFirst (findHost hostlist hn) cr cmdline (runhost hn))
        -- When continuing after a rebuild, don't want to rebuild again.
        go _ (Continue cmdline) = go NoRebuild cmdline

        withhost :: HostName -> (Host -> IO ()) -> IO ()
        withhost hn a = maybe (unknownhost hn hostlist) a (findHost hostlist hn)

        runhost hn = onlyprocess $ withhost hn mainProperties

        onlyprocess = onlyProcess (localdir </> ".lock")

unknownhost :: HostName -> [Host] -> IO a
unknownhost h hosts = errorMessage $ unlines
        [ "Propellor does not know about host: " ++ h
        , "(Perhaps you should specify the real hostname on the command line?)"
        , "(Or, edit propellor's config.hs to configure this host)"
        , "Known hosts: " ++ unwords (map hostName hosts)
        ]

-- Builds propellor (when allowed) and if it looks like a new binary,
-- re-execs it to continue.
-- Otherwise, runs the IO action to continue.
--
-- The Host should only be provided when dependencies should be installed
-- as needed to build propellor.
buildFirst :: Maybe Host -> CanRebuild -> CmdLine -> IO () -> IO ()
buildFirst h CanRebuild cmdline next = do
        oldtime <- getmtime
        buildPropellor h
        newtime <- getmtime
        if newtime == oldtime
                then next
                else continueAfterBuild cmdline
  where
        getmtime = catchMaybeIO $ getModificationTime "propellor"
buildFirst _ NoRebuild _ next = next

continueAfterBuild :: CmdLine -> IO a
continueAfterBuild cmdline = go =<< boolSystem "./propellor"
        [ Param "--continue"
        , Param (show cmdline)
        ]
  where
        go True = exitSuccess
        go False = exitWith (ExitFailure 1)

fetchFirst :: IO () -> IO ()
fetchFirst next = do
        whenM hasOrigin $
                void fetchOrigin
        next

updateFirst :: Maybe Host -> CanRebuild -> CmdLine -> IO () -> IO ()
updateFirst h canrebuild cmdline next = ifM hasOrigin
        ( updateFirst' h canrebuild cmdline next
        , next
        )

-- If changes can be fetched from origin, builds propellor (when allowed)
-- and re-execs the updated propellor binary to continue.
-- Otherwise, runs the IO action to continue.
updateFirst' :: Maybe Host -> CanRebuild -> CmdLine -> IO () -> IO ()
updateFirst' h CanRebuild cmdline next = ifM fetchOrigin
        ( do
                buildPropellor h
                continueAfterBuild cmdline
        , next
        )
updateFirst' _ NoRebuild _ next = next

-- Gets the fully qualified domain name, given a string that might be
-- a short name to look up in the DNS.
hostname :: String -> IO HostName
hostname s = go =<< catchDefaultIO [] dnslookup
  where
        dnslookup = getAddrInfo (Just canonname) (Just s) Nothing
        canonname = defaultHints { addrFlags = [AI_CANONNAME] }
        go (AddrInfo { addrCanonName = Just v } : _) = pure v
        go _
                | "." `isInfixOf` s = pure s -- assume it's a fqdn
                | otherwise =
                        error $ "cannot find host " ++ s ++ " in the DNS"