module Distribution.Ecstatic (run) where

import Control.Exception (bracket)
import Control.Monad (when)
import Data.List (stripPrefix)
import System.Directory (getTemporaryDirectory, removeFile)
import System.Exit (exitWith, ExitCode)
import System.IO (Handle, hClose, hFlush, hPutStr, openTempFile)
import System.Process (withCreateProcess, proc, waitForProcess)


-- ghc uses gcc as linker
realLinker :: String
realLinker = "gcc"

-- When setting a custom linker via -pgml, ghc doesn't pass these arguments.
-- So we restore them manually.
-- NOTE: -Wl,--no-as-needed must be the first argument
defaultLinkerArgs :: [String]
defaultLinkerArgs = ["-Wl,--no-as-needed", "-no-pie"]

debug :: Bool
debug = False

runLinker :: [String] -> IO ExitCode
runLinker args = withCreateProcess (proc realLinker args) $
    \_ _ _ h -> waitForProcess h

debugLog :: String -> IO ()
debugLog s = when debug $ appendFile "/tmp/ecstatic-log.txt" $ s ++ "\n"

-- Read all command-line arguments, processing `@file's recursively.
-- Quoting 'man gcc':
--
-- @file
--     Read command-line options from file.  The options read are inserted
--     in place of the original @file option.  If file does not exist, or
--     cannot be read, then the option will be treated literally, and not
--     removed.

--     Options in file are separated by whitespace.  A whitespace character
--     may be included in an option by surrounding the entire option in
--     either single or double quotes.  Any character (including a
--     backslash) may be included by prefixing the character to be included
--     with a backslash.  The file may itself contain additional @file
--     options; any such options will be processed recursively.
--
-- XXX: some quoted arguments may be processed incorrectly (e.g. if a quoted
-- argument contains the `@file' option).
collectArgs :: [String] -> IO [String]
collectArgs [] = pure []
collectArgs (arg:rest) = do
    fileArgs <- case arg of
        ('@':file) -> do
            fileLines <- lines <$> readFile file
            collectArgs fileLines
        _ -> pure [arg]
    restArgs <- collectArgs rest
    pure $ fileArgs ++ restArgs

isOption :: String -> String -> Bool
isOption expected actual = actual `elem` [expected, "\"" ++ expected ++ "\"", "'" ++ expected ++ "'"]

stripSuffix :: Eq a => [a] -> [a] -> Maybe [a]
stripSuffix suf s = reverse <$> stripPrefix (reverse suf) (reverse s)

getLinkedLib :: String -> Maybe String
getLinkedLib ('-':'l':rest) = Just rest
getLinkedLib ('"':'-':'l':rest) = stripSuffix "\"" rest
getLinkedLib ('\'':'-':'l':rest) = stripSuffix "'" rest
getLinkedLib _ = Nothing

shouldBuildStatically :: String -> Bool
shouldBuildStatically lib = lib `notElem` ["c", "m", "pthread", "dl", "rt"]

transformArgs :: [String] -> [String]
transformArgs = go [] False where
    toStaticLink = "-Wl,-Bstatic"
    toDynamicLink = "-Wl,-Bdynamic"
    go acc _ [] = reverse acc
    go acc isStatic (arg:rest) = case getLinkedLib arg of
      Just lib | not isStatic && shouldBuildStatically lib ->
            go (toDynamicLink:arg:toStaticLink:acc) isStatic rest
      _ | isOption toStaticLink arg  -> go (arg:acc) True rest
        | isOption toDynamicLink arg -> go (arg:acc) False rest
      _ -> go (arg:acc) isStatic rest

withTempFile :: FilePath -> (FilePath -> Handle -> IO a) -> IO a
withTempFile fileNameTemplate action = do
    tempDir <- getTemporaryDirectory
    bracket
        (openTempFile tempDir fileNameTemplate)
        (\(path, h) -> hClose h >> removeFile path)
        (\(path, h) -> action path h)

-- | Entry point to the program. Accepts the list of command-line arguments
-- to the linker (gcc).
run :: [String] -> IO ()
run selfArgs = do
    args <- collectArgs selfArgs
    debugLog $ unlines args
    let isArg s = any (isOption s) args
    when (not (isArg "-o") || isArg "-shared") $ do
        -- ghc calls linker many times, not just for the final executable.
        -- In particular, some autogen'ed files are linked into shared libraries.
        -- Don't try to statically link in that case, simply pass through the arguments.
        let newArgs = defaultLinkerArgs ++ selfArgs
        runLinker newArgs >>= exitWith

    let newArgs = defaultLinkerArgs ++ ["-static-libgcc", "-static-libstdc++"] ++ transformArgs args
    debugLog $ unlines newArgs
    exitCode <- withTempFile "ecstatic.rsp" $ \path h -> do
        hPutStr h $ unlines newArgs
        hFlush h
        hClose h
        runLinker ['@':path]
    exitWith exitCode