{-# LANGUAGE LambdaCase #-}

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

import Control.Monad (zipWithM)

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

import Language.Haskell.TH hiding (match)

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

-- | Construct a type expression for a type family.
fromTypeFamily :: Name -> Int -> [TySynEqn] -> Q TypeExp
fromTypeFamily name numParams equations =
    mkFamilyExp <$> traverse mkMapper equations
    where
        substituteBody body substitutions =
            substituteAll (Map.unions substitutions) body

        mkMapper (TySynEqn patterns body) = do
            patterns <- traverse fromTypeOnly patterns
            body <- fromTypeOnly body
            pure (fmap (substituteBody body) . zipWithM match patterns)

        mkFamilyExp mappers =
            familyExp numParams $ \ inputs ->
                case mapMaybe ($ inputs) mappers of
                    []    -> foldl Apply (Atom (Name name)) inputs
                    r : _ -> r

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

-- | Get the type expression for a 'Type'.
fromTypeOnly :: Type -> Q 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 <$> newName "wildCard"

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

-- | Get the type expression for a 'Name'.
fromNameOnly :: Name -> Q TypeExp
fromNameOnly name =
    reify name >>= \case
        TyConI (TySynD _ vars body) -> do
            body <- fromTypeOnly body
            pure (foldr Synonym body (map extractName vars))

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

        FamilyI (OpenTypeFamilyD (TypeFamilyHead _ vars _ _)) instances ->
            fromTypeFamily name (length vars) (mapMaybe synonymEquation instances)

        FamilyI (ClosedTypeFamilyD (TypeFamilyHead _ vars _ _) equations) _ ->
            fromTypeFamily name (length vars) 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")
    where
        extractName (PlainTV name)    = name
        extractName (KindedTV name _) = name

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