{-# LANGUAGE TemplateHaskell #-}
module Data.Generics.Geniplate(universeBi, universeBiT, transformBi, transformBiT, transformBiM, transformBiMT) where
import Control.Monad
import Control.Exception(assert)
import Control.Monad.State.Strict
import Data.Maybe
import Language.Haskell.TH
import Language.Haskell.TH.Syntax hiding (lift)
import System.IO

-- | Generate TH code for a function that extracts all subparts of a certain type.
-- The argument to 'universeBi' is a name with the type @S -> [T]@, for some types
-- @S@ and @T@.  The function will extract all subparts of type @T@ from @S@.
universeBi :: Name -> Q Exp
universeBi = universeBiT []

-- | Same as 'universeBi', but does not look inside any types mention in the
-- list of types.
universeBiT :: [TypeQ] -> Name -> Q Exp
universeBiT stops name = do
    (_tvs, from, tos) <- getNameType name
    let to = unList tos
--    qRunIO $ print (from, to)
    (ds, f) <- uniBiQ stops from to
    x <- newName "_x"
    let e = LamE [VarP x] $ LetE ds $ AppE (AppE f (VarE x)) (ListE [])
--    qRunIO $ do putStrLn $ pprint e; hFlush stdout
    return e

type U = StateT (Integer, Map Type [Dec], Map Type Bool) Q

newNameU :: String -> U Name
newNameU s = do
    (n, m, c) <- get
    put (n+1, m, c)
    lift $ newName $ s ++ "_" ++ show n

uniBiQ :: [TypeQ] -> Type -> Type -> Q ([Dec], Exp)
uniBiQ stops from ato = do
    ss <- sequence stops
    to <- expandSyn ato
    (f, (_, m, _)) <- runStateT (uniBi from to) (0, mEmpty, mFromList $ zip ss (repeat False))
    return (concat $ mElems m, f)

uniBi :: Type -> Type -> U Exp
uniBi afrom to = do
    (n, m, c) <- get
    from <- lift $ expandSyn afrom
    case mLookup from m of
        Just (FunD n _ : _) -> return $ VarE n
        _ -> do
            f <- newNameU "_f"
            let mkRec = do
                    put (n, mInsert from [FunD f [Clause [] (NormalB $ TupE []) []]] m, c)   -- insert something to break recursion, will be replaced below.
                    uniBiCase from to
            cs <- if from == to then do
                      b <- contains' to from
                      if b then do
                          -- Recursive data type, we need the current value and all values inside.
                          g <- newNameU "_g"
                          gcs <- mkRec
                          let dg = [FunD g gcs]
                          -- Insert with a dummy type, just to get the definition in the map for mElems.
                          modify $ \ (n', m', c') -> (n', mInsert (ConT g) dg m', c')
                          lift $ fmap unFunD [d| f _x _r = _x : $(return (VarE g)) _x _r |]
                       else
                          -- Non-recursive type, just use this value.
                          lift $ fmap unFunD [d| f _x _r = _x : _r |]
                  else do
                      -- Types differ, look inside.
                      b <- contains to from
                      if b then do
                          -- Occurrences inside, recurse.
                          mkRec
                       else
                          -- No occurrences of to inside from, so add nothing.
                          lift $ fmap unFunD [d| f _ _r = _r |]
            let d = [FunD f cs]
            modify $ \ (n', m', c') -> (n', mInsert from d m', c')
            return $ VarE f

-- Check if the second type is contained anywhere in the first type.
contains :: Type -> Type -> U Bool
contains to afrom = do
--    lift $ qRunIO $ print ("contains", to, from)
    from <- lift $ expandSyn afrom
    if from == to then
        return True
     else do
        c <- gets (\ (_,_,c) -> c)
        case mLookup from c of
            Just b  -> return b
            Nothing -> contains' to from

-- Check if the second type is contained somewhere inside the first.
contains' :: Type -> Type -> U Bool
contains' to from = do
--    lift $ qRunIO $ print ("contains'", to, from)
    let (con, ts) = splitTypeApp from
    modify $ \ (n, m, c) -> (n, m, mInsert from False c)        -- To make the fixpoint of the recursion False.
    b <- case con of
         ConT n    -> containsCon n to ts
         TupleT _  -> fmap or $ mapM (contains to) ts
         ArrowT    -> return False
         ListT     -> contains to (head ts)
         VarT _    -> return False
         t         -> genError $ "contains: unexpected type: " ++ pprint from ++ " (" ++ show t ++ ")"
    modify $ \ (n, m, c) -> (n, m, mInsert from b c)
    return b

containsCon :: Name -> Type -> [Type] -> U Bool
containsCon con to ts = do
--    lift $ qRunIO $ print ("containsCon", con, to, ts)
    (tvs, cons) <- lift $ getTyConInfo con
    let conCon (NormalC _ xs) = fmap or $ mapM (field . snd) xs
        conCon (InfixC x1 _ x2) = fmap or $ mapM field [snd x1, snd x2]
        conCon (RecC _ xs) = fmap or $ mapM field [ t | (_,_,t) <- xs ]
        conCon c = genError $ "containsCon: " ++ show c
        s = mkSubst tvs ts
        field t = contains to (subst s t)
    fmap or $ mapM conCon cons

unFunD :: [Dec] -> [Clause]
unFunD [FunD _ cs] = cs
unFunD _ = genError $ "unFunD"

uniBiCase :: Type -> Type -> U [Clause]
uniBiCase from to = do
    let (con, ts) = splitTypeApp from
    case con of
        ConT n    -> uniBiCon n ts to
        TupleT _  -> uniBiTuple ts to
--        ArrowT    -> lift $ fmap unFunD [d| f _ _r = _r |]           -- Stop at functions
        ListT     -> uniBiList (head ts) to
        t         -> genError $ "uniBiCase: unexpected type: " ++ pprint from ++ " (" ++ show t ++ ")"

uniBiList :: Type -> Type -> U [Clause]
uniBiList t to = do
    uni <- uniBi t to
    rec <- uniBi (AppT ListT t) to
    lift $ fmap unFunD [d| f [] _r = _r; f (_x:_xs) _r = $(return uni) _x ($(return rec) _xs _r) |]

uniBiTuple :: [Type] -> Type -> U [Clause]
uniBiTuple ts to = fmap (:[]) $ mkArm to [] TupP ts

uniBiCon :: Name -> [Type] -> Type -> U [Clause]
uniBiCon con ts to = do
    (tvs, cons) <- lift $ getTyConInfo con
    let genArm (NormalC c xs) = arm (ConP c) xs
        genArm (InfixC x1 c x2) = arm (\ [p1, p2] -> InfixP p1 c p2) [x1, x2]
        genArm (RecC c xs) = arm (ConP c) [ (b,t) | (_,b,t) <- xs ]
        genArm c = genError $ "uniBiCon: " ++ show c
        s = mkSubst tvs ts
        arm c xs = mkArm to s c $ map snd xs

    if null cons then
        -- No constructurs, return nothing
        lift $ fmap unFunD [d| f _ _r = _r |]
     else
        mapM genArm cons

mkArm :: Type -> Subst -> ([Pat] -> Pat) -> [Type] -> U Clause
mkArm to s c ts = do
    r <- newNameU "_r"
    vs <- mapM (const $ newNameU "_x") ts
    let sub v t = do
            let t' = subst s t
            uni <- uniBi t' to
            return $ AppE (AppE uni (VarE v))
    es <- zipWithM sub vs ts
    let body = foldr ($) (VarE r) es
    return $ Clause [c (map VarP vs), VarP r] (NormalB body) []


type Subst = [(Name, Type)]

mkSubst :: [TyVarBndr] -> [Type] -> Subst
mkSubst vs ts =
   let vs' = map un vs
       un (PlainTV v) = v
       un (KindedTV v _) = v
   in  assert (length vs' == length ts) $ zip vs' ts

subst :: Subst -> Type -> Type
subst s (ForallT v c t) = ForallT v c $ subst s t
subst s t@(VarT n) = fromMaybe t $ lookup n s
subst s (AppT t1 t2) = AppT (subst s t1) (subst s t2)
subst s (SigT t k) = SigT (subst s t) k
subst _ t = t

getTyConInfo :: Name -> Q ([TyVarBndr], [Con])
getTyConInfo con = do
    info <- qReify con
    case info of
        TyConI (DataD _ _ tvs cs _) -> return (tvs, cs)
        TyConI (NewtypeD _ _ tvs c _) -> return (tvs, [c])
        PrimTyConI{} -> return ([], [])
        i -> genError $ "unexpected TyCon: " ++ show i

getNameType :: Name -> Q ([TyVarBndr], Type, Type)
getNameType name = do
    info <- qReify name
    let split (ForallT tvs _ t) = (tvs ++ tvs', from, to) where (tvs', from, to) = split t
        split (AppT (AppT ArrowT from) to) = ([], from, to)
        split t = genError $ "Type is not an arrow: " ++ pprint t
    case info of
        VarI _ t _ _ -> return $ split t
        _            -> genError $ "Name is not variable: " ++ pprint name

unList :: Type -> Type
unList (AppT (ConT n) t) | n == ''[] = t
unList (AppT ListT t) = t
unList t = genError $ "universeBi: Type is not a list: " ++ pprint t -- ++ " (" ++ show t ++ ")"

splitTypeApp :: Type -> (Type, [Type])
splitTypeApp (AppT a r) = (c, rs ++ [r]) where (c, rs) = splitTypeApp a
splitTypeApp t = (t, [])

expandSyn :: Type -> Q Type
expandSyn (ForallT tvs ctx t) = liftM (ForallT tvs ctx) $ expandSyn t
expandSyn t@AppT{} = expandSynApp t []
expandSyn t@ConT{} = expandSynApp t []
expandSyn (SigT t k) = liftM (flip SigT k) $ expandSyn t
expandSyn t = return t

expandSynApp :: Type -> [Type] -> Q Type
expandSynApp (AppT t1 t2) ts = do t2' <- expandSyn t2; expandSynApp t1 (t2':ts)
expandSynApp t@(ConT n) ts = do
    info <- qReify n
    case info of
        TyConI (TySynD _ tvs rhs) ->
            let (ts', ts'') = splitAt (length tvs) ts
                s = mkSubst tvs ts'
                rhs' = subst s rhs
            in  expandSynApp rhs' ts''
        _ -> return $ foldl AppT t ts
expandSynApp t ts = do t' <- expandSyn t; return $ foldl AppT t' ts


genError :: String -> a
genError msg = error $ "Data.Generics.Geniplate: " ++ msg

----------------------------------------------------

-- Exp has type (S -> S) -> T -> T, for some S and T
-- | Generate TH code for a function that transforms all subparts of a certain type.
-- The argument to 'transformBi' is a name with the type @(S->S) -> T -> T@, for some types
-- @S@ and @T@.  The function will transform all subparts of type @S@ inside @T@ using the given function.
transformBi :: Name -> Q Exp
transformBi = transformBiT []

-- | Same as 'transformBi', but does not look inside any types mention in the
-- list of types.
transformBiT :: [TypeQ] -> Name -> Q Exp
transformBiT = transformBiG (Nothing, id, AppE, AppE)

transformBiM :: Name -> Q Exp
transformBiM = transformBiMT []

transformBiMT :: [TypeQ] -> Name -> Q Exp
transformBiMT = transformBiG (Just undefined, eret, eap, emap)
  where eret e = AppE (VarE 'Control.Monad.return) e
        eap f a = AppE (AppE (VarE 'Control.Monad.ap) f) a
        emap f a = AppE (AppE (VarE '(Control.Monad.=<<)) f) a

type RetAp = (Maybe Type, Exp -> Exp, Exp -> Exp -> Exp, Exp -> Exp -> Exp)

transformBiG :: RetAp -> [TypeQ] -> Name -> Q Exp
transformBiG ra stops name = do
    (_tvs, fcn, res) <- getNameType name
    f <- newName "_f"
    x <- newName "_x"
    (ds, tr, fty, argty) <-
        case (ra, fcn, res) of
            ((Nothing,_,_,_),     AppT (AppT ArrowT s) s',          AppT (AppT ArrowT t) t') | s == s' && t == t' -> trBiQ ra stops f s t
            ((Just _,ret, apl, rbn), AppT (AppT ArrowT s) (AppT m s'), AppT (AppT ArrowT t) (AppT m' t')) |
                      s == s' && t == t' && m == m' -> trBiQ (Just m, ret, apl, rbn) stops f s t
            _ -> genError $ "transformBi: malformed type: " ++ pprint (AppT (AppT ArrowT fcn) res) ++ ", should have form (S->S) -> (T->T)"
    let e = LamE [VarP f {-`SigP` fty-}, VarP x {-`SigP` argty-}] $ LetE ds $ AppE tr (VarE x)
--    qRunIO $ do putStrLn $ pprint e; hFlush stdout
    return e

trBiQ :: RetAp -> [TypeQ] -> Name -> Type -> Type -> Q ([Dec], Exp, Type, Type)
trBiQ ra@(mty,_,_,_) stops f aft st = do
    ss <- sequence stops
    ft <- expandSyn aft
    (tr, (_, m, _)) <- runStateT (trBi ra (VarE f) ft st) (0, mEmpty, mFromList $ zip ss (repeat False))
    let fty = ft `arrow` maybe ft (flip AppT ft) mty
    return (concat $ mElems m, tr, fty, st)

arrow :: Type -> Type -> Type
arrow t1 t2 = AppT (AppT ArrowT t1) t2

trBi :: RetAp -> Exp -> Type -> Type -> U Exp
trBi ra@(mty, ret, _, rbind) f ft ast = do
    (n, m, c) <- get
    st <- lift $ expandSyn ast
--    lift $ qRunIO $ print (ft, st)
    case mLookup st m of
        Just (FunD n _ : _) -> return $ VarE n
        Just (SigD n _ : _) -> return $ VarE n
        _ -> do
            tr <- newNameU "_tr"
            let mkRec = do
                    put (n, mInsert st [FunD tr [Clause [] (NormalB $ TupE []) []]] m, c)  -- insert something to break recursion, will be replaced below.
                    trBiCase ra f ft st
                trty = st `arrow` (maybe st (flip AppT st) mty)

            cs <- if ft == st then do
                      b <- contains' ft st
                      if b then do
                          g <- newNameU "_g"
                          gcs <- mkRec
                          let dg = [SigD g trty, FunD g gcs]
                          -- Insert with a dummy type, just to get the definition in the map for mElems.
                          modify $ \ (n', m', c') -> (n', mInsert (ConT g) dg m', c')
                          x <- newNameU "_x"
                          return [Clause [VarP x] (NormalB $ rbind f (AppE (VarE g) (VarE x))) []]
                       else do
                          x <- newNameU "_x"
                          return [Clause [VarP x] (NormalB $ AppE f (VarE x)) []]
                  else do
                      b <- contains ft st
--                      lift $ qRunIO $ print (b, ft, st)
                      if b then do
                          mkRec
                       else do
                          x <- newNameU "_x"
                          return [Clause [VarP x] (NormalB $ ret $ VarE x) []]
            let d = [SigD tr trty, FunD tr cs]
            modify $ \ (n', m', c') -> (n', mInsert st d m', c')
            return $ VarE tr

trBiCase :: RetAp -> Exp -> Type -> Type -> U [Clause]
trBiCase ra f ft st = do
    let (con, ts) = splitTypeApp st
    case con of
        ConT n    -> trBiCon ra f n ft st ts
        TupleT _  -> trBiTuple ra f ft st ts
--        ArrowT    -> lift $ fmap unFunD [d| f _ _r = _r |]           -- Stop at functions
        ListT     -> trBiList ra f ft st (head ts)
        _         -> genError $ "trBiCase: unexpected type: " ++ pprint st ++ " (" ++ show st ++ ")"

trBiList :: RetAp -> Exp -> Type -> Type -> Type -> U [Clause]
trBiList ra f ft st et = do
    nil <- trMkArm ra f ft st [] (const $ ListP []) (ListE []) []
    cons <- trMkArm ra f ft st [] (ConP '(:)) (ConE '(:)) [et, st]
    return [nil, cons]

trBiTuple :: RetAp -> Exp -> Type -> Type -> [Type] -> U [Clause]
trBiTuple ra f ft st ts = do
    vs <- mapM (const $ newNameU "_t") ts
    let tupE = LamE (map VarP vs) $ ListE (map VarE vs)
    c <- trMkArm ra f ft st [] TupP tupE ts
    return [c]

trBiCon :: RetAp -> Exp -> Name -> Type -> Type -> [Type] -> U [Clause]
trBiCon ra f con ft st ts = do
    (tvs, cons) <- lift $ getTyConInfo con
    let genArm (NormalC c xs) = arm (ConP c) (ConE c) xs
        genArm (InfixC x1 c x2) = arm (\ [p1, p2] -> InfixP p1 c p2) (ConE c) [x1, x2]
        genArm (RecC c xs) = arm (ConP c) (ConE c) [ (b,t) | (_,b,t) <- xs ]
        genArm c = genError $ "trBiCon: " ++ show c
        s = mkSubst tvs ts
        arm c ec xs = trMkArm ra f ft st s c ec $ map snd xs
    mapM genArm cons

trMkArm :: RetAp -> Exp -> Type -> Type -> Subst -> ([Pat] -> Pat) -> Exp -> [Type] -> U Clause
trMkArm ra@(mty, ret, apl, _) f ft st s c ec ts = do
    vs <- mapM (const $ newNameU "_x") ts
    let sub v t = do
            let t' = subst s t
            tr <- trBi ra f ft t'
            return $ AppE tr (VarE v {- `SigE` t' -}) {- `SigE` mnd t' -}
        mnd x = maybe x (flip AppT x) mty
        conTy = foldr arrow st (map (subst s) ts)
    es <- zipWithM sub vs ts
    let body = foldl apl (ret ec {- `SigE` mnd conTy -}) es
    return $ Clause [c (map VarP vs)] (NormalB body) []


----------------------------------------------------

-- Can't use Data.Map since TH stuff is not in Ord

newtype Map a b = Map [(a, b)]

mEmpty :: Map a b
mEmpty = Map []

mLookup :: (Eq a) => a -> Map a b -> Maybe b
mLookup a (Map xys) = lookup a xys

mInsert :: (Eq a) => a -> b -> Map a b -> Map a b
mInsert a b (Map xys) = Map $ (a, b) : filter ((/= a) . fst) xys

mElems :: Map a b -> [b]
mElems (Map xys) = map snd xys

mFromList :: [(a, b)] -> Map a b
mFromList xys = Map xys