content: {-# LANGUAGE DerivingVia #-}
  {-# LANGUAGE TemplateHaskell #-}
  
  -- Helpers for TemplateHaskell instance generators
  
  module Hyper.TH.Internal.Utils
      ( -- Internals for use in TH for sub-classes
        TypeInfo (..)
      , TypeContents (..)
      , CtrTypePattern (..)
      , NodeWitnesses (..)
      , makeTypeInfo
      , makeNodeOf
      , parts
      , toTuple
      , matchType
      , niceName
      , mkNiceTypeName
      , applicativeStyle
      , unapply
      , getVar
      , makeConstructorVars
      , consPat
      , simplifyContext
      , childrenTypes
      ) where
  
  import qualified Control.Lens as Lens
  import Control.Monad.Trans.Class (MonadTrans (..))
  import Control.Monad.Trans.State (State, evalState, execStateT, gets, modify)
  import qualified Data.Char as Char
  import Data.List (intercalate, nub)
  import qualified Data.Map as Map
  import Generic.Data (Generically (..))
  import Hyper.Class.Nodes (HWitness (..))
  import Hyper.Type (AHyperType (..), GetHyperType, type (:#))
  import Language.Haskell.TH
  import qualified Language.Haskell.TH.Datatype as D
  import Language.Haskell.TH.Datatype.TyVarBndr
  
  import Hyper.Internal.Prelude
  
  data TypeInfo = TypeInfo
      { tiName :: Name
      , tiInstance :: Type
      , tiParams :: [TyVarBndrUnit]
      , tiHyperParam :: Name
      , tiConstructors :: [(Name, D.ConstructorVariant, [Either Type CtrTypePattern])]
      }
      deriving (Show)
  
  data TypeContents = TypeContents
      { tcChildren :: Set Type
      , tcEmbeds :: Set Type
      , tcOthers :: Set Type
      }
      deriving (Show, Generic)
      deriving (Semigroup, Monoid) via Generically TypeContents
  
  data CtrTypePattern
      = Node Type
      | FlatEmbed TypeInfo
      | GenEmbed Type
      | InContainer Type CtrTypePattern
      deriving (Show)
  
  makeTypeInfo :: Name -> Q TypeInfo
  makeTypeInfo name =
      do
          info <- D.reifyDatatype name
          (dst, var) <- parts info
          let makeCons c =
                  traverse (matchType name var) (D.constructorFields c)
                      <&> (D.constructorName c,D.constructorVariant c,)
          cons <- traverse makeCons (D.datatypeCons info)
          pure
              TypeInfo
                  { tiName = name
                  , tiInstance = dst
                  , tiParams = D.datatypeVars info & init
                  , tiHyperParam = var
                  , tiConstructors = cons
                  }
  
  parts :: D.DatatypeInfo -> Q (Type, Name)
  parts info =
      case D.datatypeVars info of
          [] -> fail "expected type constructor which requires arguments"
          xs ->
              elimTV
                  (pure . (,) res)
                  ( \var c ->
                      case c of
                          ConT aHyper | aHyper == ''AHyperType -> pure (res, var)
                          _ -> fail "expected last argument to be a AHyperType variable"
                  )
                  (last xs)
              where
                  res =
                      foldl AppT (ConT (D.datatypeName info)) (init xs <&> VarT . D.tvName)
  
  childrenTypes :: TypeInfo -> TypeContents
  childrenTypes info = evalState (childrenTypesH info) mempty
  
  childrenTypesH ::
      TypeInfo -> State (Set Type) TypeContents
  childrenTypesH info =
      do
          did <- gets (^. Lens.contains (tiInstance info))
          if did
              then pure mempty
              else
                  modify (Lens.contains (tiInstance info) .~ True)
                      *> traverse addPat (tiConstructors info ^.. traverse . Lens._3 . traverse . Lens._Right)
                      <&> mconcat
      where
          addPat (FlatEmbed inner) = childrenTypesH inner
          addPat (Node x) = pure mempty{tcChildren = mempty & Lens.contains x .~ True}
          addPat (GenEmbed x) = pure mempty{tcEmbeds = mempty & Lens.contains x .~ True}
          addPat (InContainer _ x) = addPat x
  
  unapply :: Type -> (Type, [Type])
  unapply =
      go []
      where
          go as (SigT x _) = go as x
          go as (AppT f a) = go (a : as) f
          go as x = (x, as)
  
  matchType :: Name -> Name -> Type -> Q (Either Type CtrTypePattern)
  matchType _ var (ConT get `AppT` VarT h `AppT` (PromotedT aHyper `AppT` x))
      | get == ''GetHyperType && aHyper == 'AHyperType && h == var =
          Node x & Right & pure
  matchType _ var (InfixT (VarT h) hash x)
      | hash == ''(:#) && h == var =
          Node x & Right & pure
  matchType _ var (ConT hash `AppT` VarT h `AppT` x)
      | hash == ''(:#) && h == var =
          Node x & Right & pure
  matchType top var (x `AppT` VarT h)
      | h == var && x /= ConT ''GetHyperType =
          case unapply x of
              (ConT c, args) | c /= top ->
                  do
                      inner <- D.reifyDatatype c
                      let innerVars = D.datatypeVars inner <&> D.tvName
                      let subst =
                              args <> [VarT var]
                                  & zip innerVars
                                  & Map.fromList
                      let makeCons i =
                              traverse (matchType top var . D.applySubstitution subst) (D.constructorFields i)
                                  <&> (D.constructorName i,D.constructorVariant i,)
                      cons <- traverse makeCons (D.datatypeCons inner)
                      if var `notElem` (D.freeVariablesWellScoped (cons ^.. traverse . Lens._3 . traverse . Lens._Left) <&> D.tvName)
                          then
                              FlatEmbed
                                  TypeInfo
                                      { tiName = c
                                      , tiInstance = x
                                      , tiParams = D.datatypeVars inner & init
                                      , tiHyperParam = var
                                      , tiConstructors = cons
                                      }
                                  & pure
                          else GenEmbed x & pure
              _ -> GenEmbed x & pure
              <&> Right
  matchType top var x@(AppT f a) =
      -- TODO: check if applied over a functor-kinded type.
      matchType top var a
          <&> \case
              Left{} -> Left x
              Right pat -> InContainer f pat & Right
  matchType _ _ t = Left t & pure
  
  getVar :: Type -> Maybe Name
  getVar (VarT x) = Just x
  getVar (SigT x _) = getVar x
  getVar _ = Nothing
  
  toTuple :: Foldable t => t Type -> Type
  toTuple xs = foldl AppT (TupleT (length xs)) xs
  
  applicativeStyle :: Q Exp -> [Q Exp] -> Q Exp
  applicativeStyle f =
      foldl ap [|pure $f|]
      where
          ap x y = [|$x <*> $y|]
  
  makeConstructorVars :: String -> [a] -> [(a, Name)]
  makeConstructorVars prefix fields =
      [0 :: Int ..]
          <&> mkName . (('_' : prefix) <>) . show
          & zip fields
  
  consPat :: Name -> [(a, Name)] -> Q Pat
  consPat c = conP c . (<&> varP . snd)
  
  simplifyContext :: [Pred] -> CxtQ
  simplifyContext preds =
      execStateT (goPreds preds) (mempty :: Set (Name, [Type]), mempty :: Set Pred)
          <&> (^.. Lens._2 . Lens.folded)
      where
          goPreds = traverse_ (go . unapply)
          go (c, [VarT v]) =
              -- Work-around reifyInstances returning instances for type variables
              -- by not checking.
              yep c [VarT v]
          go (ConT c, xs) =
              Lens.use (Lens._1 . Lens.contains key)
                  >>= \case
                      True -> pure () -- already checked
                      False ->
                          do
                              Lens._1 . Lens.contains key .= True
                              reifyInstances c xs
                                  & lift
                                  >>= \case
                                      [InstanceD _ context other _] ->
                                          D.unifyTypes [other, foldl AppT (ConT c) xs]
                                              & lift
                                              <&> (`D.applySubstitution` context)
                                              >>= goPreds
                                      _ -> yep (ConT c) xs
              where
                  key = (c, xs)
          go (c, xs) = yep c xs
          yep c xs = Lens._2 . Lens.contains (foldl AppT c xs) .= True
  
  data NodeWitnesses = NodeWitnesses
      { nodeWit :: Type -> Q Exp
      , embedWit :: Type -> Q Exp
      , nodeWitCtrs :: [Name]
      , embedWitCtrs :: [Name]
      }
  
  niceName :: Name -> String
  niceName = reverse . takeWhile (/= '.') . reverse . show
  
  makeNodeOf :: TypeInfo -> ([Type -> Q Con], NodeWitnesses)
  makeNodeOf info =
      ( (nodes <&> nodeGadtType) <> (embeds <&> embedGadtType)
      , NodeWitnesses
          { nodeWit = nodes & Map.fromList & getWit <&> \x -> [|HWitness $(conE x)|]
          , embedWit = embeds & Map.fromList & getWit <&> \x -> [|HWitness . $(conE x)|]
          , nodeWitCtrs = nodes <&> snd
          , embedWitCtrs = embeds <&> snd
          }
      )
      where
          niceTypeName = tiName info & niceName
          nodeBase = "W_" <> niceTypeName <> "_"
          embedBase = "E_" <> niceTypeName <> "_"
          pats = tiConstructors info >>= (^. Lens._3)
          nodes =
              pats ^.. traverse . Lens._Right
                  >>= nodesForPat
                  & nub
                  <&> \t -> (t, mkName (nodeBase <> mkNiceTypeName t))
          nodesForPat (Node t) = [t]
          nodesForPat (InContainer _ pat) = nodesForPat pat
          nodesForPat (FlatEmbed x) = tiConstructors x ^.. traverse . Lens._3 . traverse . Lens._Right >>= nodesForPat
          nodesForPat _ = []
          nodeGadtType (t, n) c = gadtC [n] [] (pure (c `AppT` t))
          embeds =
              pats ^.. traverse . Lens._Right
                  >>= embedsForPat
                  & nub
                  <&> \t -> (t, mkName (embedBase <> mkNiceTypeName t))
          embedsForPat (GenEmbed t) = [t]
          embedsForPat (InContainer _ pat) = embedsForPat pat
          embedsForPat (FlatEmbed x) = tiConstructors x ^.. traverse . Lens._3 . traverse . Lens._Right >>= embedsForPat
          embedsForPat _ = []
          embedGadtType (t, n) c =
              gadtC
                  [n]
                  [ bangType
                      (bang noSourceUnpackedness noSourceStrictness)
                      [t|HWitness $(pure t) $nodeVar|]
                  ]
                  [t|$(pure c) $nodeVar|]
          nodeVar = mkName "node" & varT
          getWit :: Map Type Name -> Type -> Name
          getWit m h =
              m ^? Lens.ix h
                  & fromMaybe (error ("Cant find witness for " <> show h <> " in " <> show m))
  
  mkNiceTypeName :: Type -> String
  mkNiceTypeName =
      intercalate "_" . makeNiceType
      where
          makeNiceType (ConT x) =
              case niceName x of
                  n@(c : _) | Char.isAlpha c -> [n]
                  _ -> [] -- Skip operators
          makeNiceType (AppT x y) = makeNiceType x <> makeNiceType y
          makeNiceType (VarT x) = [takeWhile (/= '_') (show x)]
          makeNiceType (SigT x _) = makeNiceType x
          makeNiceType x = error ("TODO: Witness name generator is partial! Need to support " <> show x)
  