module Main where

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

import BuildIR
import CodeGen
import Defs
import LibLang
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

withTempFile :: (FilePath -> IO a) -> IO a
withTempFile = bracket mkTempFile removeFile

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 withObjfile = case laststage of
            StageObject -> ($ outname)
            StageExecutable -> withTempFile
            _ -> undefined

    let yasmFormat = case System.os of
            "linux" -> "elf64"
            "darwin" -> "macho64"
            os -> error $ "Your OS (" ++ os ++ ") is unknown, can't create binary"

    let yasmToFile outfname writer = do
            let yasmprocspec =
                    (proc "yasm" ["-w+all", "-f" ++ yasmFormat, "-", "-o", outfname])
                        {std_in = CreatePipe}
            yasmcode <- withCreateProcess yasmprocspec $ \(Just pipe) _ _ ph -> do
                _ <- writer pipe
                hFlush pipe
                hClose pipe
                waitForProcess ph
            case yasmcode of
                ExitSuccess -> return ()
                ExitFailure _ -> die "yasm failed!"

    withObjfile $ \objfname -> do
        hPutStrLn stderr "Assembling with yasm..."
        yasmToFile objfname (\pipe -> hPutStr pipe asm)

        case laststage of
            StageObject -> return ()
            StageExecutable ->
                withTempFile $ \liblangfname -> do
                    yasmToFile liblangfname (\pipe -> BS.hPut pipe libLangSource)
                    hPutStrLn stderr "Linking with ld..."
                    callProcess "ld" [objfname, liblangfname, "-o", outname]
            _ -> undefined