{-# LANGUAGE ConstraintKinds #-}
{-# LANGUAGE GADTs #-}
{-# LANGUAGE PolyKinds #-}
{-# LANGUAGE UndecidableInstances #-}
{-# OPTIONS_GHC -fno-warn-orphans #-}
module Language.Symantic.Compiling.Term
 ( module Language.Symantic.Compiling.Term
 , module Language.Symantic.Compiling.Term.Grammar
 ) where

import Data.Proxy (Proxy(..))
import qualified Data.Text as Text
import Data.Type.Equality ((:~:)(..))
import GHC.Exts (Constraint)
import qualified Data.Kind as K

import Language.Symantic.Helper.Data.Type.List
import Language.Symantic.Parsing
import Language.Symantic.Typing

import Language.Symantic.Compiling.Term.Grammar

-- * Type 'Term'
-- | An open term (i.e. depending on a 'TeCtx').
-- The data type wraps a universal quantification
-- over an interpreter @term@
-- qualified by the symantics of a term.
--
-- Moreover the term is abstracted by a /lambda term context/
-- built top-down by 'compile',
-- to enable a /Higher-Order Abstract Syntax/ (HOAS)
-- for /lambda abstractions/ ('lam').
--
-- This data type is used to keep a parsed term polymorphic enough
-- to stay interpretable by different interpreters.
--
-- * @(@'Sym_of_Ifaces'@ is term)@
-- is needed when a symantic method includes a polymorphic type
-- and thus calls: 'compile'.
--
-- * @(@'Sym_of_Ifaces'@ ls term)@ and @(@'Sym_of_Ifaces'@ rs term)@
-- make a zipper needed to be able to write the recursing 'CompileR' instance.
--
-- * @(@'Sym_of_Iface'@ (@'Proxy'@ (->)) term)@
-- is needed to handle partially applied functions.
data Term hs h is ls rs
 =   Term
 (forall term. ( Sym_of_Ifaces is term
               , Sym_of_Ifaces ls term
               , Sym_of_Ifaces rs term
               , Sym_of_Iface (Proxy (->)) term
               ) => TeCtx term hs -> term h)

-- ** Type 'ETerm'
-- | Existential 'Term', with its 'Type'.
data ETerm src hs cs is
 = forall (h:: *). ETerm (Type src cs h) (Term hs h is '[] is)

-- ** Type 'TermClosed'
-- | Closed 'Term'.
data TermClosed is h
 =   TermClosed
 (forall term. ( Sym_of_Ifaces is term
               , Sym_of_Iface (Proxy (->)) term
               ) => term h)

-- * Type 'ETermClosed'
-- | Existential 'TermClosed', with its 'Type'.
data ETermClosed src cs is
 = forall (h:: *). ETermClosed (Type src cs h) (TermClosed is h)

-- * Type 'Compile'
-- | Convenient class wrapping 'CompileR' to initiate its recursion.
class CompileR cs is '[] is => Compile cs is where
	compile :: AST_Term src is -> Compiler src hs ret cs is '[] is
instance CompileR cs is '[] is => Compile cs is where
	compile (TokTree0 (EToken tok))            = compileR tok Nothing
	compile (TokTree0 (EToken x) `TokTree1` y) = compileR x (Just y)
	{-
	compile (tok_fun `TokTree1` tok_arg) = \ctx k ->
		compile tok_fun ctx $ \ty_fun (Term fun) ->
		compile tok_arg ctx $ \ty_arg (Term arg) ->
		check_TyEq2 (ty @(->)) (At (Just tok_fun) ty_fun) $ \Refl ty_param ty_res ->
		check_TyEq
		 (At (Just tok_fun) ty_param)
		 (At (Just tok_arg) ty_arg) $ \Refl ->
		k ty_res $ Term $
		 \c -> fun c .$ arg c
	-}

-- ** Type 'CompilerWithTyCtx'
-- | Convenient type synonym defining a term compiler
-- with independant 'TyCtx' and 'TeCtx'.
-- 
-- This is the most general term compiler.
-- 
-- Useful to let the compiled term use
-- terms whose type is known before 'compile' is run
-- but actually available only after.
type CompilerWithTyCtx src hs_ty hs_te ret cs is ls rs
 = TyCtx TeName cs hs_ty
 -- ^ The bound variables in scope and their types:
 -- built top-down in the heterogeneous list @hs@,
 -- from the closest including /lambda abstraction/ to the farest.
 -> ( forall h.
       Type src cs (h:: *)
    -> Term hs_te h is ls rs
    -> Either (Error_Term src cs is) ret )
 -- ^ The accumulating continuation, called bottom-up.
 -> Either (Error_Term src cs is) ret

-- | Like 'compile' but also enables the use of a 'TyCtx'
-- whose actual terms are only known after 'compile' has been run.
-- This also enables to 'compile' a term once
-- and interpret it with different 'TeCtx's,
-- providing the types of those terms match the given 'TyCtx'.
compileWithTyCtx
 :: Compile cs is
 => AST_Term src is
 -> TyCtx TeName cs hs_ty
 -> Either (Error_Term src cs is) (ETerm src hs_ty cs is)
compileWithTyCtx tok ctx =
	compile tok ctx $ \typ -> Right . ETerm typ

-- | Return given 'Compiler' augmented with given 'TyCtx'.
compilerWithTyCtx
 :: Compile cs is
 => TyCtx TeName cs hs_ty
 -> (forall hs. Compiler src hs ret cs is ls rs)
 -> CompilerWithTyCtx src hs_te (hs_ty ++ hs_te) ret cs is ls rs
compilerWithTyCtx tyCtx comp ctx =
	comp (appendTyCtx tyCtx ctx)

-- | Compile, closing the 'TyCtx'.
compilerWithTyCtx_close
 :: CompilerWithTyCtx src '[] hs_te (ETerm src hs_te cs is) cs is '[] is
 -> Either (Error_Term src cs is) (ETerm src hs_te cs is)
compilerWithTyCtx_close comp =
	comp TyCtxZ $ \typ -> Right . ETerm typ

-- ** Type 'Compiler'
-- | Convenient type synonym defining a term compiler.
-- with matching 'TyCtx' and 'TeCtx'.
-- 
-- Useful to let the compiled term use
-- terms known before 'compile' is run.
type Compiler          src hs    ret cs is ls rs
 =   CompilerWithTyCtx src hs hs ret cs is ls rs

-- | Compile with an empty /lambda context/.
compileWithoutCtx
 :: Compile cs is
 => AST_Term src is
 -> Either (Error_Term src cs is) (ETermClosed src cs is)
compileWithoutCtx tok = compilerWithCtx_close (compile tok)

-- ** Type 'ACompileWithCtx'
-- | Wrap 'Compiler' to keep the /lambda context/ fully polymorphic.
-- 
-- Useful in 'compilerWithCtx' to help GHC's type solver, which
-- "Cannot instantiate unification variable with a type involving foralls".
data ACompileWithCtx src ret cs is ls rs
 =   ACompileWithCtx
 { unACompileWithCtx :: forall hs. Compiler src hs ret cs is ls rs }

-- | Compile with given /lambda context/.
compilerWithCtx
 :: Foldable f
 => f (TeName, ETermClosed src cs is)
 -> (forall hs'. Compiler src hs' ret cs is ls rs)
 -> Compiler src hs ret cs is ls rs
compilerWithCtx env comp = unACompileWithCtx $
	foldr
	 (\e (ACompileWithCtx c) -> ACompileWithCtx $ compilerWithCtx_push e c)
	 (ACompileWithCtx comp) env

-- | Return given 'Compiler' with a /lambda context/ augmented with given 'ETermClosed'.
compilerWithCtx_push
 :: (TeName, ETermClosed src cs is)
 -> (forall hs'. Compiler src hs' ret cs is ls rs)
 -> Compiler src hs ret cs is ls rs
compilerWithCtx_push (n, ETermClosed ty_n (TermClosed te_n)) comp ctx k =
	comp (TyCtxS n ty_n ctx) $ \ty_y (Term y) ->
	k ty_y $ Term $
	 \c -> y (te_n `TeCtxS` c)

-- | Compile with given 'Compiler', closing the /lambda context/.
compilerWithCtx_close
 :: (forall hs. Compiler src hs (ETermClosed src cs is) cs is '[] is)
 -> Either (Error_Term src cs is) (ETermClosed src cs is)
compilerWithCtx_close comp =
	comp TyCtxZ $ \typ (Term te) ->
	Right $ ETermClosed typ $ TermClosed $ te TeCtxZ

-- ** Class 'CompileR'
-- | Intermediate type class to construct an instance of 'Compile'
-- from many instances of 'CompileI', one for each item of @is@.
--
-- * @is@: starting  list of /interfaces/.
-- * @ls@: done      list of /interfaces/.
-- * @rs@: remaining list of /interfaces/.
class CompileR (cs::[*]) (is::[*]) (ls::[*]) (rs::[*]) where
	compileR :: TokenR src is rs i
	         -> Maybe (AST_Term src is)
	         -> Compiler src hs ret cs is ls rs

-- | Recurse into into the given 'TokenR'
-- to call the 'compileI' instance associated
-- to the 'TokenT' it contains.
instance
 ( CompileI cs is i
 , CompileR cs is (i ': ls) (r ': rs)
 ) => CompileR cs is ls (i ': r ': rs) where
	compileR tok arg ctx k =
		case tok of
		 TokenZ _src t -> compileI t arg ctx k
		 TokenS t ->
			compileR t arg ctx $ \typ (Term te :: Term hs h is (i ': ls) (r ': rs)) ->
				k typ (Term te :: Term hs h is ls (i ': r ': rs))
-- | End the recursion.
instance
 CompileI cs is i =>
 CompileR cs is ls (i ': '[]) where
	compileR (TokenZ _src t) arg  ctx  k = compileI t arg ctx k
	compileR TokenS{}        arg _ctx _k = error "Oops, the impossible happened..."

-- ** Class 'CompileI'
-- | Handle the work of 'Compile' for a given /interface/ @i@.
class CompileI (cs::[*]) (is::[*]) (i:: *) where
	compileI :: TokenT src is i
	         -> Maybe (AST_Term src is)
	         -> Compiler src hs ret cs is ls (i ': rs)

-- * Type family 'Sym_of_Ifaces'
type family Sym_of_Ifaces (is::[*]) (term:: * -> *) :: Constraint where
	Sym_of_Ifaces '[] term = ()
	Sym_of_Ifaces (i ': is) term = (Sym_of_Iface i term, Sym_of_Ifaces is term)

-- ** Type family 'Sym_of_Iface'
type family Sym_of_Iface (i:: *) :: {-term-}(* -> *) -> Constraint

-- * Type 'TyConsts_of_Ifaces'
type TyConsts_of_Ifaces is = Nub (TyConsts_of_IfacesR is)

-- ** Type family 'TyConsts_of_IfacesR'
type family TyConsts_of_IfacesR (is::[*]) where
	TyConsts_of_IfacesR '[]       = '[]
	TyConsts_of_IfacesR (i ': is) = TyConsts_of_Iface i ++ TyConsts_of_IfacesR is

{- NOTE: failed attempt to recursively collect TyConsts.
-- ** Type family 'TyConsts_imported_byR'
type family TyConsts_imported_byR (done::[*]) (cs::[*]) where
	TyConsts_imported_byR done '[] = done
	TyConsts_imported_byR done (c ': cs) =
		If (Exists c done)
		 (TyConsts_imported_byR done cs)
		 (TyConsts_imported_byR (c ': done) (TyConsts_imported_by c ++ TyConsts_imported_byR done cs))
-}

-- ** Type family 'TyConsts_of_Iface'
type family TyConsts_of_Iface (i:: *) :: [*]

-- ** Type family 'TyConsts_imported_by'
-- | Return the /type constant/s that a given /type constant/
-- wants to be part of the final list of /type constants/.
type family TyConsts_imported_by (c:: *) :: [*]

-- * Type 'TyCtx'
-- | GADT for a /typing context/:
-- accumulating at each /lambda abstraction/
-- the 'Type' of the introduced variable.
data TyCtx (name:: *) (cs::[*]) (hs::[*]) where
	TyCtxZ :: TyCtx name cs '[]
	TyCtxS :: name
	       -> Type src cs h
	       -> TyCtx name cs hs
	       -> TyCtx name cs (h ': hs)
infixr 5 `TyCtxS`

appendTyCtx
 :: TyCtx TeName cs hs0
 -> TyCtx TeName cs hs1
 -> TyCtx TeName cs (hs0 ++ hs1)
appendTyCtx TyCtxZ c          = c
appendTyCtx (TyCtxS n t c) c' = TyCtxS n t $ appendTyCtx c c'

-- * Type 'TeCtx'
-- | GADT for an /interpreting context/:
-- accumulating at each /lambda abstraction/
-- the @term@ of the introduced variable.
data TeCtx (term:: * -> *) (hs::[*]) where
	TeCtxZ :: TeCtx term '[]
	TeCtxS :: term h
	       -> TeCtx term hs
	       -> TeCtx term (h ': hs)
infixr 5 `TeCtxS`

-- * Type 'Error_Term'
data Error_Term src (cs::[*]) (is::[*])
 =   Error_Term_unbound TeName
 |   Error_Term_Typing (Error_Type src)
 |   Error_Term_Con_Type (Con_Type src cs)
 {- |   Error_Term_Con_Kind (Con_Kind src) -}
deriving instance (Eq src, Eq_Token src is, Inj_Source (EType src cs) src) => Eq (Error_Term src cs is)
deriving instance (Show src, Show_Token src is, Show_TyConst cs) => Show (Error_Term src cs is)

-- * Type 'Con_Type'
data Con_Type src cs
 =   Con_TyEq  (EType src cs) (EType src cs)
 |   Con_TyApp (EType src cs)
 |   Con_TyCon (KType src cs Constraint)
 |   Con_TyFam (At src TyFamName) [EType src cs]
deriving instance
 ( Eq src
 , Inj_Source (EType src cs) src
 ) => Eq (Con_Type src cs)
deriving instance
 ( Show src
 , Show_TyConst cs
 ) => Show (Con_Type src cs)

instance Inj_Error (Error_Type src) (Error_Term src cs ts) where
	inj_Error = Error_Term_Typing . inj_Error
instance Inj_Error (Error_Term src cs ts) (Error_Term src cs ts) where
	inj_Error = id
instance Inj_Error (Con_Type src cs) (Error_Term src cs is) where
	inj_Error = Error_Term_Con_Type
instance Inj_Error (Con_Kind src) (Error_Term src cs is) where
	inj_Error = Error_Term_Typing . Error_Type_Con_Kind

-- ** Checks

check_TyEq
 :: ( Inj_Error (Con_Type src cs) err
    , Inj_Source (EType src cs) src )
 => Type src cs x
 -> Type src cs y
 -> ((x :~: y) -> Either err ret) -> Either err ret
check_TyEq x y k =
	case x `eq_Type` y of
	 Just Refl -> k Refl
	 Nothing -> Left $ inj_Error $
		Con_TyEq (EType x) (EType y)

check_TyApp
 :: ( Inj_Error (Con_Type src cs) err
    , Inj_Source (EType src cs) src )
 => Type src cs (fa::kfa)
 -> (forall ka f a. (fa :~: f a)
  -> Type src cs (f::ka -> kfa)
  -> Type src cs (a::ka)
  -> Either err ret)
 -> Either err ret
check_TyApp typ k =
	case typ of
	 TyApp _src a b -> k Refl (a `source` EType typ) (b `source` EType typ)
	 _ -> Left $ inj_Error $ Con_TyApp (EType typ)

check_TyEq1
 :: ( Inj_Error (Con_Type src cs) err
    , Inj_Error (Con_Kind src) err
    , Inj_Source (EType src cs) src
    )
 => Type src cs (f:: * -> *)
 -> Type src cs fa
 -> (forall a. (fa :~: f a)
  -> Type src cs a
  -> Either err ret)
 -> Either err ret
check_TyEq1 typ ty_fa k =
	check_TyApp ty_fa $ \Refl ty_f ty_a ->
	check_Kind
	 (kind @(K.Type -> K.Type))
	 (kind_of ty_f) $ \Refl ->
	check_TyEq typ ty_f $ \Refl ->
	k Refl ty_a

check_TyEq2
 :: ( Inj_Error (Con_Type src cs) err
    , Inj_Error (Con_Kind src) err
    , Inj_Source (EType src cs) src )
 => Type src cs (f:: * -> * -> *)
 -> Type src cs fab
 -> (forall a b. (fab :~: f a b)
  -> Type src cs a
  -> Type src cs b
  -> Either err ret)
 -> Either err ret
check_TyEq2 typ ty_fab k =
	check_TyApp ty_fab  $ \Refl ty_fa ty_b ->
	check_TyApp ty_fa $ \Refl ty_f ty_a ->
	check_Kind
	 (kind @(K.Type -> K.Type -> K.Type))
	 (kind_of ty_f) $ \Refl ->
	check_TyEq typ ty_f $ \Refl ->
	k Refl ty_a ty_b

check_TyCon
 :: ( Proj_TyCon cs
    , Inj_Error (Con_Type src cs) err
    , Inj_Source (EType src cs) src )
 => Type src cs (q::Constraint)
 -> (TyCon q -> Either err ret)
 -> Either err ret
check_TyCon typ k =
	case proj_TyCon typ of
	 Just Dict -> k Dict
	 Nothing -> Left $ inj_Error $ Con_TyCon (KType typ)

check_TyCon1
 :: ( Proj_TyCon cs
    , Inj_Error (Con_Type src cs) err
    , Inj_Error (Con_Kind src) err
    , Inj_Source (EType src cs) src )
 => Type src cs con
 -> Type src cs (fa:: *)
 -> (forall f a. (fa :~: f a)
  -> TyCon (con f)
  -> Type src cs (f:: * -> *)
  -> Type src cs (a:: *)
  -> Either err ret)
 -> Either err ret
check_TyCon1 con ty_fa k =
	check_TyApp ty_fa $ \Refl ty_f ty_a ->
	check_Kind
	 (kind @(K.Type -> K.Type))
	 (kind_of ty_f) $ \Refl ->
	check_TyCon (TyApp sourceLess con ty_f) $ \Dict ->
	k Refl Dict ty_f ty_a

check_TyFam
 :: ( Inj_Error (Con_Type src cs) err
    , Proj_TyFam cs fam
    , Show fam )
 => At src fam
 -> Types src cs hs
 -> (Type src cs (TyFam fam hs) -> Either err ret)
 -> Either err ret
check_TyFam fam tys k =
	case proj_TyFam (unAt fam) tys of
	 Just t -> k t
	 Nothing -> Left $ inj_Error $
		Con_TyFam
		 (Text.pack . show <$> fam)
		 (eTypes tys)