{-# LANGUAGE LambdaCase   #-}
{-# LANGUAGE ViewPatterns #-}

module Language.Haskell.TH.TypeInterpreter.Expression
    ( TypeAtom (..)
    , TypeExp (..)
    , substitute
    , substituteAll
    , reduce
    , match
    , familyExp )
where

import Control.Monad

import qualified Data.Map as Map

import Language.Haskell.TH (Name)

-- | Type atom
data TypeAtom
    = Integer Integer
    | String String
    | Name Name
    | PromotedName Name
    deriving Eq

instance Show TypeAtom where
    show (Integer i)      = show i
    show (String s)       = show s
    show (Name n)         = show n
    show (PromotedName n) = '\'' : show n

-- | Type expression
data TypeExp
    = Atom TypeAtom
    | Apply TypeExp TypeExp
    | Variable Name
    | Synonym Name TypeExp
    | Family (TypeExp -> TypeExp)

instance Show TypeExp where
    showsPrec prec = \case
        Atom atom ->
            showsPrec prec atom

        Variable name ->
            (show name ++)

        Apply fun param ->
            showParen (prec >= 10) $ \ tail ->
                showsPrec 10 fun (' ' : showsPrec 10 param tail)

        Synonym varName body ->
            showParen (prec >= 10) $ \ tail ->
                'λ' : showsPrec 0 varName ('.' : ' ' : showsPrec 0 body tail)

        Family _ ->
            showParen (prec >= 10) ("λ?. ?" ++)

-- | @substitute name typ exp@ replaces all occurences of @name@ in @exp@ with @typ@.
substitute :: Name -> TypeExp -> TypeExp -> TypeExp
substitute name typ =
    subst
    where
        subst = \case
            Variable varName
                | varName == name -> typ

            Apply fun param ->
                Apply (subst fun) (subst param)

            Synonym subName body
                | subName == name -> subst body
                | otherwise       -> Synonym subName (subst body)

            t -> t

-- | Just like 'substitute' but for more variables.
substituteAll :: Map.Map Name TypeExp -> TypeExp -> TypeExp
substituteAll mapping exp = Map.foldrWithKey substitute exp mapping

-- | Try to reduce the given type expression as much as possible.
reduce :: TypeExp -> TypeExp
reduce = \case
    Apply (reduce -> f) (reduce -> x)
        | Synonym n b <- f -> reduce (substitute n x b)
        | Family g <- f    -> reduce (g x)
        | otherwise        -> Apply f x

    Synonym n b -> Synonym n (reduce b)

    Family f -> Family (reduce . f)

    t -> t

-- | @match pattern input@ pattern matches @input@ against the given @pattern@.
match :: TypeExp -> TypeExp -> Maybe (Map.Map Name TypeExp)
match pattern input =
    matchOnly (reduce pattern) (reduce input)
    where
        matchOnly (Variable n) v           = Just (Map.singleton n v)
        matchOnly (Apply f x ) (Apply g y) = Map.union <$> matchOnly f g <*> matchOnly x y
        matchOnly (Atom l    ) (Atom r   ) = Map.empty <$ guard (l == r)
        matchOnly _            _           = Nothing

-- | @familyExp n impl@ creates a family expression with @n@ parameters and implementation @impl@.
familyExp :: Int -> ([TypeExp] -> TypeExp) -> TypeExp
familyExp n f
    | n <= 0    = f []
    | otherwise = Family (\ t -> familyExp (n - 1) (\ ts -> f (t : ts)))