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