content: {-# LANGUAGE CPP #-}
  {-# LANGUAGE TemplateHaskell #-}
  
  module Hyper.TH.Morph
      ( makeHMorph
      ) where
  
  import qualified Control.Lens as Lens
  import qualified Data.Map as Map
  import Hyper.Class.Morph (HMorph (..))
  import Hyper.TH.Internal.Utils
  import Language.Haskell.TH
  import qualified Language.Haskell.TH.Datatype as D
  
  import Hyper.Internal.Prelude
  
  makeHMorph :: Name -> DecsQ
  makeHMorph typeName = makeTypeInfo typeName >>= makeHMorphForType
  
  {-# ANN module "HLint: ignore Use id" #-}
  makeHMorphForType :: TypeInfo -> DecsQ
  makeHMorphForType info =
      -- TODO: Contexts
      instanceD
          (pure [])
          [t|HMorph $(pure src) $(pure dst)|]
          [ D.tySynInstDCompat
              ''MorphConstraint
              (Just [pure (plainTV constraintVar)])
              ([src, dst, VarT constraintVar] <&> pure)
              (simplifyContext morphConstraint <&> toTuple)
          , dataInstD
              (pure [])
              ''MorphWitness
              [pure src, pure dst, [t|_|], [t|_|]]
              Nothing
              (witnesses ^.. traverse . Lens._2)
              []
          , funD 'morphMap (tiConstructors info <&> mkMorphCon)
          , funD 'morphLiftConstraint liftConstraintClauses
          ]
          <&> (: [])
      where
          (s0, s1) = paramSubsts info
          src = D.applySubstitution s0 (tiInstance info)
          dst = D.applySubstitution s1 (tiInstance info)
          constraintVar = mkName "constraint"
          contents = childrenTypes info
          morphConstraint =
              (tcChildren contents ^.. Lens.folded <&> appSubsts (VarT constraintVar))
                  <> ( tcEmbeds contents ^.. Lens.folded
                          <&> \x -> ConT ''MorphConstraint `appSubsts` x `AppT` VarT constraintVar
                     )
          appSubsts x t = x `AppT` D.applySubstitution s0 t `AppT` D.applySubstitution s1 t
          nodeWits =
              tcChildren contents ^.. Lens.folded
                  <&> \x ->
                      let n = witPrefix <> mkNiceTypeName x & mkName
                      in  ( x
                          , (n, gadtC [n] [] (pure (appSubsts morphWithNessOf x)))
                          )
          embedWits =
              tcEmbeds contents ^.. Lens.folded
                  <&> \x ->
                      let n = witPrefix <> mkNiceTypeName x & mkName
                      in  ( x
                          ,
                              ( n
                              , gadtC
                                  [n]
                                  [ bangType
                                      (bang noSourceUnpackedness noSourceStrictness)
                                      (pure (ConT ''MorphWitness `appSubsts` x `AppT` varA `AppT` varB))
                                  ]
                                  (pure (morphWithNessOf `AppT` varA `AppT` varB))
                              )
                          )
          witnesses = nodeWits <> embedWits & Map.fromList
          varA = VarT (mkName "a")
          varB = VarT (mkName "b")
          witPrefix = "M_" <> niceName (tiName info) <> "_"
          morphWithNessOf = ConT ''MorphWitness `AppT` src `AppT` dst
          liftConstraintClauses
              | Map.null witnesses = [clause [] (normalB (lamCaseE [])) []]
              | otherwise =
                  (nodeWits ^.. traverse . Lens._2 . Lens._1 <&> liftNodeConstraint)
                      <> (embedWits ^.. traverse . Lens._2 . Lens._1 <&> liftEmbedConstraint)
          liftNodeConstraint n = clause [conP n [], wildP] (normalB [|\x -> x|]) []
          liftEmbedConstraint n =
              clause
                  [conP n [varP varW], varP varProxy]
                  (normalB [|morphLiftConstraint $(varE varW) $(varE varProxy)|])
                  []
          varW = mkName "w"
          varProxy = mkName "p"
          mkMorphCon con =
              clause [varP varF, p] b []
              where
                  (p, b) = morphCon 0 witnesses con
  
  varF :: Name
  varF = mkName "_f"
  
  morphCon :: Int -> Map Type (Name, a) -> (Name, b, [Either c CtrTypePattern]) -> (Q Pat, Q Body)
  morphCon i witnesses (n, _, fields) =
      ( conP n (cVars <&> varP)
      , normalB (foldl appE (conE n) (zipWith bodyFor fields cVars))
      )
      where
          cVars =
              [i ..]
                  <&> mkName . ('x' :) . show
                  & take (length fields)
          f = varE varF
          bodyFor Left{} v = varE v
          bodyFor (Right x) v = [|$(bodyForPat x) $(varE v)|]
          bodyForPat (Node x) = [|$f $(conE (witnesses ^?! Lens.ix x . Lens._1))|]
          bodyForPat (InContainer _ pat) = [|fmap $(bodyForPat pat)|]
          bodyForPat (FlatEmbed x) =
              lamCaseE
                  ( tiConstructors x
                      <&> uncurry match
                          . morphCon (i + length cVars) witnesses
                          ?? []
                  )
          bodyForPat (GenEmbed t) = [|morphMap ($f . $(conE (witnesses ^?! Lens.ix t . Lens._1)))|]
  
  type MorphSubsts = (Map Name Type, Map Name Type)
  
  paramSubsts :: TypeInfo -> MorphSubsts
  paramSubsts info =
      (tiParams info <&> D.tvName) ^. traverse . Lens.to mkInfo
      where
          pinned = pinnedParams info
          mkInfo name
              | pinned ^. Lens.contains name = mempty
              | otherwise = (side name "0", side name "1")
          side name suffix = mempty & Lens.at name ?~ VarT (mkName (nameBase name <> suffix))
  
  pinnedParams :: TypeInfo -> Set Name
  pinnedParams = (^. Lens.to tiConstructors . traverse . Lens._3 . traverse . Lens.to ctrPinnedParams)
  
  ctrPinnedParams :: Either Type CtrTypePattern -> Set Name
  ctrPinnedParams (Left t) = typeParams t
  ctrPinnedParams (Right Node{}) = mempty
  ctrPinnedParams (Right GenEmbed{}) = mempty
  ctrPinnedParams (Right (FlatEmbed info)) = pinnedParams info
  ctrPinnedParams (Right (InContainer c p)) = typeParams c <> ctrPinnedParams (Right p)
  
  typeParams :: Type -> Set Name
  typeParams (VarT x) = mempty & Lens.contains x .~ True
  typeParams (AppT f x) = typeParams f <> typeParams x
  typeParams (InfixT x _ y) = typeParams x <> typeParams y
  -- TODO: Missing cases
  typeParams _ = mempty
  