-----------------------------------------------------------------------------
-- |
-- Module      :  Generics.Putlenses.TH
-- Copyright   :  (C) 2013 Hugo Pacheco
-- License     :  BSD-style (see the file LICENSE)
-- Maintainer  :  Hugo Pacheco <hpacheco@nii.ac.jp>
-- Stability   :  provisional
--
-- Template Haskell procedures to generate default putlenses for record types and default putlens constructors and destructores for @Generic@ instances.
-- 
--
--
----------------------------------------------------------------------------

module Generics.Putlenses.TH (
    makePutlensFields,
    makePutlensConstructors
) where

import Language.Haskell.TH.Syntax
import Language.Haskell.TH.Lib
import Generics.Putlenses.Putlens
import qualified Control.Lens as L
import Control.Lens.TH
import Data.Char
import Data.List
import Control.Monad

-- | Generates field putlenses for record types
-- For a record
-- @data T a = T { t1 :: String, t2 :: a }
-- $( makePutlensFields ''T )@
-- it generates two putlenses
-- @t1Put :: Monad m => PutlensM m e (T a) String
-- t2Put :: Monad m => PutlensM e (T a) a@
makePutlensFields :: Name -> Q [Dec]
makePutlensFields name = do
    lensdecs <- makeMyLenses name
    putlensdecs <- liftM concat (mapM makePutlensField (lensnames lensdecs))
    return (lensdecs ++ putlensdecs)

lensnames :: [Dec] -> [String]
lensnames [] = []
lensnames (d:ds) = let s = lensname d
                       ss = lensnames ds
                   in maybe ss (:ss) s
lensname :: Dec -> Maybe String
lensname (SigD name _) | isLens str = Just (take (length str - 4) str)
                       | otherwise = Nothing
    where str = nameBase name
          isLens xs = isSuffixOf "Lens" xs
lensname _ = Nothing

makePutlensField :: String -> Q [Dec]
makePutlensField name = do
	let pat = varP (mkName $ name++"Put")
	    body = normalB $ appE (varE $ mkName "Generics.Putlenses.Putlens.simplelens2put") (varE $ mkName $ name++"Lens")
	dec <- valD pat body []
	return [dec]

mylensRules :: LensRules
mylensRules = L.set lensField lensName lensRules
lensName :: [Name] -> Name -> [DefName]
lensName ns n = case nameBase n of
	"" -> [TopName $ mkName ""]
	x:xs -> [TopName (mkName $ toLower x : xs ++ "Lens")]

makeMyLenses :: Name -> Q [Dec]
makeMyLenses name = makeLensesWith mylensRules name

-- | Generates constructor and destructor putlenses for data types that are @Generic@ instances
-- For an algebraic data type
-- @data List a = Nil | Cons a [a] deriving Generic
-- $( makePutlensConstructors ''List )@
-- it generates two constructors
-- @nilPut :: MonadPlus m => PutlensM m e (List a) ()
-- nilPut = innPut .< injlPut
-- consPut :: Monad m => PutlensM m e (List a) (a,List a)
-- consPut = innPut .< injrPut@
-- and two destructors
-- @unnilPut :: Monad m => PutlensM m e () (List a)
-- unnilPut = uninjlPut .< outPut
-- unconsPut :: Monad m => PutlensM m e (a,List a) (List a)
-- unconsPut = uninjrPut .< outPut@
makePutlensConstructors :: Name -> Q [Dec]
makePutlensConstructors name = do
    info <- reify name
    case info of
      (TyConI decl) -> case deNewtype decl of
        (DataD ctx tyConName args cons _) -> do
	      let srcvars = map tyVar args
	          src = appsT (return (ConT tyConName)) (map varT srcvars)
	      makeConstructors srcvars src cons
        _ -> fail "makeLensesWith: Unsupported data type"
      _ -> fail "makeLensesWith: Expected the name of a data type or newtype"
  where
    deNewtype (NewtypeD ctx tyConName args c d) = DataD ctx tyConName args [c] d
    deNewtype d = d
    tyVar (PlainTV n) = n
    tyVar (KindedTV n k) = n

constructorType :: Con -> TypeQ
constructorType c = genNestedPair (constructorArgs c)

