{-# LANGUAGE LambdaCase #-}

module Language.Haskell.TH.TypeInterpreter.Import
    ( fromType
    , fromName )
where

import Control.Monad.State

import qualified Data.Map   as Map
import           Data.Maybe (mapMaybe)

import Language.Haskell.TH

import Language.Haskell.TH.TypeInterpreter.Expression
import Language.Haskell.TH.TypeInterpreter.Names

-- | Importer monad
type Importer = StateT (Map.Map Name TypeExp) Q

-- | Register a name with the name cache.
registerName :: Name -> TypeExp -> Importer ()
registerName name exp = modify (Map.insert name exp)

-- | Construct a type expression using the equations of a type family.
fromTypeFamily :: Name -> [TySynEqn] -> Importer TypeExp
fromTypeFamily familyName synonymEquations = do
    -- Register a dummy that we can easily substitute later
    registerName familyName (Variable familyName)

    equations <- traverse mkEquation synonymEquations
    let result = substitute familyName result (Function equations)

    result <$ registerName familyName result
    where
        mkEquation (TySynEqn patterns body) =
            TypeEquation <$> traverse fromTypeOnly patterns <*> fromTypeOnly body

-- | Get the type expression for a 'Type'. Attempts to reduce the resulting type expression.
fromType :: Type -> Q TypeExp
fromType typ = reduce <$> evalStateT (fromTypeOnly typ) Map.empty

-- | Get the type expression for a 'Type'.
fromTypeOnly :: Type -> Importer TypeExp
fromTypeOnly = \case
    AppT f x          -> Apply <$> fromTypeOnly f <*> fromTypeOnly x
    ArrowT            -> pure (Atom (Name arrowTypeName))
    ConstraintT       -> pure (Atom (Name constraintTypeName))
    ConT n            -> fromNameOnly n
    EqualityT         -> pure (Atom (Name equalityTypeName))
    ForallT _ _ t     -> fromTypeOnly t
    InfixT l n r      -> Apply <$> (Apply <$> fromNameOnly n <*> fromTypeOnly l) <*> fromTypeOnly r
    ListT             -> pure (Atom (Name listTypeName))
    LitT (NumTyLit n) -> pure (Atom (Integer n))
    LitT (StrTyLit s) -> pure (Atom (String s))
    ParensT t         -> fromTypeOnly t
    PromotedConsT     -> pure (Atom (PromotedName consName))
    PromotedNilT      -> pure (Atom (PromotedName nilName))
    PromotedT n       -> pure (Atom (PromotedName n))
    PromotedTupleT n  -> pure (Atom (PromotedName (tupleDataName n)))
    SigT t _          -> fromTypeOnly t
    StarT             -> pure (Atom (Name starTypeName))
    TupleT n          -> pure (Atom (Name (tupleTypeName n)))
    UInfixT l n r     -> Apply <$> (Apply <$> fromNameOnly n <*> fromTypeOnly l) <*> fromTypeOnly r
    UnboxedSumT n     -> pure (Atom (Name (unboxedSumTypeName n)))
    UnboxedTupleT n   -> pure (Atom (Name (unboxedTupleTypeName n)))
    VarT n            -> pure (Variable n)
    WildCardT         -> Variable <$> lift (newName "wildCard")

-- | Get the type expression for a 'Name'. Attempts to reduce the resulting type expression.
fromName :: Name -> Q TypeExp
fromName name = reduce <$> evalStateT (fromNameOnly name) Map.empty

-- | Get the type expression for a 'Name'.
fromNameOnly :: Name -> Importer TypeExp
fromNameOnly name =
    gets (Map.lookup name) >>= \case
        Just x -> pure x
        Nothing -> do
            info <- lift (reify name)
            fromInfo info
    where
        extractName (PlainTV name)    = name
        extractName (KindedTV name _) = name

        synonymEquation (TySynInstD _ equation) = Just equation
        synonymEquation _                       = Nothing

        foldTypeSynonym body var = Function [TypeEquation [Variable (extractName var)] body]

        fromInfo = \case
            TyConI (TySynD _ vars body) ->
                (\ body -> foldl foldTypeSynonym body vars) <$> fromTypeOnly body

            TyConI {} ->
                pure (Atom (Name name))

            FamilyI (OpenTypeFamilyD _) instances ->
                fromTypeFamily name (mapMaybe synonymEquation (reverse instances))

            FamilyI (ClosedTypeFamilyD _ equations) _ ->
                fromTypeFamily name equations

            PrimTyConI {} ->
                pure (Atom (Name name))

            TyVarI {} ->
                pure (Variable name)

            ClassI {} ->
                pure (Atom (Name name))

            -- The following errors should generally only occur when the user misuses this function.
            ClassOpI {} -> fail ("Cannot turn class method " ++ show name ++ " into a TypeExp")
            FamilyI {}  -> fail ("Cannot turn family " ++ show name ++ " into a TypeExp")
            DataConI {} -> fail ("Cannot turn data constructor " ++ show name ++ " into TypeExp")
            PatSynI {}  -> fail ("Cannot turn pattern synonym " ++ show name ++ " into TypeExp")
            VarI {}     -> fail ("Cannot turn variable " ++ show name ++ " into TypeExp")