module Hydrogen.CliArgs (
    Option
  , switch
  , optarg
  , (~:)
  , (~?)
  , (~=)
  , getOpts
  , getOpts'
  , Set
  , Has (..)
  , MultiMap
  , Container (..)
  , OptArgs
 ) where

import Hydrogen.Prelude.System
import qualified Hydrogen.MultiMap as MultiMap
import qualified Data.Map as Map
import qualified Data.Set as Set

type OptArgs = (MultiMap String String, Set String, [String])

switch, optarg :: String -> Option

switch = OptSwitch
optarg = OptArg

infixl 2 ~:
infixl 1 ~?

(~:) :: Char -> Option -> Option
(~:) = OptShort

(~?) :: Option -> (String -> Bool) -> Option
(~?) = flip OptCheck

(~=) :: Option -> String -> Option
(~=) = \opt pattern -> OptCheck (=~ pattern) opt

data Option =
    OptArg String
  | OptSwitch String
  | OptCheck (String -> Bool) Option
  | OptShort Char Option
  deriving (Show, Generic, Typeable)

isArg :: Option -> Bool
isArg = \case
    OptArg _ -> True
    OptShort _ x -> isArg x
    OptCheck _ x -> isArg x
    _ -> False

shorts :: Option -> [Char]
shorts = \case
    OptShort c xs -> c : shorts xs
    OptCheck _ xs -> shorts xs
    _ -> []

long :: Option -> String
long = \case
    OptArg x -> x
    OptSwitch x -> x
    OptShort _ xs -> long xs
    OptCheck _ xs -> long xs

check :: Option -> Maybe (String -> Bool)
check = \case
    OptArg _ -> Nothing
    OptSwitch _ -> Nothing
    OptShort _ xs -> check xs
    OptCheck x _ -> Just x

getOpts :: [Option] -> IO OptArgs
getOpts opts = getOpts' opts <$> getArgs

getOpts' :: [Option] -> [String] -> OptArgs
getOpts' opts = readArgs MultiMap.empty Set.empty

  where

    readArgs :: MultiMap String String -> Set String -> [String]
        -> (MultiMap String String, Set String, [String])
    readArgs args switches = \case
        x : xs
            | x =~ "^--[^ =-][^ =]*=" -> break (== '=') x |> \case
                (key, arg) | opt `elem` longArgs && test opt val
                    -> readArgs (MultiMap.insert opt val args) switches xs
                  where
                    val = tail arg
                    opt = drop 2 key

                _ -> readArg

            | x =~ "^--[^-]" -> drop 2 x |> \case
                opt
                    | opt `elem` longArgs && not (null xs) && test opt val
                        -> readArgs (MultiMap.insert opt val args) switches (tail xs)
                    | opt `elem` longSwitches
                        -> readArgs args (Set.insert opt switches) xs
                  where
                    val = head xs

                _ -> readArg

            | x =~ "^-[^ =-]" -> drop 1 x |> \case
                optString@(shortOpt : val)
                    | shortOpt `elem` shortArgs && not (null val) && test opt val
                        -> readArgs (MultiMap.insert opt val args) switches xs
                    | shortOpt `elem` shortArgs && not (null xs) && test opt val'
                        -> readArgs (MultiMap.insert opt val' args) switches (tail xs)
                    | otherwise
                        -> readOpts
                  where
                    val' = head xs
                    opt = aliasMap ! shortOpt
                    (recog, nrecog) = partition (`elem` shortSwitches) optString
                    shortOpts = map (aliasMap !) recog
                    readOpts
                      | null nrecog = (args'', switches'', xs'')
                      | otherwise   = (args'', switches'', ('-' : nrecog) : xs'')
                    (args'', switches'', xs'')
                        = readArgs args (foldr Set.insert switches shortOpts) xs

                _ -> readArg

            | x == "--" -> (args, switches, xs)

            | otherwise -> readArg

          where
            readArg = (args', switches', x : xs')
            (args', switches', xs') = readArgs args switches xs

        _ -> (args, switches, [])


    (optArgs, optSwitches) = partition isArg opts
    
    (longArgs, longSwitches) = (map long optArgs, map long optSwitches)
    (shortArgs, shortSwitches) = (concatMap shorts optArgs, concatMap shorts optSwitches)
    
    aliasMap = foldr aliases Map.empty opts
    aliases opt = flip (foldr (flip Map.insert (long opt))) (shorts opt)

    patternMap = foldr patterns Map.empty opts
    patterns opt = maybe id (Map.insert (long opt)) (check opt)

    test opt val = maybe True (val |>) (Map.lookup opt patternMap)