    {-# 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