{-# LANGUAGE TemplateHaskell #-}

-- | Generate 'HFunctor' instances via @TemplateHaskell@

module Hyper.TH.Functor
    ( makeHFunctor
    ) where

import qualified Control.Lens as Lens
import           Hyper.Class.Functor (HFunctor(..))
import           Hyper.TH.Internal.Utils
import           Language.Haskell.TH
import           Language.Haskell.TH.Datatype (ConstructorVariant)

import           Hyper.Internal.Prelude

-- | Generate a 'HFunctor' instance
makeHFunctor :: Name -> DecsQ
makeHFunctor typeName = makeTypeInfo typeName >>= makeHFunctorForType

makeHFunctorForType :: TypeInfo -> DecsQ
makeHFunctorForType info =
    instanceD (makeContext info >>= simplifyContext) [t|HFunctor $(pure (tiInstance info))|]
    [ InlineP 'hmap Inline FunLike AllPhases & PragmaD & pure
    , funD 'hmap (tiConstructors info <&> makeCtr)
    ]
    <&> (:[])
    where
        (_, wit) = makeNodeOf info
        makeCtr ctr =
            clause [varP varF, pat] body []
            where
                (pat, body) = makeHMapCtr 0 wit ctr

varF :: Name
varF = mkName "_f"

makeContext :: TypeInfo -> Q [Pred]
makeContext info =
    tiConstructors info ^.. traverse . Lens._3 . traverse . Lens._Right
    & traverse ctxForPat <&> mconcat
    where
        ctxForPat (InContainer t pat) = (:) <$> [t|Functor $(pure t)|] <*> ctxForPat pat
        ctxForPat (GenEmbed t) = [t|HFunctor $(pure t)|] <&> (:[])
        ctxForPat (FlatEmbed t) = makeContext t
        ctxForPat _ = pure []

makeHMapCtr :: Int -> NodeWitnesses -> (Name, ConstructorVariant, [Either Type CtrTypePattern]) -> (Q Pat, Q Body)
makeHMapCtr i wit (cName, _, cFields) =
    (conP cName (cVars <&> varP), body)
    where
        cVars =
            [i ..] <&> show <&> ('x':) <&> mkName
            & take (length cFields)
        body =
            zipWith bodyFor cFields cVars
            & foldl appE (conE cName)
            & normalB
        bodyFor (Right x) v = bodyForPat x `appE` varE v
        bodyFor Left{} v = varE v
        f = varE varF
        bodyForPat (Node t) = [|$f $(nodeWit wit t)|]
        bodyForPat (GenEmbed t) = [|hmap ($f . $(embedWit wit t))|]
        bodyForPat (InContainer _ pat) = [|fmap $(bodyForPat pat)|]
        bodyForPat (FlatEmbed x) =
            lamCaseE
            (tiConstructors x
                <&> makeHMapCtr (i + length cVars) wit
                <&> \(p, b) -> match p b []
            )