    {-# LANGUAGE PatternGuards, ScopedTypeVariables, ExistentialQuantification, DeriveDataTypeable #-}

    -- | This module captures annotations on a value, and builds a 'Capture' value.
    --   This module has two ways of writing annotations:
    --
    --   /Impure/: The impure method of writing annotations is susceptible to over-optimisation by GHC
    --   - sometimes @\{\-\# OPTIONS_GHC -fno-cse \#\-\}@ will be required.
    --
    --   /Pure/: The pure method is more verbose, and lacks some type safety.
    --
    --   As an example of the two styles:
    --
    -- > data Foo = Foo {foo :: Int, bar :: Int}
    --
    --   @ impure = 'capture' $ Foo {foo = 12, bar = 'many' [1 '&=' \"inner\", 2]} '&=' \"top\"@
    --
    --   @ pure = 'capture_' $ 'record' Foo{} [foo := 12, bar :=+ ['atom' 1 '+=' \"inner\", 'atom' 2]] '+=' \"top\"@
    --
    --   Both evaluate to:
    --
    -- > Capture (Ann "top") (Ctor (Foo 12 1) [Value 12, Many [Ann "inner" (Value 1), Value 2]]
    module System.Console.CmdArgs.Annotate(
        -- * Capture framework
        Capture(..), Any(..), fromCapture, defaultMissing,
        -- * Impure
        capture, many, (&=),
        -- * Pure
        capture_, many_, (+=), atom, record, Annotate((:=),(:=+))
        ) where

    import Control.Monad
    import Control.Monad.Trans.State
    import Data.Data(Data,Typeable)
    import Data.List
    import Data.Maybe
    import Data.IORef
    import System.IO.Unsafe
    import Control.Exception
    import Data.Generics.Any

    infixl 2 &=, +=
    infix 3 :=

    -- | The result of capturing some annotations.
    data Capture ann
        = Many [Capture ann] -- ^ Many values collapsed ('many' or 'many_')
        | Ann ann (Capture ann) -- ^ An annotation attached to a value ('&=' or '+=')
        | Value Any -- ^ A value (just a value, or 'atom')
        | Missing Any -- ^ A missing field (a 'RecConError' exception, or missing from 'record')
        | Ctor Any [Capture ann] -- ^ A constructor (a constructor, or 'record')
          deriving Show

    instance Functor Capture where
        fmap f (Many xs) = Many $ map (fmap f) xs
        fmap f (Ann a x) = Ann (f a) $ fmap f x
        fmap f (Value x) = Value x
        fmap f (Missing x) = Missing x
        fmap f (Ctor x xs) = Ctor x $ map (fmap f) xs

    -- | Return the value inside a capture.
    fromCapture :: Capture ann -> Any
    fromCapture (Many (x:_)) = fromCapture x
    fromCapture (Ann _ x) = fromCapture x
    fromCapture (Value x) = x
    fromCapture (Missing x) = x
    fromCapture (Ctor x _) = x

    -- | Remove all Missing values by using any previous instances as default values
    defaultMissing :: Capture ann -> Capture ann
    defaultMissing x = evalState (f Nothing Nothing x) []
        where
            f ctor field (Many xs) = fmap Many $ mapM (f ctor field) xs
            f ctor field (Ann a x) = fmap (Ann a) $ f ctor field x
            f ctor field (Value x) = return $ Value x
            f (Just ctor) (Just field) (Missing x) = do
                s <- get
                return $ head $
                    [x2 | (ctor2,field2,x2) <- s, typeOf ctor == typeOf ctor2, field == field2] ++
                    err ("missing value encountered, no field for " ++ field ++ " (of type " ++ show x ++ ")")
            f _ _ (Missing x) = err $ "missing value encountered, but not as a field (of type " ++ show x ++ ")"
            f _ _ (Ctor x xs) | length (fields x) == length xs = do
                ys <- zipWithM (g x) (fields x) xs
                return $ Ctor (recompose x $ map fromCapture ys) ys
            f _ _ (Ctor x xs) = fmap (Ctor x) $ mapM (f Nothing Nothing) xs

            g ctor field x = do
                y <- f (Just ctor) (Just field) x
                modify ((ctor,field,y):)
                return y

            err x = error $ "System.Console.CmdArgs.Annotate.defaultMissing, " ++ x

    ---------------------------------------------------------------------
    -- IMPURE BIT

    -- test = show $ capture $ many [Just ((66::Int) &= P 1 &= P 2), Nothing &= P 8] &= P 3

    {-
    Notes On Purity
    ---------------

    There is a risk that things that are unsafe will be inlined. That can generally be
    removed by NOININE on everything.

    There is also a risk that things get commoned up. For example:

    foo = trace "1" 1
    bar = trace "1" 1
    main = do
        evaluate foo
        evaluate bar

    Will print "1" only once, since foo and bar share the same pattern. However, if
    anything in the value is a lambda they are not seen as equal. We exploit this by
    defining const_ and id_ as per this module.

    Now anything wrapped in id_ looks different from anything else.
    -}

    {-
    The idea is to keep a stack of either continuations, or values
    If you encounter 'many' you become a value
    If you encounter '&=' you increase the continuation
    -}

    {-# NOINLINE ref #-}
    ref :: IORef [Either (Capture Any -> Capture Any) (Capture Any)]
    ref = unsafePerformIO $ newIORef []

    push = modifyIORef ref (Left id :)
    pop = do x:xs <- readIORef ref; writeIORef ref xs; return x
    change f = modifyIORef ref $ \x -> case x of Left g : rest -> f g : rest ; _ -> error "Internal error in Capture"
    add f = change $ \x -> Left $ x . f
    set x = change $ \f -> Right $ f x

    -- | Collapse multiple values in to one.
    {-# NOINLINE many #-}
    many :: Data val => [val] -> val
    many xs = unsafePerformIO $ do
        ys <- mapM (force . Any) xs
        set $ Many ys
        return $ head xs

    {-# NOINLINE addAnn #-}
    addAnn :: (Data val, Data ann) => val -> ann -> val
    addAnn x y = unsafePerformIO $ do
        add (Ann $ Any y)
        evaluate x
        return x

    -- | Capture a value. Note that if the value is evaluated
    --   more than once the result may be different, i.e.
    --
    -- > capture x /= capture x
    {-# NOINLINE capture #-}
    capture :: (Data val, Data ann) => val -> Capture ann
    capture x = unsafePerformIO $ fmap (fmap fromAny) $ force $ Any x

    force :: Any -> IO (Capture Any)
    force x@(Any xx) = do
        push
        res <- try $ evaluate xx
        y <- pop
        case y of
            _ | Left (_ :: RecConError) <- res -> return $ Missing x
            Right r -> return r
            Left f | not $ isAlgType x -> return $ f $ Value x
                   | otherwise -> do
                cs <- mapM force $ children x
                return $ f $ Ctor x cs

    -- | Add an annotation to a value.
    --
    --   It is recommended that anyone making use of this function redefine
    --   it with a more restrictive type signature to control the type of the
    --   annotation (the second argument). Any redefinitions of this function
    --   should add an INLINE pragma, to reduce the chance of incorrect
    --   optimisations.
    {-# INLINE (&=) #-}
    (&=) :: (Data val, Data ann) => val -> ann -> val
    (&=) x y = addAnn (id_ x) (id_ y)

    {-# INLINE id_ #-}
    id_ :: a -> a
    id_ x = case unit of () -> x
        where unit = reverse "" `seq` ()

    ---------------------------------------------------------------------
    -- PURE PART

    -- | This type represents an annotated value. The type of the underlying value is not specified.
    data Annotate ann
        = forall c f . (Data c, Data f) => (c -> f) := f -- ^ Construct a field, @fieldname := value@.
        | forall c f . (Data c, Data f) => (c -> f) :=+ [Annotate ann] -- ^ Add annotations to a field.
        | AAnn ann (Annotate ann)
        | AMany [Annotate ann]
        | AAtom Any
        | ACtor Any [Annotate ann]
          deriving Typeable
          -- specifically DOES NOT derive Data, to avoid people accidentally including it

    -- | Add an annotation to a value.
    (+=) :: Annotate ann -> ann -> Annotate ann
    (+=) = flip AAnn

    -- | Collapse many annotated values in to one.
    many_ :: [Annotate a] -> Annotate a
    many_ = AMany

    -- | Lift a pure value to an annotation.
    atom :: Data val => val -> Annotate ann
    atom = AAtom . Any

    -- | Create a constructor/record. The first argument should be
    --   the type of field, the second should be a list of fields constructed
    --   originally defined by @:=@ or @:=+@.
    --
    --   This operation is not type safe, and may raise an exception at runtime
    --   if any field has the wrong type or label.
    record :: Data a => a -> [Annotate ann] -> Annotate ann
    record a b = ACtor (Any a) b

    -- | Capture the annotations from an annotated value.
    capture_ :: Show a => Annotate a -> Capture a
    capture_ (AAnn a x) = Ann a (capture_ x)
    capture_ (AMany xs) = Many (map capture_ xs)
    capture_ (AAtom x) = Value x
    capture_ (_ := c) = Value $ Any c
    capture_ (_ :=+ c) = Many $ map capture_ c
    capture_ (ACtor x xs)
        | not $ null rep = error $ "Some fields got repeated under " ++ show x ++ "." ++ ctor x ++ ": " ++ show rep
        | otherwise = Ctor x2 xs2
        where
            x2 = recompose x $ map fromCapture xs2
            xs2 = [fromMaybe (Missing c) $ lookup i is | let is = zip inds $ map capture_ xs, (i,c) <- zip [0..] $ children x]
            inds = zipWith fromMaybe [0..] $ map (fieldIndex x) xs
            rep = inds \\ nub inds

    fieldIndex :: Any -> Annotate a -> Maybe Int
    fieldIndex ctor (AAnn a x) = fieldIndex ctor x
    fieldIndex ctor (f := _) = fieldIndex ctor (f :=+ [])
    fieldIndex ctor (f :=+ _) | isJust res = res
                              | otherwise = error $ "Couldn't resolve field for " ++ show ctor
        where c = recompose ctor [Any $ throwInt i `asTypeOf` x | (i,Any x) <- zip [0..] (children ctor)]
              res = catchInt $ f $ fromAny c
    fieldIndex _ _ = Nothing

    data ExceptionInt = ExceptionInt Int deriving (Show, Typeable)
    instance Exception ExceptionInt

    throwInt :: Int -> a
    throwInt i = throw (ExceptionInt i)

    {-# NOINLINE catchInt #-}
    catchInt :: a -> Maybe Int
    catchInt x = unsafePerformIO $ do
        y <- try (evaluate x)
        return $ case y of
            Left (ExceptionInt z) -> Just z
            _ -> Nothing