{-# LANGUAGE DeriveFunctor, GeneralizedNewtypeDeriving #-}
module OptionParser(optionParser, Options(..), CompilerStage(..)) where

import Control.Monad
import Control.Monad.Except
import Control.Monad.State.Strict
import Data.List
import Data.Maybe
import qualified Data.Map.Strict as Map
import Debug.Trace

import Defs (Error)
import Optimiser (OptimiserLevel(..))


data Options = Options
    { oOutputName :: String
    , oSourceFiles :: [String]
    , oLastStage :: CompilerStage
    , oShowHelp :: Bool
    , oAppendExtension :: Bool
    , oOptimise :: OptimiserLevel
    }
  deriving Show

initOptions :: Options
initOptions = Options
    { oOutputName = "z_output"
    , oSourceFiles = []
    , oLastStage = StageExecutable
    , oShowHelp = False
    , oAppendExtension = True
    , oOptimise = Level0
    }

data CompilerStage = StageTypeCheck | StageIR | StageAsm | StageObject | StageExecutable
  deriving (Show, Eq, Ord, Enum)


data DBEntry
    = WithArgument (String -> Options -> Error Options)
    | NoArgument (Options -> Error Options)
    | MaybeArgument (Maybe String -> Options -> Error Options)

optionDatabase :: Map.Map String DBEntry
optionDatabase = Map.fromList $ map swap . concatMap sequence . map swap $
    [ (["--help", "-h"], NoArgument $ \o -> return $ o {oShowHelp = True})
    , (["--output", "-o"], WithArgument $ \s o -> return $ o {oOutputName = s, oAppendExtension = False})
    , (["-c"], NoArgument $ \o -> return $ o {oLastStage = StageObject})
    , (["-S"], NoArgument $ \o -> return $ o {oLastStage = StageAsm})
    , (["--optimise", "-O"], MaybeArgument $ \ms o -> case ms of
            Nothing -> return $ o {oOptimise = Level1}
            Just "0" -> return $ o {oOptimise = Level0}
            Just "1" -> return $ o {oOptimise = Level1}
            Just val -> Left $ "Invalid optimiser level '" ++ val ++ "'")
    , (["--stage"], WithArgument $ \s o -> case s of
            "typecheck" -> return $ o {oLastStage = StageTypeCheck}
            "ir" -> return $ o {oLastStage = StageIR}
            "asm" -> return $ o {oLastStage = StageAsm}
            "obj" -> return $ o {oLastStage = StageObject}
            "exe" -> return $ o {oLastStage = StageExecutable}
            _ -> Left $ "Unrecognised value for '--stage'")
    ]

swap :: (a, b) -> (b, a)
swap (a, b) = (b, a)


optionParser :: [String] -> Error Options
optionParser args = foldM foldfunc initOptions kvs
  where
    kvs = collectKeyValues args
    foldfunc opts (arg@('-':_), mval) = case (Map.lookup arg optionDatabase, mval) of
        (Nothing, _) -> Left $ "Unknown command-line option '" ++ arg ++ "'"
        (Just (NoArgument f), Nothing) -> f opts
        (Just (NoArgument _), Just val) ->
            Left $ "Option '" ++ arg ++ "' takes no argument, but got '" ++ val ++ "'"
        (Just (WithArgument f), Just val) -> f val opts
        (Just (WithArgument _), Nothing) ->
            Left $ "Option '" ++ arg ++ "' takes an argument, but got none"
        (Just (MaybeArgument f), _) -> f mval opts
    foldfunc opts (arg, Nothing) = return $ opts {oSourceFiles = oSourceFiles opts ++ [arg]}
    foldfunc _ (_, Just _) = undefined

collectKeyValues :: [String] -> [(String, Maybe String)]
collectKeyValues [] = []
collectKeyValues (arg@('-':'-':_:_):rest) = case findIndex (== '=') arg of
    Nothing -> case Map.lookup arg optionDatabase of
        Nothing -> (arg, Nothing) : collectKeyValues rest
        Just (NoArgument _) -> (arg, Nothing) : collectKeyValues rest
        Just (WithArgument _) -> case rest of
            [] -> [(arg, Nothing)]
            (hd:tl) -> (arg, Just hd) : collectKeyValues tl
        Just (MaybeArgument _) -> (arg, Nothing) : collectKeyValues rest
    Just idx -> (take idx arg, Just $ drop (idx + 1) arg) : collectKeyValues rest
collectKeyValues (arg@('-':_:_:_):rest) = case Map.lookup (take 2 arg) optionDatabase of
    Nothing -> (take 2 arg, Just $ drop 2 arg) : collectKeyValues rest
    Just (NoArgument _) -> (take 2 arg, Nothing) : collectKeyValues (('-' : drop 2 arg) : rest)
    Just (WithArgument _) -> (take 2 arg, Just $ drop 2 arg) : collectKeyValues rest
    Just (MaybeArgument _) -> (take 2 arg, Just $ drop 2 arg) : collectKeyValues rest
collectKeyValues (arg@['-',_]:rest) = case Map.lookup arg optionDatabase of
    Nothing -> (arg, Nothing) : collectKeyValues rest
    Just (NoArgument _) -> (arg, Nothing) : collectKeyValues rest
    Just (WithArgument _) -> case rest of
        [] -> [(arg, Nothing)]
        (next:rest2) -> (arg, Just next) : collectKeyValues rest2
    Just (MaybeArgument _) -> (arg, Nothing) : collectKeyValues rest
collectKeyValues (arg:rest) = (arg, Nothing) : collectKeyValues rest