content: module Optics.TH.Internal.Utils where
  
  import Control.Monad
  import Data.Maybe
  import Language.Haskell.TH
  import Language.Haskell.TH.Datatype.TyVarBndr
  import qualified Data.List as L
  import qualified Data.Map as M
  import qualified Data.Set as S
  import qualified Language.Haskell.TH.Datatype as D
  
  import Data.Set.Optics
  import Language.Haskell.TH.Optics.Internal
  import Optics.Core
  
  -- | Apply arguments to a type constructor
  appsT :: TypeQ -> [TypeQ] -> TypeQ
  appsT = foldl appT
  
  -- | Apply arguments to a function
  appsE1 :: ExpQ -> [ExpQ] -> ExpQ
  appsE1 = foldl appE
  
  -- | Construct a tuple type given a list of types.
  toTupleT :: [TypeQ] -> TypeQ
  toTupleT [x] = x
  toTupleT xs = appsT (tupleT (length xs)) xs
  
  -- | Construct a tuple value given a list of expressions.
  toTupleE :: [ExpQ] -> ExpQ
  toTupleE [x] = x
  toTupleE xs = tupE xs
  
  -- | Construct a tuple pattern given a list of patterns.
  toTupleP :: [PatQ] -> PatQ
  toTupleP [x] = x
  toTupleP xs = tupP xs
  
  -- | Apply arguments to a type constructor.
  conAppsT :: Name -> [Type] -> Type
  conAppsT conName = foldl AppT (ConT conName)
  
  -- | Generate many new names from a given base name.
  newNames :: String {- ^ base name -} -> Int {- ^ count -} -> Q [Name]
  newNames base n = sequence [ newName (base++show i) | i <- [1..n] ]
  
  -- We substitute concrete types with type variables and match them with concrete
  -- types in the instance context. This significantly improves type inference as
  -- GHC can match the instance more easily, but costs dependence on TypeFamilies
  -- and UndecidableInstances.
  eqSubst :: Type -> String -> Q (Type, Pred)
  eqSubst ty n = do
    placeholder <- VarT <$> newName n
    pure (placeholder, D.equalPred placeholder ty)
  
  addKindInfo :: D.DatatypeInfo -> Type -> Type
  addKindInfo = addKindInfo' []
  
  -- | Fill in kind variables using info from datatype type parameters.
  addKindInfo' :: [Type] -> D.DatatypeInfo -> Type -> Type
  addKindInfo' additionalInfo di =
    substType . M.fromList . mapMaybe var $ additionalInfo ++ D.datatypeInstTypes di
    where
      -- If the type is a data/newtype family instance, we need to fill in all of
      -- the kinds for weird cases such as:
      --
      -- data family KDF (a :: k)
      -- data instance KDF (a :: Type) = Kinded3 { _kdf :: Proxy a }
      --
      -- Otherwise we only need info about kind variables.
      --
      -- More info at https://github.com/ekmett/lens/pull/945.
      isDataFamily = D.datatypeVariant di == D.DataInstance
                  || D.datatypeVariant di == D.NewtypeInstance
  
      var t@(SigT (VarT n) k)
        | isDataFamily   = Just (n, t)
        | has typeVars k = Just (n, t)
        | otherwise      = Nothing
      var _              = Nothing
  
  -- | Template Haskell wants type variables declared in a forall, so
  -- we find all free type variables in a given type and declare them.
  quantifyType :: [TyVarBndrSpec] -> Cxt -> Type -> Type
  quantifyType = quantifyType' S.empty
  
  -- | This function works like 'quantifyType' except that it takes
  -- a list of variables to exclude from quantification.
  quantifyType' :: S.Set Name -> [TyVarBndrSpec] -> Cxt -> Type -> Type
  quantifyType' exclude vars cx t = ForallT vs cx t
    where
      vs = filter (\v -> D.tvName v `S.notMember` exclude)
         . changeTVFlags SpecifiedSpec
         . D.freeVariablesWellScoped
         $ map tyVarBndrToType vars ++ S.toList (setOf typeVarsKinded t)
  
  -- | Transform 'TyVarBndr' into a 'Type' so it's suitable e.g. for
  -- freeVariablesWellScoped or type substitution.
  tyVarBndrToType :: TyVarBndr_ flag -> Type
  tyVarBndrToType = elimTV VarT (\n k -> SigT (VarT n) k)
  
  -- | Pass in a list of lists of extensions, where any of the given extensions
  -- will satisfy it. For example, you might need either GADTs or
  -- ExistentialQuantification, so you'd write:
  --
  -- > requireExtensions [[GADTs, ExistentialQuantification]]
  --
  -- But if you need TypeFamilies and MultiParamTypeClasses, then you'd write:
  --
  -- > requireExtensions [[TypeFamilies], [MultiParamTypeClasses]]
  --
  requireExtensions :: String -> [[Extension]] -> Q ()
  requireExtensions what extLists = do
    -- Taken from the persistent library
    required <- filterM (fmap (not . or) . traverse isExtEnabled) extLists
    case mapMaybe listToMaybe required of
      [] -> pure ()
      [extension] -> fail $ mconcat
        [ "Generating " ++ what ++ " requires the "
        , show extension
        , " language extension. Please enable it by copy/pasting this line to the top of your file:\n\n"
        , extensionToPragma extension
        , "\n\nTo enable it in a GHCi session, use the following command:\n\n"
        , ":seti -X" ++ show extension
        ]
      extensions -> fail $ mconcat
        [ "Generating " ++ what ++ " requires the following language extensions:\n\n"
        , L.intercalate "\n" (map (("- " ++) . show) extensions)
        , "\n\nPlease enable the extensions by copy/pasting these lines into the top of your file:\n\n"
        , L.intercalate "\n" (map extensionToPragma extensions)
        , "\n\nTo enable them in a GHCi session, use the following command:\n\n"
        , ":seti " ++ unwords (map (("-X" ++) . show) extensions)
        ]
    where
      extensionToPragma ext = "{-# LANGUAGE " ++ show ext ++ " #-}"
  
  requireExtensionsForLabels :: Q ()
  requireExtensionsForLabels = requireExtensions "LabelOptic instances"
    [ [DataKinds]
    , [FlexibleInstances]
    , [MultiParamTypeClasses]
    , [TypeFamilies, GADTs]
    , [UndecidableInstances]
    ]
  
  requireExtensionsForFields :: Q ()
  requireExtensionsForFields = requireExtensions "field optics"
    [ [FlexibleInstances]
    , [FunctionalDependencies]
    ]
  
  ------------------------------------------------------------------------
  -- Support for generating inline pragmas
  ------------------------------------------------------------------------
  
  inlinePragma :: Name -> [DecQ]
  inlinePragma methodName = [pragInlD methodName Inline FunLike AllPhases]
  