constructorArgs :: Con -> [TypeQ]
constructorArgs (NormalC _ ts) = map (return . snd) ts
constructorArgs (RecC _ ts) = map (\(x,y,z) -> return z) ts
constructorArgs (ForallC _ _ con) = constructorArgs con
constructorArgs _ = fail "makePutlensConstructors: Unsupported infix constructor"

-- GHC.Generics splits nested pairs in half
genNestedPair :: [TypeQ] -> TypeQ
genNestedPair [] = tupleT 0
genNestedPair [t] = t
genNestedPair ts = appT (appT (tupleT 2) lp) rp
    where (l,r) = splitAt (length ts `div` 2) ts
          (lp,rp) = (genNestedPair l,genNestedPair r)

appArgs :: Type -> [Type] -> Type
appArgs t [] = t
appArgs t (x:xs) = foldl AppT (AppT t x) xs

appsT :: TypeQ -> [TypeQ] -> TypeQ
appsT = Prelude.foldl appT

makeConstructors :: [Name] -> TypeQ -> [Con] -> Q [Dec]
makeConstructors srcvars src cons = liftM concat $ mapM (makeConstructor srcvars src) $ zip cons (makeChoices (length cons))

makeConstructor :: [Name] -> TypeQ -> (Con,Choice) -> Q [Dec]
makeConstructor srcvars src (c@(NormalC name _),cs) = makeConstructor' srcvars src (constructorType c) (nameBase name) cs
makeConstructor srcvars src (c@(RecC name _),cs) = makeConstructor' srcvars src (constructorType c) (nameBase name) cs
makeConstructor srcvars src (c@(ForallC _ _ con),cs) = makeConstructor srcvars src (con,cs)
makeConstructor srcvars src _ = fail "makePutlensConstructors: Unsupported infix constructor"

makeConstructor' :: [Name] -> TypeQ -> TypeQ -> String -> Choice -> Q [Dec]
makeConstructor' srcvars src tgt name cs = do
    sigc <- genSig (constructorName name) srcvars src tgt
    c <- genConstructor name cs
    sigd <- genSig (destructorName name) srcvars tgt src
    d <- genDestructor name cs
    return [sigc,c,sigd,d]

genSig :: Name -> [Name] -> TypeQ -> TypeQ -> DecQ
genSig name srcvars src tgt = do
	m <- return $ mkName "m"
	monad <- return $ mkName "Monad"
	let putlens = appsT (return (ConT ''PutlensM)) [varT m,src,tgt]
	sigD name (forallT (map plainTV (m:srcvars)) (liftM (:[]) $ classP monad [varT m]) putlens)

genConstructor :: String -> Choice -> Q Dec
genConstructor name c = do
    let pat = varP (constructorName name)
        body = case c of
          [] -> normalB (varE $ mkName "innPut")
          otherwise -> normalB $ infixApp (varE $ mkName "innPut") (varE $ mkName ".<") (genInjs c)
    valD pat body []

genInjs :: Choice -> ExpQ
genInjs [b] = genInj b
genInjs (b:bs) = infixApp (genInj b) (varE (mkName ".<")) (genInjs bs)
genInj True = varE (mkName "injlPut")
genInj False = varE (mkName "injrPut")

constructorName :: String -> Name
constructorName (n:ns) = mkName $ toLower n : ns ++ "Put"

genDestructor :: String -> Choice -> Q Dec
genDestructor name c = do
    let pat = varP (destructorName name)
        body = case c of
          [] -> normalB (varE $ mkName "outPut")
          otherwise -> normalB $ infixApp (genUninjs (reverse c)) (varE $ mkName ".<") (varE $ mkName "outPut")
    valD pat body []

genUninjs :: Choice -> Q Exp
genUninjs [b] = genUninj b
genUninjs (b:bs) = infixApp (genUninj b) (varE (mkName ".<")) (genUninjs bs)
genUninj True = varE (mkName "uninjlPut")
genUninj False = varE (mkName "uninjrPut")

destructorName :: String -> Name
destructorName (n:ns) = mkName $ "un" ++ toLower n : ns ++ "Put"

type Choice = [Bool]

-- GHC.Generics splits nested sums in half
makeChoices :: Int -> [Choice]
makeChoices n = makeChoices' n []

makeChoices' :: Int -> Choice -> [Choice]
makeChoices' 1 args = [args]
makeChoices' n args = makeChoices' l (args++[True]) ++ makeChoices' r (args++[False])
    where l = div n 2
          r = n - l