{-# LANGUAGE TypeFamilies #-}

-- |
-- Module: Options.Tokenize
-- License: MIT
module Options.Tokenize
        ( Token(..)
        , tokenFlagName
        , Tokens(..)
        , tokensMap
        , tokenize
        ) where

import           Control.Applicative
import           Control.Monad.Error hiding (throwError)
import qualified Control.Monad.Error
import           Control.Monad.State
import           Data.Functor.Identity
import qualified Data.Map

import           Options.Types
import           Options.Util

data Token
        = TokenUnary String -- flag name
        | Token String String -- flag name, flag value
        deriving (Eq, Show)

tokenFlagName :: Token -> String
tokenFlagName (TokenUnary s) = s
tokenFlagName (Token s _) = s

data Tokens = Tokens
        { tokensList :: [([OptionKey], Token)]
        , tokensArgv :: [String]
        }
        deriving (Show)

tokensMap :: Tokens -> Data.Map.Map OptionKey [Token]
tokensMap tokens = Data.Map.fromListWith (\xs ys -> ys ++ xs) $ do
        (keys, token) <- tokensList tokens
        key <- keys
        return (key, [token])

data TokState = TokState
        { stArgv :: [String]
        , stArgs :: [String]
        , stOpts :: [([OptionKey], Token)]
        , stShortKeys :: Data.Map.Map Char ([OptionKey], OptionInfo)
        , stLongKeys :: Data.Map.Map String ([OptionKey], OptionInfo)
        , stSubcommands :: [(String, [OptionInfo])]
        , stSubCmd :: Maybe String
        }

newtype Tok a = Tok { unTok :: ErrorT String (StateT TokState Identity) a }

instance Functor Tok where
        fmap = liftM

instance Applicative Tok where
        pure = return
        (<*>) = ap

instance Monad Tok where
        return = Tok . return
        m >>= f = Tok (unTok m >>= unTok . f)

instance MonadState Tok where
        type StateType Tok = TokState
        get = Tok get
        put = Tok . put

tokenize :: OptionDefinitions -> [String] -> (Maybe String, Either String Tokens)
tokenize (OptionDefinitions options subcommands) argv = runIdentity $ do
        let st = TokState
                { stArgv = argv
                , stArgs = []
                , stOpts = []
                , stShortKeys = toShortKeys options
                , stLongKeys = toLongKeys options
                , stSubcommands = subcommands
                , stSubCmd = Nothing
                }
        (err, st') <- runStateT (runErrorT (unTok loop)) st
        return (stSubCmd st', case err of
                Left err' -> Left err'
                Right _ -> Right (Tokens (reverse (stOpts st')) (stArgs st')))

loop :: Tok ()
loop = do
        ms <- nextItem
        st <- get
        case ms of
                Nothing -> return ()
                Just s -> (>> loop) $ case stringToGhc704 s of
                        '-':'-':[] -> put (st { stArgv = [], stArgs = stArgs st ++ stArgv st })
                        '-':'-':opt -> parseLong opt
                        '-':optChar:optValue -> parseShort optChar optValue
                        '-':[] -> addArg s
                        decoded -> case (stSubcommands st, stSubCmd st) of
                                ([], _) -> addArg s
                                (_, Just _) -> addArg s
                                (_, Nothing) -> case lookup decoded (stSubcommands st) of
                                        Nothing -> throwError ("Unknown subcommand " ++ show decoded ++ ".")
                                        Just subOptions -> mergeSubcommand decoded subOptions

nextItem :: Tok (Maybe String)
nextItem = do
        st <- get
        case stArgv st of
                [] -> return Nothing
                (x:xs) -> do
                        put (st { stArgv = xs })
                        return (Just x)

addArg :: String -> Tok ()
addArg s = modify (\st -> st { stArgs = stArgs st ++ [s] })

addOpt :: [OptionKey] -> Token  -> Tok ()
addOpt keys val = modify (\st -> st
        { stOpts = (keys, val) : stOpts st
        })

mergeSubcommand :: String -> [OptionInfo] -> Tok ()
mergeSubcommand name opts = modify $ \st -> st
        { stSubCmd = Just name
        , stShortKeys = Data.Map.unionWith unionKeys (stShortKeys st) (toShortKeys opts)
        , stLongKeys = Data.Map.unionWith unionKeys (stLongKeys st) (toLongKeys opts)
        }

-- note: unionKeys assumes that the OptionInfo is equivalent in both maps.
unionKeys :: ([OptionKey], OptionInfo) -> ([OptionKey], OptionInfo) -> ([OptionKey], OptionInfo)
unionKeys (keys1, info) (keys2,_) = (keys1++keys2, info)

parseLong :: String -> Tok ()
parseLong optName = do
        longKeys <- gets stLongKeys
        case break (== '=') optName of
                (before, after) -> case after of
                        '=' : value -> case Data.Map.lookup before longKeys of
                                Nothing -> throwError ("Unknown flag --" ++ before)
                                Just (keys, info) -> if optionInfoUnaryOnly info
                                        then throwError ("Flag --" ++ before ++ " takes no parameters.")
                                        else addOpt keys (Token ("--" ++ before) value)
                        _ -> case Data.Map.lookup optName longKeys of
                                Nothing -> throwError ("Unknown flag --" ++ optName)
                                Just (keys, info) -> if optionInfoUnary info
                                        then addOpt keys (TokenUnary ("--" ++ optName))
                                        else do
                                                next <- nextItem
                                                case next of
                                                        Nothing -> throwError ("The flag --" ++ optName ++ " requires a parameter.")
                                                        Just value -> addOpt keys (Token ("--" ++ optName) value)

parseShort :: Char -> String -> Tok ()
parseShort optChar optValue = do
        let optName = '-' : [optChar]
        shortKeys <- gets stShortKeys
        case Data.Map.lookup optChar shortKeys of
                Nothing -> throwError ("Unknown flag " ++ optName)
                Just (keys, info) -> if optionInfoUnary info
                        -- don't check optionInfoUnaryOnly, because that's only set by --help
                        -- options and they define no short flags.
                        then do
                                addOpt keys (TokenUnary optName)
                                case optValue of
                                        [] -> return ()
                                        nextChar:nextValue -> parseShort nextChar nextValue
                        else case optValue of
                                "" -> do
                                        next <- nextItem
                                        case next of
                                                Nothing -> throwError ("The flag " ++ optName ++ " requires a parameter.")
                                                Just value -> addOpt keys (Token optName value)
                                _ -> addOpt keys (Token optName optValue)

toShortKeys :: [OptionInfo] -> Data.Map.Map Char ([OptionKey], OptionInfo)
toShortKeys opts = Data.Map.fromListWith (\(keys1, info) (keys2, _) -> (keys2 ++ keys1, info)) $ do
        opt <- opts
        flag <- optionInfoShortFlags opt
        return (flag, ([optionInfoKey opt], opt))

toLongKeys :: [OptionInfo] -> Data.Map.Map String ([OptionKey], OptionInfo)
toLongKeys opts = Data.Map.fromListWith (\(keys1, info) (keys2, _) -> (keys2 ++ keys1, info)) $ do
        opt <- opts
        flag <- optionInfoLongFlags opt
        return (flag, ([optionInfoKey opt], opt))

throwError :: String -> Tok a
throwError = Tok . Control.Monad.Error.throwError