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

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

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.Quantification

import Debug.Trace (trace)

dbg :: Show a => [Char] -> a -> a
dbg msg x = trace (msg ++ " = " ++ show x) x

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

deriving instance
 ( Show src
 , Show_TyConst cs
 , Proj_TyConst cs (->)
 , Proj_TyConst cs (#>)
 ) => Show (Error_TeApp src cs)

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

-- | Return the next free 'Var' considering all the 'KType's of given 'BinTree'.
freeTyVarTreeKT :: BinTree (KType src 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 cs.
    Proj_TyConst cs (->)
 => Proj_TyConst cs (#>)
 => Proj_TyCon cs
 => Inj_Source (EType src cs) src
 => Show_TyConst cs
 => KTerm src cs
 -> KTerm src cs
 -> Either (Error_TeApp src cs) (KTerm src cs)
teApp (KT fun) (KT arg) = go fun arg
	where
	go
	 :: Term src cs f
	 -> Term src cs a
	 -> Either (Error_TeApp src cs) (KTerm src cs)
	go f@TyConst{} _a = Left $ Error_TeApp_Term_missing (EType $ f `source` EType fun)
	go (TyApp _sb (TyApp _sa (TyConst _src' a2b) a2b_a) a2b_b) a
	 | Just KRefl <- projKiTyConst @_ @(->) a2b
	 = case eqTy a2b_a a of
		 Just Refl -> Right $ KT a2b_b
		 Nothing -> Left $ Error_TeApp_Type_mismatch (EType a2b_a) (EType a)
	go f@TyApp{} _a = Left $ Error_TeApp_Term_not_applicable (EType $ f `source` EType fun)
	go f@TyVar{} _a = Left $ Error_TeApp_Term_missing (EType $ f `source` EType 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 (TeSym te_fun) (TyApp _sb (TyApp _sa (TyConst _src' a2b) a2b_a) a2b_b)) a
	 | Just KRefl <- projKiTyConst @_ @(->) a2b
	 = case a of
		 Term (TeSym te_arg) a' ->
			case eqTy a2b_a a' of
			 Just Refl -> Right $ KT $ Term (TeSym $ sym_app te_fun te_arg) a2b_b
			 Nothing -> Left $ Error_TeApp_Type_mismatch (EType a2b_a) (EType a')
		 _ -> Left $ Error_TeApp_Term_missing (EType $ a `source` EType arg)
	go (Term _te x) _a = Left $ Error_TeApp_Term_not_applicable (EType $ x `source` EType fun)

-- | Monomorphize a 'BinTree' of 'KT'.
monoTreeKT :: BinTree (KT src cs k) -> BinTree (KT src 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 cs) src
 => Show_TyConst cs
 => Proj_TyConst cs (->)
 => Proj_TyConst cs (#>)
 => Proj_TyCon cs
 => BinTree (KTerm src cs)
 -> Either (Error_TeApp src cs) (MapTyVar src cs, Maybe (KTerm src 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 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 cs h
 -> Maybe (TeSym h)
evalTe (Term te _ty) = Just te
evalTe TyQuant{}     = Nothing
evalTe TyConst{}     = Nothing
evalTe TyVar{}       = Nothing
evalTe TyApp{}       = Nothing