{-# LANGUAGE ScopedTypeVariables, TypeFamilies, MultiParamTypeClasses, DeriveDataTypeable,
             Rank2Types, FlexibleContexts, UndecidableInstances #-}
module System.Console.CmdLib.Record where
import System.Console.CmdLib.Attribute
import System.Console.CmdLib.Command
import System.Console.GetOpt
import Data.Data
import Data.List( nub, elemIndex )
import Control.Monad.State( evalState, get, put, State, execState )
import Data.Maybe( fromMaybe, fromJust )
import System.Exit

import System.Console.CmdLib.Flag

append :: forall a. (Data a) => a -> a -> a
append a = any <+< list
  where any = const a
        list :: (Typeable b) => b -> [String]
        list = (++(fromJust $ cast a)) . fromJust . cast

-- | This wrapper type allows use of record types (single or multi-constructor)
-- for handling flags. Each field of the record is made into a single flag of
-- the corresponding type. The record needs to be made an instance of the
-- "Attributes" class. That way, attributes can be attached to the field
-- selectors, although when used with RecordCommand, its "rec_options" method
-- can be used as well and the Attributes instance left empty.
--
-- > data Flags = Flags { wibblify :: Int, simplify :: Bool }
-- > instance Attributes Flags where
-- >     attributes _ =
-- >        wibblify %> Help "Add a wibblification pass." %+ ArgHelp "intensity" %%
-- >        simplify %> Help "Enable a two-pass simplifier."
--
-- A single value of the Flags type will then be passed to the "Command"
-- instances (those that use @Record Flags@ as their second type parameter),
-- containing the value of the rightmost occurence for each of the flags.
--
-- TODO: List-based option types should be accumulated instead of overriden.
data (Attributes rec) => Record rec = Record rec String deriving Eq

instance (Eq rec, Eq (Record rec), Data rec, Attributes rec) => FlagType (Record rec) where
  type Folded (Record rec) = rec

  flag_attrkey (Record _ f) = KeyF (typeOf (undefined :: rec)) f

  flag_list = [ Record undefined field |
                field <- nub $ concatMap constrFields
                         (dataTypeConstrs $ dataTypeOf (undefined :: rec)) ]

  flag_type f@(Record rec field) attr =
      fixup $ optionType [gmapQi idx dataTypeOf x] [gmapQi idx typeOf x]
    where constr = head [ ctor | ctor <- dataTypeConstrs $ dataTypeOf rec
                               , field `elem` constrFields ctor ]
          idx = fromMaybe (error $ "BUG: Getting type of nonexistent field " ++ field) $
                    elemIndex field $ constrFields constr
          x :: rec = fromConstr constr
          fixup BooleanOption | False <- invertible attr = SimpleOption
          fixup x = x

  flag_defaults attr = map setfield flag_list
    where setdef :: forall a. (Data a) => Record rec -> a
          setdef flag = defvalue $ attrs flag attr
          setfield flag@(Record _ field) rec = setField field rec (setdef flag)

  flag_attrs (Record _ name) = (attrFun $ attributes (undefined :: rec))
                               (KeyF (typeOf (undefined :: rec)) name) %+
                               long (hyphenate name) %+ defaults

  flag_value (Record _ field) folded = getField field folded
  flag_set (Record _ field) v = \x -> setField field x (errcast v)
     where errcast :: (Typeable a, Typeable b) => a -> b
           errcast x = case cast x of
                           Just x -> x
                           Nothing -> error "BUG: flag_set in Record used with wrong value type"

  flag_parse f@(Record _ field) str v = setField field v (readFlag (undefined :: rec) str)

  flag_args f@(Record flag field) attr = case flag_type (Record flag field) attr of
    BooleanOption -> OptArg setoptional ""
    OptionalArgument -> OptArg setoptional ""
    RequiredArgument -> ReqArg setlist ""
    SimpleOption -> NoArg $ flag_set f True
   where set str v = setField field v (readFlag (undefined :: rec) str)
         setlist str v = setField field v (append (readFlag (undefined :: rec) str) (flag_value f v))
         setoptional str = set (fromMaybe "" str)

  flag_empty _ = fromConstr $ head $ dataTypeConstrs $ dataTypeOf (undefined :: rec)

-- | A bridge that allows multi-constructor record types to be used as a
-- description of a command set. In such a type, each constructor corresponds
-- to a single command and its fields to its options. To describe a program
-- with two commands, @foo@ and @bar@, each taking a @--wibble@ boolean option
-- and @bar@ also taking a @--text=<string>@ option, you can write:
--
-- > data Commands = Foo { wibble :: Bool }
-- >               | Bar { wibble :: Bool, text :: String }
-- >
-- > instance RecordCommand Commands where (...)
--
-- You should at least implement @run'@, @rec_options@ and @mode_summary@ are optional.

class (Data cmd) => RecordCommand cmd where
  -- | @run'@ is your entrypoint into the whole set of commands. You can
  -- dispatch on the command by looking at the constructor in @cmd@:
  --
  -- > run' cmd@(Foo {}) _ = putStrLn $ "Foo running. Wibble = " ++ show (wibble cmd)
  -- > run' cmd@(Bar {}) _ = putStrLn "This is bar."
  run' :: cmd -> [String] -> IO ()

  -- | You can also provide extra per-command flag attributes (match on the
  -- constructor like with @run'@). The attributes shared by various commands
  -- can be set in "rec_attrs" in "Attributes" instead.
  rec_options :: cmd -> AttributeMap Key
  rec_options _ = EmptyMap

  -- | Provide a summary help string for each mode. Used in help output. Again,
  -- pattern match like in @run'@.
  mode_summary :: cmd -> String
  mode_summary _ = ""

  -- | Provide a help blurb for each mode. Use patterns like in @run'@.
  mode_help :: cmd -> String
  mode_help _ = ""

data RecordMode cmd = RecordMode { rec_cmdname :: String
                                 , rec_initial :: Constr }
                    deriving (Typeable)

instance (Eq cmd, Eq (Record cmd), RecordCommand cmd, Data cmd, Attributes cmd)
         => Command (RecordMode cmd) (Record cmd) where
  cmdname = rec_cmdname
  run _ = run'

  summary cmd = mode_summary $ (fromConstr $ rec_initial cmd :: cmd)
  help cmd = mode_help $ (fromConstr $ rec_initial cmd :: cmd)

  options cmd = rec_options ctor %% available %% everywhere disable
    where available = [ (KeyF (typeOf (undefined :: cmd)) opt, [enable])
                      | opt <- constrFields $ rec_initial cmd]
          ctor = (fromConstr $ rec_initial cmd :: cmd)

  cmd_flag_empty cmd = fromConstr $ rec_initial cmd

-- | Construct a command list (for "dispatch"/"helpCommands") from a
-- multi-constructor record data type. See also "RecordCommand". Alternatively,
-- you can use "dispatchR" directly.
recordCommands :: forall cmd. (Eq cmd, Eq (Record cmd), Data cmd, RecordCommand cmd, Attributes cmd)
               => cmd -> [CommandWrap]
recordCommands _ = go $ dataTypeConstrs $ dataTypeOf (undefined :: cmd)
    where go = map $ CommandWrap . rec_cmd
          rec_cmd :: Constr -> RecordMode cmd
          rec_cmd x = RecordMode { rec_cmdname = nameFromConstr x
                                 , rec_initial = x }

-- | Record field update using a string field name. Sets a field value in a
-- record, using a (string) name of the field,
setField :: forall rec a. (Data rec) => String -> rec -> (forall b. (Data b) => b) -> rec
setField field current value = evalState (gmapM subst current) (constrFields $ toConstr current)
    where subst :: Data x => x -> State [String] x
          subst f = do x:xs <- get
                       put xs
                       if x == field then return value else return f

data Imp = Imp Int (forall a. (Typeable a) => a)

getField :: forall rec a. (Data rec, Typeable a) => String -> rec -> a
getField field value = case execState (gmapM find value) (Imp 0 (error "")) of Imp _ v -> v
    where idx = elemIndex field $ constrFields (toConstr value)
          find :: Data x => x -> State Imp x
          find f = do Imp i val <- get
                      if (Just i == idx) then put $ Imp (i + 1) (errcast f)
                                         else put $ Imp (i + 1) val
                      return $ error "find"
          errcast :: forall a b. (Typeable a, Typeable b) => a -> b
          errcast = fromMaybe (error $ "BUG: getField used with wrong type on " ++ field) . cast

-- | A command parsing & dispatch entry point for record-based
-- commands. Ex. (see "RecordCommand"):
--
-- > main = getArgs >>= dispatchR [] >>= \x -> case x of
-- >   Foo {} -> putStrLn $ "You asked for foo. Wibble = " ++ show (wibble x)
-- >   Bar {} -> putStrLn $ "You asked for bar. ..."
dispatchR :: forall cmd f. (Eq cmd, Eq (Record cmd), Attributes cmd, RecordCommand cmd,
                            Command (RecordMode cmd) f,
                            Folded f ~ cmd) => [DispatchOpt] -> [String] -> IO cmd
dispatchR dopt opts = dispatch' die dopt (recordCommands (undefined :: cmd)) opts >>= \c -> case c of
  Nothing -> exitWith ExitSuccess >> return undefined
  Just (CommandWrap x, opts') -> execute' x opts' >>= \c -> case c of
    Just (command, opts') -> case (cast command) of
      Just comm -> return comm
      Nothing -> execute x opts' >> exitWith ExitSuccess >> return undefined
    Nothing -> exitWith ExitSuccess >> return undefined

-- | Like "execute", but you get the flags as a return value. This is useful to
-- implement non-modal applications with record-based flags, eg.:
--
-- > data Main = Main { greeting :: String, again :: Bool }
-- >     deriving (Typeable, Data, Eq)
-- > instance Attributes Main where -- (...)
-- > instance RecordCommand Main
-- > main = getArgs >>= executeR Main {} >>= \opts -> do
-- >    putStrLn (greeting opts) -- (...)
executeR :: forall cmd. (Eq cmd, Eq (Record cmd), Attributes cmd, RecordCommand cmd)
         => cmd -> [String] -> IO cmd
executeR cmd opts = execute' cmd' opts >>= \c -> case c of
  Just (command, _) -> return command
  Nothing -> exitWith ExitSuccess >> return undefined
  where cmd' = RecordMode (nameFromConstr $ toConstr cmd) (toConstr cmd)