module Main where

import Control.Exception (bracket)
import Control.Monad
import System.Directory (removeFile)
import System.Environment
import System.Exit
import System.IO
import System.Process
import qualified System.Posix.Temp as Posix

import BuildIR
import CodeGen
import Defs
import Optimiser
import OptionParser
import Pretty
import ProgramParser
import TypeCheck
import Verify


infix 2 <?>
(<?>) :: (a -> Error b) -> String -> a -> Error b
f <?> pre = \a -> case f a of
    Left e -> Left $ pre ++ ": " ++ e
    Right x -> Right x


eitherToIO :: Either String a -> IO a
eitherToIO = either die return

extensionForStage :: CompilerStage -> String
extensionForStage StageTypeCheck = ".lang"
extensionForStage StageIR = ".ir"
extensionForStage StageAsm = ".asm"
extensionForStage StageObject = ".o"
extensionForStage StageExecutable = ""

inputFile :: String -> IO String
inputFile "-" = getContents
inputFile fname = readFile fname

outputFile :: String -> String -> IO ()
outputFile "-" str = putStr str
outputFile fname str = writeFile fname str

mkTempFile :: IO FilePath
mkTempFile = do
    (path, handle) <- Posix.mkstemp "/tmp/tmp.lang."
    hClose handle
    return path

main :: IO ()
main = do
    opts <- getArgs >>= eitherToIO . optionParser

    when (oShowHelp opts) $ do
        name <- getProgName
        putStrLn $ "Usage: " ++ name ++ " [options] files..."
        putStrLn ""
        putStrLn "-h, --help         Show this help"
        putStrLn "-o, --output FILE  Send output to the specified file instead of z_output[.ext]."
        putStrLn "-O, --optimise     Optimise the output. (Happens at IR stage.)"
        putStrLn "--stage typecheck  Only parse and typecheck the program, outputting a pretty-"
        putStrLn "                   printed version."
        putStrLn "--stage ir         Compile until (possibly optimised) IR, then output that."
        putStrLn "-S, --stage asm    Compile until assembly, then output that."
        putStrLn "-c, --stage obj    Compile to an object file."
        putStrLn "--stage exe        Default mode; compile and link."
        exitSuccess

    source <- case oSourceFiles opts of
        [] -> die "No source files given. Run with --help for usage information."
        [name] -> inputFile name
        _ -> die "Only one source file supported at the moment."

    when (oOutputName opts == "-" && oLastStage opts >= StageObject) $
        die "Will not write binary file to stdout."

    let outname = if oAppendExtension opts
                    then oOutputName opts ++ extensionForStage (oLastStage opts)
                    else oOutputName opts

    let laststage = oLastStage opts

    ast' <- eitherToIO $ (parseProgram <?> "Parse error") source
    ast <- eitherToIO $ (typeCheck <?> "Type error") ast'
    when (laststage == StageTypeCheck) $ do
        outputFile outname $ pretty ast
        exitSuccess

    ir <- eitherToIO $ (buildIR <?> "IR building error") ast
    iropt' <- eitherToIO $ (optimise (oOptimise opts) <?> "Error while optimising") ir
    iropt <- eitherToIO $ (verify <?> "Verify error") iropt'
    when (laststage == StageIR) $ do
        outputFile outname $ pretty iropt
        exitSuccess

    asm <- eitherToIO $ (codegen <?> "Codegen error") iropt
    when (laststage == StageAsm) $ do
        outputFile outname asm
        exitSuccess


    let objfnameIO = case laststage of
            StageObject -> return outname
            StageExecutable -> mkTempFile
            _ -> undefined

    let rmObjfileIO = case outname of
            "-" -> const $ return ()
            _ -> removeFile

    bracket objfnameIO rmObjfileIO $ \objfname -> do
        hPutStrLn stderr "Assembling with yasm..."
        let yasmprocspec = (proc "yasm" ["-w+all", "-fmacho64", "-", "-o", objfname]) {std_in = CreatePipe}
        yasmcode <- withCreateProcess yasmprocspec $ \(Just pipe) _ _ ph -> do
            hPutStr pipe asm
            hFlush pipe
            hClose pipe
            waitForProcess ph
        case yasmcode of
            ExitSuccess -> return ()
            ExitFailure _ -> die "yasm failed!"
        when (laststage == StageObject) exitSuccess

        hPutStrLn stderr "Linking with ld..."
        callProcess "ld" [objfname, "liblang.o", "-o", outname]