{-# LANGUAGE UndecidableInstances #-}
{-# OPTIONS_GHC -fno-warn-orphans #-}
module Language.Symantic.Typing.Term where

import Control.Arrow (left)
import Data.List (foldl')
import qualified Control.Monad.Trans.State.Strict as SS
import qualified Data.Map.Strict as Map

import Language.Symantic.Helper.Data.Type.Equality
import Language.Symantic.Parsing
import Language.Symantic.Typing.Kind
import Language.Symantic.Typing.Constant
import Language.Symantic.Typing.Constraint
import Language.Symantic.Typing.Type
import Language.Symantic.Typing.Show ()
import Language.Symantic.Typing.Quantification

-- * Type 'Error_TeApp'
data Error_TeApp src ss cs
 =   Error_TeApp_FunArg (EType src '[] ss cs)
     -- ^ Cannot get the expected 'Type' of the argument of the function.
 |   Error_TeApp_MGU (Error_MGU src ss cs)
     -- ^ Cannot unify the expected 'Type' of the argument of the function,
     -- with the 'Type' of the argument.
 |   Error_TeApp_Term_missing (EType src '[] ss cs)
     -- ^ Cannot find a 'TeTerm', either in the function or in the argument.
 {-
 |   Error_TeApp_TyQuant_remaining (EType src ctx ss cs)
     -- ^ A 'TyQuant' remains to be instantiated.
 -}
 |   Error_TeApp_Term_not_applicable (EType src '[] ss cs)
 |   Error_TeApp_Type_mismatch (EType src '[] ss cs) (EType src '[] ss cs)

deriving instance
 ( Eq src
 , Eq_Token ss
 , Inj_Source (EType src '[] ss cs) src
 ) => Eq (Error_TeApp src ss cs)
deriving instance
 ( Show src
 , Show_TyConst cs
 , Fixity_TyConst cs
 ) => Show (Error_TeApp src ss cs)

instance Inj_Error (Error_TeApp src ss cs) (Error_TeApp src ss cs) where
	inj_Error = id
instance Inj_Error (Error_MGU src ss cs) (Error_TeApp src ss cs) where
	inj_Error = Error_TeApp_MGU

-- | Return the next free 'Var' considering all the 'KType's of given 'BinTree'.
freeTyVarTreeKT :: BinTree (KType src ctx ss cs k) -> Var
freeTyVarTreeKT t = foldr max 0 $ ofKT freeTyVar <$> t

-- | Term application: apply second given 'KTerm' to the first,
-- applying embedded 'TeSym's if possible.
teApp
 :: forall src ctx ss cs.
    Proj_TyConst cs (->)
 => Proj_TyConst cs (#>)
 => Proj_TyCon cs
 => Inj_Source (EType src '[] ss cs) src
 => Show_TyConst cs
 => KTerm src ctx ss cs
 -> KTerm src ctx ss cs
 -> Either (Error_TeApp src ss cs) (KTerm src ctx ss cs)
teApp (KT fun) (KT arg) = go fun arg
	where
	go
	 :: Term src ctx ss cs f
	 -> Term src ctx ss cs a
	 -> Either (Error_TeApp src ss cs)
	           (KTerm src ctx ss cs)
	go f@TyConst{} _a = Left $ Error_TeApp_Term_missing (EType $ unCtxTe f `source` EType (unCtxTe fun))
	go (TyApp _sb (TyApp _sa (TyConst _src' a2b) a2b_a) a2b_b) a
	 | Just KRefl <- proj_TyConstKi @_ @(->) a2b
	 = case eqTy a2b_a a of
		 Just Refl -> Right $ KT a2b_b
		 Nothing -> Left $ Error_TeApp_Type_mismatch (EType $ unCtxTe a2b_a) (EType $ unCtxTe a)
	go f@TyApp{} _a = Left $ Error_TeApp_Term_not_applicable (EType $ unCtxTe f `source` EType (unCtxTe fun))
	go f@TyVar{} _a = Left $ Error_TeApp_Term_missing (EType $ unCtxTe f `source` EType (unCtxTe fun))
	go t@(TyQuant src n kv f) a =
		let v = freeTyVar t in
		case f $ TyVar src kv v of
		 KT f' -> bindKT (polyTyVar n kv v) <$> go f' a
	go (Term (TyApp _sb (TyApp _sa (TyConst _src' a2b) a2b_a) a2b_b) (TeSym te_fun)) a
	 | Just KRefl <- proj_TyConstKi @_ @(->) a2b
	 = case a of
		 Term a' (TeSym te_arg) ->
			case eqTy a2b_a a' of
			 Just Refl -> Right $ KT $ Term a2b_b (TeSym $ \ctx -> te_fun ctx `app` te_arg ctx)
			 Nothing -> Left $ Error_TeApp_Type_mismatch (EType $ unCtxTe a2b_a) (EType $ unCtxTe a')
		 _ -> Left $ Error_TeApp_Term_missing (EType $ unCtxTe a `source` EType (unCtxTe arg))
	go (Term x _te) _a = Left $ Error_TeApp_Term_not_applicable (EType $ unCtxTe x `source` EType (unCtxTe fun))

-- | Monomorphize a 'BinTree' of 'KT'.
monoTreeKT :: BinTree (KT src ctx ss cs k) -> BinTree (KT src ctx ss cs k)
monoTreeKT t =
	(`SS.evalState` freeTyVarTreeKT t) $
	traverse (monoTy `ofKT`) t

-- | Fold given 'BinTree' of 'KTerm' in order to compute
-- a resulting 'KTerm', if possible.
evalTreeKT
 :: Inj_Source (EType src '[] ss cs) src
 => Show_TyConst cs
 => Proj_TyConst cs (->)
 => Proj_TyConst cs (#>)
 => Proj_TyCon cs
 => BinTree (KTerm src ctx ss cs)
 -> Either (Error_TeApp src ss cs)
           ( MapTyVar src ctx ss cs
           , Maybe (KTerm src ctx ss cs) )
evalTreeKT t =
	foldl' (\acc (KT arg) -> do
		(mgu, may_fun) <- acc
		case may_fun of
		 Nothing -> Right (mgu, Just (KT arg))
		 Just (KT fun) ->
			case unTyFun fun of
			 Nothing -> inj_Error $ Error_TeApp_FunArg $ EType $ unCtxTe fun
			 Just (KT fun_arg, KT _fun_res) -> do
				mgu' <- inj_Error `left` mguTy mgu fun_arg arg
				let fun' = substTyVar mgu' fun
				let arg' = substTyVar mgu' arg
				KT res <- teApp fun' arg'
				Right (mgu', Just (KT res)))
	 (Right (Map.empty, Nothing)) $
	monoTreeKT t

evalTe
 :: Source src
 => Proj_TyCon cs
 => Type src ctx ss cs h
 -> Maybe (TeSym ctx ss h)
evalTe (Term _ty te) = Just te
evalTe TyQuant{}     = Nothing
evalTe TyConst{}     = Nothing
evalTe TyVar{}       = Nothing
evalTe TyFam{}       = Nothing
evalTe TyApp{}       = Nothing