module Hydrogen.Util.CliArgs (
    Option
  , switch
  , optarg
  , (~:)
  , getOpts
  , getOpts'
 ) where

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

switch, optarg :: String -> Option

switch = OptSwitch
optarg = OptArg

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

data Option =
    OptArg String
  | OptSwitch String
  | OptShort Char Option
  deriving (Eq, Show, Generic, Typeable)

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

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

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

getOpts :: [Option] -> IO (Map String String, Set String, [String])
getOpts opts = getOpts' opts <$> getArgs

getOpts' :: [Option] -> [String] -> (Map String String, Set String, [String])
getOpts' opts = readArgs Map.empty Set.empty

  where

    readArgs :: Map String String -> Set String -> [String]
        -> (Map String String, Set String, [String])
    readArgs args switches = \case
        x : xs
            | x =~ "^--[^-]" -> drop 2 x |> \case
                opt
                    | opt `elem` longArgs && not (null xs) -> readOptArg opt
                    | opt `elem` longSwitches -> readOptSwitch opt

                _ -> readArg

            | x =~ "^-[^-]$" -> x !! 1 |> \case
                shortOpt
                    | shortOpt `elem` shortArgs && not (null xs) -> readOptArg opt
                    | shortOpt `elem` shortSwitches -> readOptSwitch opt
                  where
                    opt = aliasMap ! shortOpt

                _ -> readArg

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

            | otherwise -> readArg

          where
            readArg = (args', switches', x : xs')
            readOptArg opt = readArgs (Map.insert opt (head xs) args) switches (tail xs)
            readOptSwitch opt = readArgs args (Set.insert opt switches) 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)