{-# LANGUAGE TypeFamilies, DeriveDataTypeable, ScopedTypeVariables #-}
module System.Console.CmdLib.Flag where

import System.Console.CmdLib.Attribute
import System.Console.GetOpt
import Data.Typeable
import Data.Data

import Data.List.Split( splitOn )
import Data.List( nub, intercalate, isSuffixOf, sortBy, partition )
import Data.Ord ( comparing )
import Data.Char( isUpper, toLower, toUpper )
import Data.Generics.Aliases( extR )
import Data.Generics.Text( gshow )
import Data.Maybe ( fromJust, isJust )

newtype PathF = PathF String deriving (Typeable, Data)

defaults = [ Default ""
           , Default (0 :: Int)
           , Default False
           , Default ([] :: [Int])
           , Default ([] :: [String])
           , Enabled True ]

data OptionType = SimpleOption | BooleanOption | OptionalArgument | RequiredArgument | BadOption
                deriving Eq

class Attributes a where
  attributes :: a -> AttributeMap Key
  attributes _ = EmptyMap
  readFlag :: Data b => a -> String -> b
  readFlag _ = readCommon

class (Eq flag) => FlagType flag where
  type Folded flag :: *

  flag_attrkey :: flag -> Key
  flag_attrs :: flag -> [Attribute]

  flag_args :: flag -> [Attribute] -> ArgDescr (Folded flag -> Folded flag)
  flag_set :: (Typeable a) => flag -> a -> Folded flag -> Folded flag
  flag_parse :: flag -> String -> Folded flag -> Folded flag
  flag_value :: (Typeable a) => flag -> Folded flag -> a
  flag_eval :: flag -> Folded flag -> IO ()
  flag_type :: flag -> [Attribute] -> OptionType
  flag_list :: [flag]
  flag_defaults :: (flag -> [Attribute]) -> Folded flag -> Folded flag
  flag_fold :: flag -> [Folded flag -> Folded flag] -> Folded flag -> Folded flag
  flag_empty :: flag -> Folded flag

instance FlagType () where
  type Folded () = ()
  flag_list = []

optionType :: [DataType] -> [TypeRep] -> OptionType
optionType dt t = case (dt, t) of
    ([], []) -> SimpleOption
    ([_], [t]) | typeRepTyCon t == typeRepTyCon (typeOf True) -> BooleanOption
    ([_], [t]) | typeRepTyCon t == typeRepTyCon (typeOf (Just ())) -> OptionalArgument
    ([_], _) -> RequiredArgument
    _ -> BadOption

nonoption attr x = isExtra attr x || isPositional attr x

isExtra attr      = extra . attr
isPositional attr = isJust . positional . attr
isRequired attr   = required . attr

formatOptDescrOptions (Option ss ls _ _) = (ss', ls') where
 ss' = map (\s -> '-' : [s]) ss
 ls' = map ("--" ++) ls

usageDescr :: forall flag. FlagType flag => (flag -> [Attribute]) -> [String]
usageDescr attr = optGroups ++ reqs' ++ nonopts where
  partitionFlag (os, ps, es) f | isExtra attr f      = (os, ps, f : es)
  partitionFlag (os, ps, es) f | isPositional attr f = (os, f : ps, es)
  partitionFlag (os, ps, es) f                       = (f : os, ps, es)

  (opts, positionals, extras) = foldl partitionFlag ([],[],[])
    [f | f <- flag_list, enabled $ attr f]

  (reqs, opts') = partition (isRequired attr) opts
  reqs' = fmtRequired `map` concatMap (flagToOptDescr MergeSuffix attr) reqs

  fmtRequired :: OptDescr a -> String
  fmtRequired od@(Option ss ls ad _) = separated ++ args ad where
    (ss', ls') = formatOptDescrOptions od
    all = ss' ++ ls'
    separated = case length all of
      0 -> ""
      1 -> head all
      _ -> "<" ++ intercalate "|" all ++ ">"

    args (NoArg  _    ) = ""
    args (ReqArg _  ad) = '=' : ad
    args (OptArg _  ad) = "[=" ++ ad ++ "]"

  optGroups = [ "[" ++ map toUpper g ++ "]" |
                g <- nub $ map (getgroup . attr) opts' ]

  nonopts = upperReqPositionals ++ upperPositionals ++ upperExtras where
    orderedPositionals = orderByPosition positionals
    splitAtLastRequired p (req, nonreq) = case req of
      [] -> if isRequired attr p
              then ([p], nonreq)
              else ([] , p : nonreq)
      _  -> (p : req, nonreq)
    (requiredPos, nonRequiredPos) = foldr splitAtLastRequired ([],[])
      orderedPositionals
    upperReqPositionals = map nonoptStr requiredPos
    upperPositionals =
      case nestedOptional $ map nonoptStr nonRequiredPos of
        [] -> []
        nOpts -> [nOpts]
    upperExtras = map ((\s -> "[" ++ s ++ "]" ) . nonoptStr)
      extras

  nonoptStr :: flag -> String
  nonoptStr f = extractArgDesc ad where
     (Option _ _ ad _) = head . flagToOptDescr MergeSuffix attr $ f
     extractArgDesc (ReqArg _  ad) =
         getattr (requiredFlagStr f) fromArgHelp (attr f)
     extractArgDesc _ = ""
     fromArgHelp (ArgHelp x) = Just x
     fromArgHelp _           = Nothing

  -- | Turn ["a","b","c"] into "[a [b [c]]]"
  nestedOptional [] = ""
  nestedOptional (x : xs) = "[" ++ x ++ rest ++ "]" where
    rest = case nestedOptional xs of
      "" -> ""
      r  -> ' ' : r

  orderByPosition = sortBy cmpPos where
    cmpPos = comparing getPos
    getPos = fromJust . positional . attr

requiredFlagStr :: FlagType flag => flag -> String
requiredFlagStr = underscorate . keyToStr' . flag_attrkey where
  keyToStr' (KeyF _ s) = s
  keyToStr' (KeyC c) = show c

  underscorate (x:xs) | isUpper x = '_' : x : underscorate xs
                      | otherwise = toUpper x : underscorate xs
  underscorate [] = []

optDescr :: FlagType flag => (flag -> [Attribute]) -> [OptDescr (Folded flag -> Folded flag)]
optDescr attr = concat [ one x | x <- flag_list, enabled $ attr x ]
  where one x = case flag_type x (attr x) of
          BooleanOption -> [ Option [] (invlongs $ attr x) (NoArg $ flag_set x False) ""
                           , Option (shorts $ attr x) (longs $ attr x) (flag_args x (attr x)) ""]
          _ | nonoption attr x -> []
          _ -> [Option (shorts $ attr x) (longs $ attr x) (flag_args x (attr x)) ""]

-- | Controls how to display boolean options in help output: MergePrefix shows
-- --[no-]foo, MergeSuffix shows --foo[=yes|no] and NoMerge shows two lines,
-- first with --foo and help, second with --no-foo (and no further help).
data HelpStyle = MergePrefix | MergeSuffix | NoMerge deriving Eq

flagToOptDescr :: FlagType flag => HelpStyle -> (flag -> [Attribute]) -> flag -> [OptDescr ()]
flagToOptDescr style attr x | BooleanOption <- flag_type x (attr x), NoMerge <- style =
                                [ Option (shorts $ attr x) (longs $ attr x) (args x) (help x)
                                , Option [] (invlongs $ attr x) (NoArg ()) "" ]
                            | BooleanOption <- flag_type x (attr x), MergePrefix <- style =
                                [ Option (shorts $ attr x) (ziplongs (longs $ attr x) (invlongs $ attr x))
                                         (NoArg ()) (help x) ]
                            | otherwise = [Option (shorts $ attr x) (longs $ attr x) (args x) (help x)] where
  args x = case flag_type x (attr x) of
      BooleanOption | MergeSuffix <- style -> OptArg (const ()) "yes|no"
                    | otherwise -> NoArg ()
      OptionalArgument -> OptArg (const ()) (arghelp $ attr x)
      RequiredArgument -> ReqArg (const ()) (arghelp $ attr x)
      SimpleOption -> NoArg ()
  ziplongs [] [] = []
  ziplongs (p:ps) ns = merge p (takeWhile (p `isSuffixOf`) ns)
                         : ziplongs ps (dropWhile (p `isSuffixOf`) ns)
  help x | BooleanOption <- flag_type x (attr x) = summary ++ " (default: " ++ booldef ++ ")"
         | defvalue (attr x) /= "" && flag_type x (attr x) /= SimpleOption =
             summary ++ " (default: " ++ defvalue (attr x) ++ ")"
         | otherwise = summary
    where summary = helpattr $ attr x
          booldef = if defvalue $ attr x then "yes" else "no"
  merge p ns =
    "[" ++ intercalate "/" [ take (length n - length p) n | n <- ns ] ++ "]" ++ p


helpDescr :: FlagType flag => HelpStyle -> (flag -> [Attribute]) -> [(String, [OptDescr ()])]
helpDescr style attr = requireds : [ grp g | g <- nub $ map getgroup
                         [ attr f | f <- flag_list, enabled $ attr f ] ]
  where requireds = ("Required Flags", requiredOpts)
        requiredOpts = concat
          [ flagToOptDescr style attr x | x <- flag_list,
                                          not $ nonoption attr x,
                                          isRequired attr x,
                                          enabled $ attr x ]
        grp g = (g, grpOptions g)
        grpOptions g = concat
          [ flagToOptDescr style attr x | x <- flag_list,
                                          not $ nonoption attr x,
                                          not $ isRequired attr x,
                                          enabled $ attr x,
                                          getgroup (attr x) == g ]

-- | The default parser for option arguments. Handles strings, string lists
-- (always produces single-element list), integers, booleans (@yes|true|1@ vs
-- @no|false|0@), PathF and integer lists (@--foo=1,2,3@).
readCommon :: (Data a) => String -> a
readCommon = (\x -> error $ "readflag: " ++ gshow x)
             <+< string <+< optional string
             <+< int <+< optional int
             <+< bool <+< optional bool
             <+< path <+< optional path
             <+< strlist <+< optional strlist
             <+< list int <+< optional (list int)
    where string = id
          int = read :: String -> Int
          bool b | b `elem` ["yes", "true", "1", ""] = True
                 | b `elem` ["no", "false", "0"] = False
                 | otherwise = error $ "Boolean expected in place of '" ++ b ++ "'"
          path = PathF
          strlist = (:[])
          list w str = map w (splitOn "," str)

(<+<) :: (Typeable a, Typeable b, Monad m) => m a -> m b -> m a
(<+<) = extR
infixl 8 <+<

optional _ "" = Nothing
optional f x = Just $ f x

hyphenate (x:xs) | isUpper x = '-' : toLower x : hyphenate xs
                 | otherwise = x : hyphenate xs
hyphenate [] = []

-- | Extract a long option name from an ADT constructor using its name. For
-- @Foo@, you will get @foo@ and on @FooBar@ will get @foo-bar@.
nameFromConstr :: Constr -> String
nameFromConstr ctor = map toLower (take 1 name) ++ hyphenate (drop 1 name)
  where name = showConstr ctor