content: {-# LANGUAGE DataKinds #-}
  {-# LANGUAGE FlexibleContexts #-}
  {-# LANGUAGE LambdaCase #-}
  {-# LANGUAGE NamedFieldPuns #-}
  {-# LANGUAGE TemplateHaskell #-}
  {-# LANGUAGE TypeOperators #-}
  {-# LANGUAGE ViewPatterns #-}
  
  {-|
  Module      :  Data.Aeson.Schema.TH.Utils
  Maintainer  :  Brandon Chinn <brandonchinn178@gmail.com>
  Stability   :  experimental
  Portability :  portable
  -}
  module Data.Aeson.Schema.TH.Utils (
    reifySchema,
    lookupSchema,
    loadSchema,
    resolveSchemaType,
    schemaVToTypeQ,
    schemaTypeVToTypeQ,
  ) where
  
  import Control.Applicative (empty)
  import Control.Monad (forM)
  import Data.Bifunctor (bimap)
  import Language.Haskell.TH
  
  import Data.Aeson.Schema.Internal (Object)
  import Data.Aeson.Schema.Key (SchemaKey' (..), SchemaKeyV)
  import Data.Aeson.Schema.Type (
    Schema' (..),
    SchemaObjectMapV,
    SchemaType' (..),
    SchemaTypeV,
    SchemaV,
    fromSchemaV,
    toSchemaObjectV,
   )
  import Data.Aeson.Schema.Utils.Invariant (unreachable)
  import Data.Aeson.Schema.Utils.NameLike (NameLike (..), resolveName)
  
  {- Loading schema from TH -}
  
  reifySchema :: String -> Q SchemaV
  reifySchema name = lookupSchema (NameRef name) >>= loadSchema
  
  data ReifiedSchema = ReifiedSchema
    { reifiedSchemaName :: Name
    , reifiedSchemaType :: TypeWithoutKinds
    }
  
  -- | Look up a schema with the given name. Errors if the name doesn't exist or if the name does
  --  not refer to a schema.
  lookupSchema :: NameLike -> Q ReifiedSchema
  lookupSchema nameLike = do
    name <- lookupSchemaName nameLike
    ReifiedSchema name <$> reifySchemaType name
    where
      lookupSchemaName = \case
        NameRef name -> lookupTypeName name >>= maybe (fail $ "Unknown schema: " ++ name) return
        NameTH name -> return name
  
      reifySchemaType :: Name -> Q TypeWithoutKinds
      reifySchemaType schemaName =
        reify schemaName >>= \case
          TyConI (TySynD _ _ (stripKinds -> ty))
            -- `type MySchema = 'Schema '[ ... ]`
            | isPromotedSchema ty ->
                return ty
            -- `type MySchema = Object ('Schema '[ ... ])`
            | Just inner <- unwrapObject ty
            , isPromotedSchema inner ->
                return inner
            -- `type MySchema = Object OtherSchema`
            | Just (ConT schemaName') <- unwrapObject ty ->
                reifySchemaType schemaName'
          _ -> fail $ "'" ++ show schemaName ++ "' is not a Schema"
  
      -- If the given type is of the format `Object a`, return `a`.
      unwrapObject :: TypeWithoutKinds -> Maybe TypeWithoutKinds
      unwrapObject = \case
        AppT (ConT name) inner | name == ''Object -> Just inner
        _ -> Nothing
  
      -- Return True if the given type is of the format: 'Schema '[ ... ]
      isPromotedSchema :: TypeWithoutKinds -> Bool
      isPromotedSchema = \case
        AppT (PromotedT name) _ | name == 'Schema -> True
        _ -> False
  
  loadSchema :: ReifiedSchema -> Q SchemaV
  loadSchema ReifiedSchema{reifiedSchemaType} =
    maybe (fail $ "Could not parse schema: " ++ show reifiedSchemaType) return $ parseSchema reifiedSchemaType
    where
      -- should be the inverse of schemaVToTypeQ
      parseSchema :: TypeWithoutKinds -> Maybe SchemaV
      parseSchema ty = do
        schemaObjectType <- case ty of
          AppT (PromotedT name) schemaType | name == 'Schema -> return schemaType
          _ -> empty
  
        Schema <$> parseSchemaObjectMap schemaObjectType
  
      -- should be the inverse of schemaObjectMapVToTypeQ
      parseSchemaObjectMap :: TypeWithoutKinds -> Maybe SchemaObjectMapV
      parseSchemaObjectMap schemaObjectType = do
        schemaObjectListOfPairs <- mapM typeToPair =<< typeToList schemaObjectType
        forM schemaObjectListOfPairs $ \(schemaKeyType, schemaTypeType) -> do
          schemaKey <- parseSchemaKey schemaKeyType
          schemaType <- parseSchemaType schemaTypeType
          return (schemaKey, schemaType)
  
      -- should be the inverse of schemaKeyVToTypeQ
      parseSchemaKey :: TypeWithoutKinds -> Maybe SchemaKeyV
      parseSchemaKey = \case
        AppT (PromotedT ty) (LitT (StrTyLit key))
          | ty == 'NormalKey -> return $ NormalKey key
          | ty == 'PhantomKey -> return $ PhantomKey key
        _ -> empty
  
      -- should be the inverse of schemaTypeVToTypeQ
      parseSchemaType :: TypeWithoutKinds -> Maybe SchemaTypeV
      parseSchemaType = \case
        AppT (PromotedT name) (ConT inner)
          | name == 'SchemaScalar -> return $ SchemaScalar $ NameTH inner
        AppT (PromotedT name) inner
          | name == 'SchemaMaybe -> SchemaMaybe <$> parseSchemaType inner
          | name == 'SchemaTry -> SchemaTry <$> parseSchemaType inner
          | name == 'SchemaList -> SchemaList <$> parseSchemaType inner
          | name == 'SchemaUnion -> do
              schemas <- typeToList inner
              SchemaUnion <$> mapM parseSchemaType schemas
          | name == 'SchemaObject -> SchemaObject <$> parseSchemaObjectMap inner
        AppT (PromotedT name) (AppT (PromotedT right) (ConT inner))
          | name == 'SchemaInclude
          , right == 'Right ->
              return $ SchemaInclude $ Left $ NameTH inner
        _ -> empty
  
  -- | Resolve SchemaInclude, if present. (Not recursive)
  resolveSchemaType :: SchemaTypeV -> Q SchemaTypeV
  resolveSchemaType = \case
    SchemaInclude (Left name) -> fmap toSchemaObjectV . loadSchema =<< lookupSchema name
    SchemaInclude (Right _) -> unreachable "Found 'SchemaInclude Right' when resolving schema type"
    schemaType -> pure schemaType
  
  {- Splicing schema into TH -}
  
  schemaVToTypeQ :: SchemaV -> TypeQ
  schemaVToTypeQ = appT [t|'Schema|] . schemaObjectMapVToTypeQ . fromSchemaV
  
  schemaObjectMapVToTypeQ :: SchemaObjectMapV -> TypeQ
  schemaObjectMapVToTypeQ = promotedListT . map schemaObjectPairVToTypeQ
    where
      schemaObjectPairVToTypeQ :: (SchemaKeyV, SchemaTypeV) -> TypeQ
      schemaObjectPairVToTypeQ = promotedPairT . bimap schemaKeyVToTypeQ schemaTypeVToTypeQ
  
      schemaKeyVToTypeQ :: SchemaKeyV -> TypeQ
      schemaKeyVToTypeQ = \case
        NormalKey key -> [t|'NormalKey $(litT $ strTyLit key)|]
        PhantomKey key -> [t|'PhantomKey $(litT $ strTyLit key)|]
  
  schemaTypeVToTypeQ :: SchemaTypeV -> TypeQ
  schemaTypeVToTypeQ = \case
    SchemaScalar name -> [t|'SchemaScalar $(resolveName name >>= conT)|]
    SchemaMaybe inner -> [t|'SchemaMaybe $(schemaTypeVToTypeQ inner)|]
    SchemaTry inner -> [t|'SchemaTry $(schemaTypeVToTypeQ inner)|]
    SchemaList inner -> [t|'SchemaList $(schemaTypeVToTypeQ inner)|]
    SchemaUnion schemas -> [t|'SchemaUnion $(promotedListT $ map schemaTypeVToTypeQ schemas)|]
    SchemaObject pairs -> [t|'SchemaObject $(schemaObjectMapVToTypeQ pairs)|]
    SchemaInclude (Left name) -> [t|'SchemaInclude ('Right $(conT . reifiedSchemaName =<< lookupSchema name))|]
    SchemaInclude (Right _) -> unreachable "Found 'SchemaInclude Right' when converting to TypeQ"
  
  {- TH utilities -}
  
  -- | Same as 'Type' except without any kind signatures or applications at any depth.
  --
  --  Provides no actual guarantees. The caller is responsible for making sure the value
  --  has been run through 'stripKinds' at one point.
  type TypeWithoutKinds = Type
  
  {- FOURMOLU_DISABLE -}
  
  -- | Recursively strip all kind signatures and applications.
  stripKinds :: Type -> TypeWithoutKinds
  stripKinds ty =
    case ty of
      -- cases that strip + recurse
      SigT ty1 _ -> stripKinds ty1
      AppKindT ty1 _ -> stripKinds ty1
  
      -- cases that recurse
      ForallT tyVars ctx ty1 -> ForallT tyVars ctx (stripKinds ty1)
      ForallVisT tyVars ty1 -> ForallVisT tyVars (stripKinds ty1)
      AppT ty1 ty2 -> AppT (stripKinds ty1) (stripKinds ty2)
      InfixT ty1 name ty2 -> InfixT (stripKinds ty1) name (stripKinds ty2)
      UInfixT ty1 name ty2 -> UInfixT (stripKinds ty1) name (stripKinds ty2)
      PromotedInfixT ty1 name ty2 -> PromotedInfixT (stripKinds ty1) name (stripKinds ty2)
      PromotedUInfixT ty1 name ty2 -> PromotedUInfixT (stripKinds ty1) name (stripKinds ty2)
      ParensT ty1 -> ParensT (stripKinds ty1)
      ImplicitParamT str ty1 -> ImplicitParamT str (stripKinds ty1)
  
      -- base cases
      VarT _ -> ty
      ConT _ -> ty
      PromotedT _ -> ty
      TupleT _ -> ty
      UnboxedTupleT _ -> ty
      UnboxedSumT _ -> ty
      ArrowT -> ty
      MulArrowT -> ty
      EqualityT -> ty
      ListT -> ty
      PromotedTupleT _ -> ty
      PromotedNilT -> ty
      PromotedConsT -> ty
      StarT -> ty
      ConstraintT -> ty
      LitT _ -> ty
      WildCardT -> ty
  
  {- FOURMOLU_ENABLE -}
  
  typeToList :: TypeWithoutKinds -> Maybe [TypeWithoutKinds]
  typeToList = \case
    PromotedNilT -> Just []
    AppT (AppT PromotedConsT x) xs -> (x :) <$> typeToList xs
    _ -> Nothing
  
  typeToPair :: TypeWithoutKinds -> Maybe (TypeWithoutKinds, TypeWithoutKinds)
  typeToPair = \case
    AppT (AppT (PromotedTupleT 2) a) b -> Just (a, b)
    _ -> Nothing
  
  promotedListT :: [TypeQ] -> TypeQ
  promotedListT = foldr consT promotedNilT
    where
      -- nb. https://stackoverflow.com/a/34457936
      consT x xs = appT (appT promotedConsT x) xs
  
  promotedPairT :: (TypeQ, TypeQ) -> TypeQ
  promotedPairT (a, b) = [t|'($a, $b)|]
  