{-# LANGUAGE DeriveDataTypeable, MultiParamTypeClasses, ScopedTypeVariables #-}
{-# LANGUAGE FunctionalDependencies, GADTs, FlexibleInstances, UndecidableInstances #-}
{-# LANGUAGE OverlappingInstances, Rank2Types #-}

module CmdLib (
  -- * Handling commandline options.
  FlagType(..), optsFromADT, fallback, (<+<), optional, defaultLongs

  -- * Handling subcommands.
  , Command(..), Group(..), Combine(..), Commands(..)
  , dispatch, helpCommands, execute, help

  -- * Convenience functions.
  , accept, (<|>), die

  -- * Convenience re-exports.
  , Data, Typeable
  ) where

import System.Console.GetOpt
import Control.Monad
import Control.Monad.Instances
import Data.Typeable
import Data.Data
import Data.Char( toLower, isUpper )
import Data.Maybe
import Data.Generics.Aliases
import Data.Generics.Text
import System.Exit
import System.IO

-- | A class that annotates an ADT with extra information required or useful
-- for constructing a commandline processor. Each of the ADT constructors
-- should be given a one-line description and possibly a list of short
-- (single-letter) option names that activate it. The ADT may look like this:
--
-- > newtype Path = Path String
-- > data Flag = Wibblify Path | UnBarify Int | Fooize (Maybe Int)
--
-- Moreover, a single per-ADT "readFlag" needs to be provided. This is a
-- generic read for the *contents* (arguments) of the flags. The easiest way to
-- construct it is to use "fallback" and @"<+<"@ from this module, e.g.:
--
-- > readFlag _ = fallback <+< Path <+< (read :: String -> Int)
-- >                       <+< optional (read :: String -> Int)
--
-- Usually, you will be only using few types of flag arguments, but possibly
-- many flags. You can also use the (default) @fallback <+< id <+< optional id@
-- if you prefer to keep your options as Strings (and Maybe Strings, for
-- optional arguments).
--
-- Then, "optsFromADT" will let yo uconvert this annotated ADT into a
-- @getOpt@-compatible description of options, with nullary constructors taking
-- role of @NoArg@ options, while unary ones of @ReqArg@, or @OptArg@ when the
-- parameter is Maybe X.
--
-- NB. When implementing "describe", "longs" and "shorts", make sure you only
-- ever look at the *head* (the constructor) of the flag value, never at its
-- argument (i.e. use matches like @describe (Wibblify _) = \"something
-- clever\"@. The methods are usually called on values like @Wibblify
-- undefined@.
class (Typeable flag, Data flag) => FlagType flag where
  -- | From a flag value, obtain its one-line description.
  describe :: flag -> String
  describe _ = "(no description available)"

  -- | From a flag value, obtain its long option name(s). (Reasonable default
  -- available by default, only override this if you have special
  -- requirements. See "defaultLongs".)
  longs :: flag -> [String]
  longs = defaultLongs

  -- | From a flag value, obtain its short option name(s).
  shorts :: flag -> [Char]
  shorts _ = []

  -- | De-serialize an option argument. See "FlagType" description for details.
  readFlag :: Data a => flag -> String -> a
  readFlag _ = fallback <+< id <+< optional id
  -- describeFlag :: Data a => flag -> a -> String

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

-- | A description of a single option group. Comes with a name (title) and a
-- filter for deciding which options belong to the group. The filter can be
-- built, preferably, using "accept" and "any" or @"\<|\>"@.
data Group flag where
  Group :: (FlagType flag) => String -> (flag -> Bool) -> Group flag

groupname (Group n _) = n

-- | It is often useful to process all options into a single record or similar
-- structure, that can then be used to simply pick up the values relevant for
-- the command, instead of using a list, which could contain duplicates that
-- need extra care. This class provides an interface that allows the list of
-- flags to be folded into a single value. See "execute" for details.
class (Data setup) => Combine setup flag | setup -> flag where
  -- | Take a single flag and merge it into the resulting @setup@.
  combine :: flag -> setup -> setup

  -- | The initial (default) value for the fold. If not implemented, the first
  -- constructor with all fields undefined will be used. Works well if you use
  -- records and provide defaults via a [flag] list in "execute" or "dispatch".
  initial :: setup
  initial = fromConstr $ head $ dataTypeConstrs $ dataTypeOf (undefined :: setup) :: setup

instance Combine () flag where
  combine _ _ = ()

instance (Data flag) => Combine [flag] flag where
  combine = (:)

-- | A class that describes a single (sub)command. The @cmd@ type parameter is
-- just for dispatch (and the default command name is derived from this type's
-- name, but this can be overriden). It could be an empty data decl as far as
-- this library is concerned, although you may choose to store information in
-- it.
--
-- (TODO) Nothing should prevent an ADT (record) that reflects this class to be
-- implemented and made instance of the class, for reasonable programmatic
-- manipulation.
--
-- To parse the commandline for a given
class (Typeable cmd, FlagType flag, Data setup, Combine setup flag) =>
      Command cmd flag setup | cmd -> flag setup where
  -- | List of option groups that this command takes.
  options :: cmd -> [Group flag]
  options _ = []

  -- | Set this to True if the command is a supercommand (i.e. expects another
  -- subcommand). Defaults to False. Supercommands can come with their own
  -- options, which need to appear between the supercommand and its
  -- subcommand. Any later options go to the subcommand. The "run" (and
  -- "description") method of a supercommand should use "dispatch" and
  -- "helpCommands" respectively (on its list of subcommands) itself.
  supercommand :: cmd -> Bool
  supercommand _ = False

  -- | The handler that actually runs the command. Gets the @setup@ value as
  -- folded from the processed options (see "Combine") and a list of non-option
  -- arguments.
  run :: cmd -> setup -> [String] -> IO ()

  -- | Provides the commands' short synopsis.
  synopsis :: cmd -> String
  synopsis _ = "(no synopsis available)"

  -- | Provides a short (one-line) description of the command. Used in help
  -- output.
  description :: cmd -> String
  description _ = "(no description available)"

  -- | The name of the command. Normally derived automatically from @cmd@, but
  -- may be overriden.
  cmdname :: cmd -> String
  cmdname c = map toLower $ reverse . takeWhile (/= '.') . reverse . show $ typeOf c

  -- | A convenience "undefined" of the command, for use with "Commands".
  cmd :: cmd
  cmd = undefined

-- | A (heterogenous) list of commands, with grouping. Typical use:
--
-- > commands = CmdGroup "Basic commands" (
-- >                (cmd :: Wibblify) :&:
-- >                (cmd :: Fooize) :&: EndGroup
-- >            ) :&&:
-- >            CmdGroup "Advanced commands" (
-- >                (cmd :: UnBarify) :&:
-- >                (cmd :: Poke) :&: EndGroup
-- >            ) :&&: EndCommands
--
-- Serves as input to "dispatch" and "helpCommands".
data Commands c f x where
  (:&:) :: (Command c f a) => c -> Commands c' f x' -> Commands c f x
  (:&&:) :: Commands c' f x' -> Commands c'' f x'' -> Commands c f x
  CmdGroup :: String -> Commands c' f x' -> Commands c f x
  EndGroup :: Commands c f x
  EndCommands :: Commands c f x
 deriving Typeable

infixr :&:
infixl :&&:

-- | Process an ADT into a list of one "OptDescr" per constructor, as used by
-- getOpt. Also takes an acceptance filter, which decides which constructors to
-- include in the list.
optsFromADT :: forall a. (FlagType a) => (a -> Bool) -> [OptDescr a]
optsFromADT accept = [ one x | x <- dataTypeConstrs $ dataTypeOf (undefined :: a)
                             , accept (fromConstr x) ]
  where one x = Option (shorts bit) (longs bit) args (describe bit)
          where bit = fromConstr x
                args = case (gmapQ dataTypeOf bit, gmapQ typeOf bit) of
                    ([a], [t]) | typeRepTyCon t == typeRepTyCon (typeOf (Just ())) ->
                      OptArg (\mstr -> case mstr of Just str -> reifyFlag str
                                                    Nothing -> reifyFlag "") "X"
                    ([a], _) -> ReqArg reifyFlag "X"
                    ([], []) -> NoArg bit
                    _ -> error "Can't handle multi-argument options."
                reifyFlag y = fromConstrB (readFlag (undefined :: a) y) x

-- | A simple getOpt-like function that uses "optsFromADT" to feed an
-- (annotated) ADT into getOpt. Useful for simple programs without subcommands.
parseOpt :: forall a. (FlagType a) => [String] -> ([a], [String], [String])
parseOpt opts = getOpt Permute (optsFromADT $ const True) opts

-- | A "optsFromADT"-based alternative to GetOpt's "usageInfo".
usage :: forall a. (FlagType a) => (a -> Bool) -> String -> String
usage acc str = usageInfo str $ optsFromADT acc

class (Typeable x) => GetConstr x where
  getConstr :: x -> Constr

instance (Data x) => GetConstr x where
  getConstr = toConstr

instance (Typeable a, Data x) => GetConstr (a -> x) where
  getConstr = toConstr . ($ undefined)

-- | Make a filter from a constructor. @accept Wibblify@ will give True on any
-- values of the form (Wibblify _) and False otherwise.
accept :: (GetConstr x, Data flag) => x -> (flag -> Bool)
accept ctor x = getConstr ctor == toConstr x

-- | Just a binary "any". Convenient when writing acceptance filters for "Group"s.
(<|>) :: (a -> Bool) -> (a -> Bool) -> a -> Bool
(f <|> g) a = f a || g a

acceptable :: Group a -> (a -> Bool)
acceptable (Group _ f) = f

-- | Parse options for and execute a single command (see "Command"). May be
-- useful for programs that do not need command-based "dispatch", but still
-- make use of the "Command" class to describe themselves. Handles @--help@
-- internally.
execute :: forall cmd a f. Command cmd f a => cmd -> [String] -> [f] -> IO ()
execute cmd opts defaults
  | "--help" `elem` opts = printHelp cmd
  | otherwise = do when (not $ null errs) $ die (concat errs)
                   run cmd (foldr combine initial (flags ++ defaults)) opts'
  where getopts = optsFromADT (foldr (<|>) (const False) (map acceptable $ options cmd))
        order = if supercommand cmd then RequireOrder else Permute
        (flags, opts', errs) = getOpt order getopts opts

-- | Format a help string for a single command.
help :: forall cmd a f. Command cmd f a => cmd -> String
help cmd = unlines $ ("Usage: " ++ synopsis cmd) : description cmd : "" :
           [ usageInfo (groupname group) (optsFromADT $ acceptable group)
           | group <- options cmd ]

-- | Helper for dying with an error message (nicely, at least compared to
-- "fail" in IO).
die :: String -> IO a
die msg = do hPutStrLn stderr ("FATAL: " ++ trim msg)
             exitWith (ExitFailure 1)
             return (error "unreachable")
  where trim msg | last msg == '\n' = trim $ init msg
                 | otherwise = msg

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

fallback x = error $ "readflag: " ++ gshow x

-- | From "Commands" (a list of commands), produce a help text describing each
-- of the commands, sorted into groups.
helpCommands :: forall c f x. Commands c f x -> String
helpCommands (CmdGroup name comms) = name ++ ":\n" ++ helpCommands comms
helpCommands (comm :&: comms) = "  " ++ pad (cmdname comm) ++ description comm ++ "\n"
                                ++ helpCommands comms
  where pad str = (take 15 $ str ++ replicate 15 ' ') ++ " "
helpCommands (comms :&&: comms') = helpCommands comms ++ "\n" ++ helpCommands comms'
helpCommands EndGroup = ""
helpCommands EndCommands = ""

newtype Imp f b = Imp (forall d y. (Command d f y) => d -> b) -- impredicative

-- Admittedly, this is a little scary...
find :: forall c f x a b. Commands c f x -> String -> (a, Imp f b) -> Either a b
find (CmdGroup _ comms) cmd act = find comms cmd act
find (comm :&: comms) cmd (notfound, Imp found)
  | cmd == cmdname comm = Right $ found comm
  | otherwise = find comms cmd (notfound, Imp found)
find (comms :&&: comms') cmd act = case find comms cmd act of
  Left _ -> find comms' cmd act
  Right res -> Right res
find EndGroup _ (notfound, _) = Left notfound
find EndCommands _ (notfound, _) = Left notfound

printHelp c = putStr $ help c
printCommands comms = putStr $ helpCommands comms

joinEither :: Either a a -> a
joinEither (Left x) = x
joinEither (Right x) = x

-- | Given a list of commands (see "Commands"), a list of commandline arguments
-- and a list of default flags, dispatch on the command name, parse the
-- commandline options (see "execute") and transfer control to the command.
-- This function also implements the @help@ pseudocommand.
dispatch :: forall c f x. Commands c f x -> [String] -> [f] -> IO ()
dispatch comms [] _ = printCommands comms >> die "Command required."
dispatch comms ["help"] _ = printCommands comms
dispatch comms ("help":cmd:_) _ = joinEither $ find comms cmd (notfound, Imp printHelp)
  where notfound = printCommands comms >> die ("No such command " ++ cmd)
dispatch comms (cmd:opts) def = joinEither $ find comms cmd (notfound, Imp found)
  where notfound = printCommands comms >> die ("No such command " ++ cmd)
        found :: forall c a. (Command c f a) => c -> IO ()
        found x = execute x opts def

-- | Extract default long option name from an ADT constructor using its
-- name. For @data Foo = Bar | FooBar String@, calling @defaultLongs@ on @Bar@
-- will give @bar@ and on @FooBar undefined@ will give @foo-bar@.
defaultLongs t = [name]
  where ctr = show $ toConstr t
        name = map toLower (take 1 ctr) ++ hyphen (drop 1 ctr)
        hyphen (x:xs) | isUpper x = '-' : toLower x : hyphen xs
                      | otherwise = x : hyphen xs
        hyphen [] = []