{-# LANGUAGE TemplateHaskell            #-}
{-# LANGUAGE FlexibleContexts           #-}
{-# LANGUAGE TypeFamilies               #-}
{-# LANGUAGE UndecidableInstances       #-}

module Generics.Regular.Transformations.TH ( 
  deriveRefRep, prefix, postfix
  ) where

import Generics.Regular
import Language.Haskell.TH
import Generics.Regular.Transformations.Explicit

-- Code taken from regular library and adapted to work with transformations

-- | Derive data type with references and 'HasRef' instance. For a data type
--   N the name of the constructor for a reference is RefN, and the given
--   function is used to change the rest of the constructors and the data
--   type name itself. For example, for the following definition:
--
-- > data Tree = Leaf Int | Bin Tree Tree
-- > $(deriveRefRep ''Tree (postfix "R"))
--
-- The following data type is generated:
--
-- > data TreeR = LeafR Int | BinR TreeR TreeR | RefTree Path
-- > instance HasRef Tree
deriveRefRep :: Name -> (Name -> Name) -> Q [Dec]
deriveRefRep t namef = do
  d <- deriveData t namef
  ins <- deriveInst t namef
  return [d,ins]

deriveData :: Name -> (Name -> Name) -> Q Dec
deriveData t namef = do
  let nm = namef t
  i <- reify t
  cons <- case i of
    TyConI (DataD _ _ _ cs _) -> mapM (mkCon t namef nm) cs
  r <- normalC (prefix "Ref" t) [return (NotStrict, ConT ''Path)]
  dataD (cxt []) nm (typeVariables i) (map return $ r : cons) []

mkCon :: Name -> (Name -> Name) -> Name -> Con -> Q Con
mkCon t namef repname (NormalC a b) = normalC (namef a) (map f b) where
  f :: (Strict, Type) -> Q (Strict, Type)
  f (s,t') | t' == ConT t = return (s, ConT repname)
           | otherwise    = return (s, t')

prefix :: String -> Name -> Name
prefix pref n = mkName $ pref ++ nameBase n

postfix :: String -> Name -> Name
postfix post n = mkName $ nameBase n ++ post

deriveInst :: Name -> (Name -> Name) -> Q Dec
deriveInst t namef =
  do
    i <- reify t
    let typ = foldl (\a -> AppT a . VarT . tyVarBndrToName) (ConT t) (typeVariables i)
    let rn = prefix "Ref" t
    fcs <- mkFrom t 1 0 namef rn t
    tcs <- mkTo   t 1 0 namef rn t
    let typ' = return $ foldl (\a -> AppT a . VarT . tyVarBndrToName) (ConT $ namef t) (typeVariables i)      
    instanceD (cxt []) (conT ''HasRef `appT` return typ)
        [tySynInstD ''RefRep [return typ] typ', funD 'toRef tcs, funD 'fromRef fcs]

lrE :: Int -> Int -> (Q Exp -> Q Exp)
lrE 1 0 e = e
lrE m 0 e = conE 'L `appE` e
lrE m i e = conE 'R `appE` lrE (m-1) (i-1) e

tyVarBndrToName :: TyVarBndr -> Name
tyVarBndrToName (PlainTV  name)   = name
tyVarBndrToName (KindedTV name _) = name

typeVariables :: Info -> [TyVarBndr]
typeVariables (TyConI (DataD    _ _ tv _ _)) = tv
typeVariables (TyConI (NewtypeD _ _ tv _ _)) = tv
typeVariables _                           = []

mkFrom :: Name -> Int -> Int -> (Name -> Name) -> Name -> Name -> Q [Q Clause]
mkFrom ns m i namef refname n = do
  let wrapE e = conE 'InR `appE` lrE m i e
  i <- reify n
  let r = clause [conP refname [varP $ field 0]] (normalB $ conE 'Ref `appE` varE (field 0)) []
  let b = case i of
        TyConI (DataD _ dt vs cs _) ->
          zipWith (fromCon wrapE ns namef (dt, map tyVarBndrToName vs) (length cs)) [0..] cs
        TyConI (NewtypeD _ dt vs c _) ->
          [fromCon wrapE ns namef (dt, map tyVarBndrToName vs) 1 0 c]
        TyConI (TySynD t _ _) ->
          [clause [varP (field 0)] (normalB (wrapE $ conE 'K `appE` varE (field 0))) []]
        _ -> error "unknown construct"
  return $ r : b

mkTo :: Name -> Int -> Int -> (Name -> Name) -> Name -> Name -> Q [Q Clause]
mkTo ns m i namef refname n = do
  let wrapP p = conP 'InR [lrP m i p]
  i <- reify n
  let r = clause [conP 'Ref [varP $ field 0]] (normalB $ conE refname `appE` varE (field 0)) []
  let b = case i of
                TyConI (DataD _ dt vs cs _) ->
                  zipWith (toCon wrapP ns namef (dt, map tyVarBndrToName vs) (length cs)) [0..] cs
                TyConI (NewtypeD _ dt vs c _) ->
                  [toCon wrapP ns namef (dt, map tyVarBndrToName vs) 1 0 c]
                TyConI (TySynD t _ _) ->
                  [clause [wrapP $ conP 'K [varP (field 0)]] (normalB $ varE (field 0)) []]
                _ -> error "unknown construct" 
  return $ r : b

fromCon :: (Q Exp -> Q Exp) -> Name -> (Name -> Name) -> (Name, [Name]) -> Int -> Int -> Con -> Q Clause
fromCon wrap ns namef (dt, vs) m i (NormalC cn []) =
    clause
      [conP (namef cn) []]
      (normalB $ wrap $ lrE m i $ conE 'C `appE` (conE 'U)) []
fromCon wrap ns namef (dt, vs) m i (NormalC cn fs) =
    clause
      [conP (namef cn) (map (varP . field) [0..length fs - 1])]
      (normalB $ wrap $ lrE m i $ conE 'C `appE` foldr1 prod (zipWith (fromField (dt, vs)) [0..] (map snd fs))) []
  where
    prod x y = conE '(:*:) `appE` x `appE` y
fromCon wrap ns namef (dt, vs) m i r@(RecC cn []) =
    clause
      [conP (namef cn) []]
      (normalB $ wrap $ lrE m i $ conE 'C `appE` (conE 'U)) []
fromCon wrap ns namef (dt, vs) m i r@(RecC cn fs) =
    clause
      [conP (namef cn) (map (varP . field) [0..length fs - 1])]
      (normalB $ wrap $ lrE m i $ conE 'C `appE` foldr1 prod (zipWith (fromField' (dt, vs)) [0..] fs)) []
  where
    prod x y = conE '(:*:) `appE` x `appE` y
fromCon wrap ns namef (dt, vs) m i (InfixC t1 cn t2) =
  fromCon wrap ns namef (dt, vs) m i (NormalC cn [t1,t2])

fromField :: (Name, [Name]) -> Int -> Type -> Q Exp
fromField (dt, vs) nr t | t == dataDeclToType (dt, vs) = 
  conE 'I `appE` varE (field nr)
fromField (dt, vs) nr t                                = 
  conE 'K `appE` varE (field nr)

fromField' :: (Name, [Name]) -> Int -> (Name, Strict, Type) -> Q Exp
fromField' (dt, vs) nr (_, _, t) | t == dataDeclToType (dt, vs) =
  conE 'S `appE` (conE 'I `appE` varE (field nr))
fromField' (dt, vs) nr (_, _, t)                                =
  conE 'S `appE` (conE 'K `appE` varE (field nr))

toCon :: (Q Pat -> Q Pat) -> Name -> (Name -> Name) -> (Name, [Name]) -> Int -> Int -> Con -> Q Clause
toCon wrap ns namef (dt, vs) m i (NormalC cn []) =
    clause
      [wrap $ lrP m i $ conP 'C [conP 'U []]]
      (normalB $ conE $ namef cn) []
toCon wrap ns namef (dt, vs) m i (NormalC cn fs) =
    -- runIO (putStrLn ("constructor " ++ show ix)) >>
    clause
      [wrap $ lrP m i $ conP 'C [foldr1 prod (zipWith (toField (dt, vs)) [0..] (map snd fs))]]
      (normalB $ foldl appE (conE $ namef cn) (map (varE . field) [0..length fs - 1])) []
  where
    prod x y = conP '(:*:) [x,y]
toCon wrap ns namef (dt, vs) m i r@(RecC cn []) =
    clause
      [wrap $ lrP m i $ conP 'C [conP 'U []]]
      (normalB $ conE $ namef cn) []
toCon wrap ns namef (dt, vs) m i r@(RecC cn fs) =
    clause
      [wrap $ lrP m i $ conP 'C [foldr1 prod (zipWith (toField' (dt, vs)) [0..] fs)]]
      (normalB $ foldl appE (conE $ namef cn) (map (varE . field) [0..length fs - 1])) []
  where
    prod x y = conP '(:*:) [x,y]
toCon wrap ns namef (dt, vs) m i (InfixC t1 cn t2) =
  toCon wrap ns namef (dt, vs) m i (NormalC cn [t1,t2])

toField :: (Name, [Name]) -> Int -> Type -> Q Pat
toField (dt, vs) nr t | t == dataDeclToType (dt, vs) = 
  conP 'I [varP (field nr)]
toField (dt, vs) nr t                                = 
  conP 'K [varP (field nr)]

toField' :: (Name, [Name]) -> Int -> (Name, Strict, Type) -> Q Pat
toField' (dt, vs) nr (_, _, t) | t == dataDeclToType (dt, vs) = conP 'S [conP 'I [varP (field nr)]]
toField' (dt, vs) nr (_, _, t)                                = conP 'S [conP 'K [varP (field nr)]]

field :: Int -> Name
field n = mkName $ "f" ++ show n

lrP :: Int -> Int -> (Q Pat -> Q Pat)
lrP 1 0 p = p
lrP m 0 p = conP 'L [p]
lrP m i p = conP 'R [lrP (m-1) (i-1) p]

dataDeclToType :: (Name, [Name]) -> Type
dataDeclToType (dt, vs) = foldl (\a b -> AppT a (VarT b)) (ConT dt) vs