{-# LANGUAGE LambdaCase #-}

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

import Control.Monad (zipWithM)

import           Data.Function (fix)
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 :: Int -> [TySynEqn] -> Q TypeExp
fromTypeFamily numParams equations =
    fix . mkFamilyExp <$> traverse mkMapper equations
    where
        substituteBody body substitutions =
            substituteAll (Map.unions substitutions) body

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

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

-- | Get the type expression for a 'Type'.
fromType :: Type -> Q TypeExp
fromType = fmap reduce . \case
    AppT f x          -> Apply <$> fromType f <*> fromType x
    ArrowT            -> pure (Atom (Name arrowTypeName))
    ConstraintT       -> pure (Atom (Name constraintTypeName))
    ConT n            -> fromName n
    EqualityT         -> pure (Atom (Name equalityTypeName))
    ForallT _ _ t     -> fromType t
    InfixT l n r      -> Apply <$> (Apply <$> fromName n <*> fromType l) <*> fromType r
    ListT             -> pure (Atom (Name listTypeName))
    LitT (NumTyLit n) -> pure (Atom (Integer n))
    LitT (StrTyLit s) -> pure (Atom (String s))
    ParensT t         -> fromType 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 _          -> fromType t
    StarT             -> pure (Atom (Name starTypeName))
    TupleT n          -> pure (Atom (Name (tupleTypeName n)))
    UInfixT l n r     -> Apply <$> (Apply <$> fromName n <*> fromType l) <*> fromType 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'.
fromName :: Name -> Q TypeExp
fromName name =
    reify name >>= fmap reduce . \case
        TyConI (TySynD _ vars body) -> do
            body <- fromType body
            pure (foldr Synonym body (map extractName vars))

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

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

        FamilyI (ClosedTypeFamilyD (TypeFamilyHead _ vars _ _) equations) _ ->
            fromTypeFamily (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