{-# LANGUAGE ConstraintKinds #-}
{-# LANGUAGE GADTs #-}
{-# LANGUAGE GeneralizedNewtypeDeriving #-}
{-# LANGUAGE TypeInType #-}
{-# LANGUAGE UndecidableInstances #-}
{-# OPTIONS_GHC -fno-warn-orphans #-}
module Language.Symantic.Typing.Type
 ( module Language.Symantic.Typing.Type
 , Any
 , Proxy(..)
 , (:~:)(..)
 ) where

import Data.Map.Strict (Map)
import Data.Maybe (isJust)
import Data.Proxy
import Data.Semigroup ((<>))
import Data.Set (Set)
import Data.Text (Text)
import Data.Type.Equality
import GHC.Exts (Constraint)
import GHC.Prim (Any)
-- import Data.Typeable (Typeable, eqT)
import qualified Data.Kind as K
import qualified Data.List as List
import qualified Data.Map.Strict as Map
import qualified Data.Set as Set
import qualified Data.Text as Text

import Language.Symantic.Grammar
import Language.Symantic.Typing.Kind
import Language.Symantic.Typing.Constant
import Language.Symantic.Parsing

-- * Type 'T', for both 'Type' and 'Term'

-- | /Type of terms/ embedding… /terms/ to be able
-- to have /polymorphic terms/ (ie. handle quantified variables).
--
-- * @src@: /source/.
-- * @cs@: /type constants/, fixed at compile time.
-- * @h@: indexed Haskell type.
-- * @k@: indexed Haskell kind of @h@.
--
-- * 'TyConst': /type constant/, selected amongst @cs@.
--
-- * 'TyApp': /type parameter application/.
--
-- * 'TyQuant': /type universal quantification/ (aka. /type abstraction/):
--   used to introduce (and bind) a /type variable/.
--   
--   Note that the type indexing is lost (though not kind @k@),
--   because GHC cannot represent a quantified type in a type
--   aka. impredicative polymorphism.
--
-- * 'TyVar': /type variable/:,
--   used to represent a /bounded type variable/,
--   when passing through a 'TyQuant'.
--   
--   Note that only the kind @k@ is flexible,
--   @h@ being set to 'Any' to loose the type indexing
--   and enabling the return of a type equality when the 'Var's are the same.
--   
--   Note also that there is no DeBruijn encoding done to link 'TyQuant's and 'TyVar's.
--
-- * 'Term': /term/, useful to embed a term
--   requiring quantifications and qualifications.
--
-- See also: https://ghc.haskell.org/trac/ghc/wiki/Typeable
-- Which currently concludes:
-- "Why kind equalities, then? Given the fact that Richard's branch
-- doesn't solve this problem, what is it good for?
-- It still works wonders in the mono-kinded case,
-- such as for decomposing ->.
-- It's just that poly-kinded constructors are still a pain."
data T (src::K.Type) (cs::[K.Type]) (h::k) where
	TyConst :: src -> TyConst src cs c
	               -> Type src cs c
	TyApp   :: src -> Type src cs f
	               -> Type src cs a
	               -> Type src cs (f a)
	               -- NOTE: @f a@ decomposition made possible
	               -- at the cost of requiring @TypeInType@.
	TyQuant :: src -> NameHint
	               -> Kind src kv
	               -> (forall (v::kv). Type src cs v -> KType src cs k)
	               -> Type src cs (h::k)
	TyVar   :: src -> Kind src kv
	               -> Var
	               -> Type src cs (Any::kv)
	Term    :: TeSym h
	               -> Type src cs h
	               -> Term src cs h

-- ** Type 'Type'
type Type = T

-- ** Type 'Term'
type Term src cs h = T src cs (h::K.Type)

-- | Remove all top 'Term's.
unTerm :: Type src cs (h::k) -> Type src cs (h::k)
unTerm (Term _te x) = unTerm x
unTerm t = t

-- | 'Type' constructor to be used with @TypeApplications@ and @NoMonomorphismRestriction@.
ty :: forall c src cs. (Source src, Inj_TyConst cs c) => Type src cs c
ty = TyConst sourceLess inj_TyConst
 -- NOTE: The forall brings @c@ first in the type variables,
 -- which is convenient to use @TypeApplications@.

(~>) :: Source src => Inj_TyConst cs (->)
 => Type src cs a -> Type src cs b -> Type src cs (a -> b)
(~>) a b = (ty @(->) `tyApp` a) `tyApp` b
infixr 5 ~>

-- | Like 'ty', but for a known 'Source'.
tySourced :: forall c src cs. (Source src, Inj_TyConst cs c) => src -> Type src cs c
tySourced src = TyConst src inj_TyConst

-- | Like 'TyApp', but 'sourceLess'.
tyApp :: Source src => Type src cs f -> Type src cs a -> Type src cs (f a)
tyApp = TyApp sourceLess
infixl 9 `tyApp`

-- | Like 'TyQuant', but 'sourceLess'.
tyQuant
 :: Source src
 => NameHint
 -> Kind src kv
 -> (forall (v::kv). Type src cs v -> KT src cs k)
 -> KT src cs k
tyQuant n kv f = KT $ TyQuant sourceLess n kv f

-- | Return the 'Kind' of the given 'Type'.
kindTy :: Inj_Source (EType src cs) src => Type src cs (h::k) -> Kind src k
kindTy t = kindTyOS t `source` EType t

-- | Like 'kindTy' but keep the old 'Source'.
kindTyOS :: Type src cs (h::k) -> Kind src k
kindTyOS = go
	where
	go :: forall cs src k h. Type src cs (h::k) -> Kind src k
	go = \case
	 TyConst _src c -> kindTyConst c
	 TyApp _src f _x ->
		case kindTyOS f of
		 KiArrow _src_kf _kx kf -> kf
	 TyQuant src _n k f -> unKT kindTyOS $ f (TyVar src k 0)
	 TyVar  _src k _v   -> k
	 Term _te x         -> kindTyOS x

instance -- TestEquality
 Inj_Source (EType src cs) src
 => TestEquality (Type src cs) where
	testEquality = eqTy
instance Eq (Type src cs h) where
	_x == _y = True

-- | Return a proof when two 'Type's are equals.
--
-- NOTE: 'TyQuant's are not compared:
-- they first have to be monomorphized.
eqTy :: Type src cs (x::k) -> Type src cs (y::k) -> Maybe (x:~:y)
eqTy = go
	where
	go :: Type src cs (x::k) -> Type src cs (y::k) -> Maybe (x:~:y)
	go (TyConst _sx x) (TyConst _sy y)
	 | Just Refl <- testEquality x y
	 = Just Refl
	go (TyApp _sx fx xx) (TyApp _sy fy xy)
	 | Just Refl <- eqKi (kindTyOS fx) (kindTyOS fy)
	 , Just Refl <- go fx fy
	 , Just Refl <- go xx xy
	 = Just Refl
	go TyQuant{} TyQuant{} = Nothing
	{-
	go v (TyQuant sx _nx kx x::Type src cs vx)
	     (TyQuant sy _ny ky y::Type src cs vy)
	 | Just Refl <- eqKi kx ky
	 , Just Refl <- go (v - 1) (x $ TyVar sx kx v) (y $ TyVar sy ky v)
	 = Just Refl
	-}
	go (TyVar _sx kx x) (TyVar _sy ky y)
	 | Just Refl <- eqKi kx ky
	 , True <- x == y
	 = Just Refl
	-- NOTE: 'Term' are transparent to 'eqTy'.
	go (Term _tx x) (Term _ty y) | Just Refl <- go x y = Just Refl
	go (Term _tx x) y            | Just Refl <- go x y = Just Refl
	go x            (Term _ty y) | Just Refl <- go x y = Just Refl
	go _x _y = Nothing

-- | Return two proofs when two 'Type's are equals:
-- one for the type and one for the kind.
eqKiTy :: Type src cs (x::kx) -> Type src cs (y::ky) -> Maybe (x:~~:y)
eqKiTy = go
	where
	go :: Type src cs (x::kx) -> Type src cs (y::ky) -> Maybe (x:~~:y)
	go (TyConst _sx x) (TyConst _sy y) = eqKiTyConst x y
	go (TyApp _sx fx xx) (TyApp _sy fy xy)
	 | Just KRefl <- go fx fy
	 , Just KRefl <- go xx xy
	 = Just KRefl
	go TyQuant{} TyQuant{} = Nothing
	go (TyVar _sx kx _x) (TyVar _sy ky _y)
	 | Just Refl <- eqKi kx ky
	 = Just KRefl
	go (Term _tx x) (Term _ty y) = go x y
	go _x _y = Nothing

-- * Type '(#>)'
-- | 'TyConst' to qualify a 'Type'.
data (#>) (q::Constraint) (a::K.Type)
instance Fixity_TyConstC (#>) where
	fixity_TyConstC _ = FixityInfix $ infixR 0

(#>)
 :: Source src
 => Inj_TyConst cs (#>)
 => Type src cs a -> Type src cs b -> Type src cs (a #> b)
(#>) a b = (ty @(#>) `tyApp` a) `tyApp` b
infixr 0 #>
 -- NOTE: should actually be (-1)
 -- to compose well with @infixr 0 (->)@
 -- but this is not allowed by GHC.


instance
 ( Show_TyConst cs
 , Fixity_TyConst cs
 ) => Show (Type src cs h) where
	showsPrec = showTy

showTy :: forall cs src h.
 Show_TyConst cs
 => Fixity_TyConst cs
 => Int -> Type src cs h -> ShowS
showTy pr typ =
	go Nothing
	 (infixB L pr, L)
	 (freeTyVar typ)
	 (Map.empty, namesTy typ) typ
	where
	go :: forall p x.
	    Maybe (Type src cs p) -- parent Type, used to prettify
	 -> (Infix, LR)           -- fixity and associativity
	 -> Var                   -- variable counter (DeBruijn)
	 -> ( Map Var Name        -- names of bound variables
	    , Names               -- used names : bound variables', free variables' and constants'
	    )
	 -> Type src cs x
	 -> ShowS
	go _prev (po, _) _v _vs (TyConst _src c) = showsPrec (precedence po) c
	go _prev po v vs t@(TyApp _ (TyApp _ (TyConst _ f) a) b)
	 | FixityInfix op <- fixity_TyConst f =
		showParen (parenInfix po op) $
		go (Just t) (op, L) v vs a .
		showChar ' ' .
		showString (unParen $ show_TyConst f) .
		showChar ' ' .
		go (Just t) (op, R) v vs b
		where
		unParen ('(':s) | ')':s' <- reverse s = reverse s'
		unParen s = s
	go _prev po v vs t@(TyApp _src f a) =
		let op = infixR 11 in
		showParen (parenInfix po op) $
		go (Just t) (op, L) v vs f .
		showChar ' ' .
		go (Just t) (op, R) v vs a
	go prev po v (vs, ns) t@(TyQuant src nv kv f) =
		let op = infixR 0 in
		let var = TyVar src kv v in
		let n   = freshifyName ns nv in
		let vs' = Map.insert v n vs in
		let ns' = Set.insert n ns in
		case f var of
		 KT x ->
			showParen (parenInfix po op) $
			(case prev of
			 Just TyQuant{} -> id
			 _ -> showString "forall") .
			showChar ' ' .
			go Nothing po v (vs', ns') var .
			(case x of
			 TyQuant{} -> id
			 _ -> showString ". ") .
			go (Just t) (op, R) (v + 1) (vs', ns') x
	go _prev _po _v (vs, _ns) (TyVar _src _kv v) =
		case Map.lookup v vs of
		 Nothing -> showChar '#' . showsPrec 0 v
		 Just n  -> showString $ Text.unpack n
	go _prev po v vs t@(Term _te x) =
		showBracket True $
		go (Just t) po v vs x

showBracket :: Bool -> ShowS -> ShowS
showBracket b p =  if b then showChar '{' . p . showChar '}' else p

-- ** Type 'TyCon'
-- | Captures the proof of a 'Constraint'
-- (and its dictionaries when it contains type classes):
-- pattern matching on the 'TyCon' constructor
-- brings the 'Constraint' into scope.
data TyCon c where
	Dict :: c => TyCon c

-- ** Type 'TyFamName'
type TyFamName = Text

-- ** Class 'Sym'
class Sym (term::K.Type -> K.Type) where
	sym_lam :: (term a -> term b) -> term (a -> b)
	sym_app :: term (a -> b) -> (term a -> term b)
	sym_int :: Int -> term Int
	sym_id  :: term (a -> a)
	sym_uc  :: term (Char -> Char)
	-- sym_bind :: Monad m => term (m a) -> term (a -> m b) -> term (m b)
	-- sym_bind :: term (Monad m #> (m a -> (a -> m b) -> m b))
	sym_bind :: Monad m => term (m a -> (a -> m b) -> m b)
	sym_const :: term a -> term b -> term a
	-- sym_comp :: term (b -> c) -> term (a -> c) -> term (a -> b)
	sym_comp :: term ((b -> c) -> (a -> b) -> a -> c)

-- ** Type 'TeSym'
data TeSym (h::K.Type) = TeSym (forall term. Sym term => term h)

instance Source src => Sourced (Type src cs h) where
	type Source_of (Type src cs h) = src
	
	source_of (TyConst src _c)        = src
	source_of (TyApp   src _f _x)     = src
	source_of (TyQuant src _n _kv _f) = src
	source_of (TyVar   src _kv _v)    = src
	source_of (Term  _te x)           = source_of x
	
	source_is (TyConst _src c)      src = TyConst src c
	source_is (TyApp   _src f x)    src = TyApp   src f x
	source_is (TyQuant _src n kv f) src = TyQuant src n kv f
	source_is (TyVar   _src kv v)   src = TyVar   src kv v
	source_is (Term  te x)          src = Term te $ source_is x src

-- ** Type 'Var'
type Var = Int

-- | Return the smallest 'Var' not used in given 'Type'.
freeTyVar :: Type src cs (h::k) -> Var
freeTyVar = (+ 1) . maxTyVar (-1)
	where
	maxTyVar :: Var -> Type src cs h -> Var
	maxTyVar m TyConst{}             = m
	maxTyVar m (TyApp   _src f a)    = maxTyVar m f `max` maxTyVar m a
	maxTyVar m (TyQuant src _n kv f) = ofKT (maxTyVar m) $ f $ TyVar src kv (-1)
	maxTyVar m (TyVar   _src _kv v)  = max m $ if v < 0 then -1 else v + 1
	maxTyVar m (Term  _te x)         = maxTyVar m x

-- ** Type 'Name'
type Name     = Text
type NameHint = Name
type Names    = Set Name

namesTy :: Show_TyConst cs => Type src cs h -> Names
namesTy (TyConst _src c)      = Set.singleton $ Text.pack $ show_TyConst c
namesTy (TyApp   _src f x)    = namesTy f `Set.union` namesTy x
namesTy (TyQuant src _n kv f) = ofKT namesTy $ f $ TyVar src kv (-1)
namesTy (TyVar   _src _kv v)  | v < 0     = Set.empty
                              | otherwise = Set.singleton $ "#" <> Text.pack (show v)
namesTy (Term  _te x) = namesTy x

-- | Infinite list of unique 'Name's:
-- @a, b, .., z, a1, b1 .., z1, a2, ..@
poolNames :: [Name]
poolNames =
	[ Text.singleton n     | n <- ['a'..'z'] ] <>
	[ Text.pack (n:show i) | n <- ['a'..'z']
	                       , i <- [1 :: Int ..] ]

-- | Return given 'Name' renamed a bit to avoid
-- conflicting with any given 'Names'.
freshifyName :: Names -> Name -> Name
freshifyName ns "" = freshName ns
freshifyName ns n =
	let ints = [1..] :: [Int] in
	List.head
	 [ fresh
	 | suffix <- "" : (show <$> ints)
	 , fresh <- [n <> Text.pack suffix]
	 , fresh `Set.notMember` ns
	 ]

freshName :: Names -> Name
freshName ns = List.head $ poolNames List.\\ Set.toList ns

-- * Type 'Con_Type'
data Con_Type src cs
 =   Con_EqType  (EType src cs) (EType src cs)
 |   Con_TyApp   (EType src cs)
 |   Con_TyCon   (KT 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
 , Fixity_TyConst cs
 ) => Show (Con_Type src cs)

if_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
if_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)

if_EqType
 :: Inj_Error (Con_Type src cs) err
 => Type src cs x
 -> Type src cs y
 -> ((x :~: y) -> Either err ret) -> Either err ret
if_EqType x y k =
	case x `eqTy` y of
	 Just Refl -> k Refl
	 Nothing -> Left $ inj_Error $
		Con_EqType (EType x) (EType y)

if_EqType1
 :: ( Inj_Error (Con_Type src cs) err
    , Inj_Error (Con_Kind src) err
    , Inj_Source (EType src cs) src
    )
 => Type src cs (f:: K.Type -> K.Type)
 -> Type src cs fa
 -> (forall a. (fa :~: f a)
  -> Type src cs a
  -> Either err ret)
 -> Either err ret
if_EqType1 typ ty_fa k =
	if_TyApp ty_fa $ \Refl ty_f ty_a ->
	if_EqKind (kind @(K.Type -> K.Type))
	          (kindTy ty_f) $ \Refl ->
	if_EqType typ ty_f $ \Refl ->
	k Refl ty_a

if_EqType2
 :: ( Inj_Error (Con_Type src cs) err
    , Inj_Error (Con_Kind src) err
    , Inj_Source (EType src cs) src )
 => Type src cs (f:: K.Type -> K.Type -> K.Type)
 -> 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
if_EqType2 typ ty_fab k =
	if_TyApp ty_fab $ \Refl ty_fa ty_b ->
	if_TyApp ty_fa  $ \Refl ty_f  ty_a ->
	if_EqKind (kind @(K.Type -> K.Type -> K.Type))
	          (kindTy ty_f) $ \Refl ->
	if_EqType typ ty_f $ \Refl ->
	k Refl ty_a ty_b

if_TyFun ::
 ( Inj_TyConst cs (->)
 , Inj_Source (EType src cs) src
 , Inj_Error (Con_Kind src) err
 , Inj_Error (Con_Type src cs) err
 ) => Type src cs fab
 -> (forall a b. (:~:) fab (a -> b)
    -> Type src cs a -> Type src cs b -> Either err ret)
 -> Either err ret
if_TyFun = if_EqType2 (ty @(->))

-- * Type 'EType'
-- | Existential for 'Type'.
data EType src cs = forall h. EType (Type src cs h)

instance Inj_Source (EType src cs) src => Eq (EType src cs) where
	EType x == EType y = isJust $ eqKiTy x y
instance
 ( Show_TyConst cs
 , Fixity_TyConst cs
 ) => Show (EType src cs) where
	showsPrec p (EType t) = showsPrec p t

mapEType
 :: (forall k (h::k). Type src cs h -> Type src cs h)
 -> EType src cs
 -> EType src cs
mapEType f (EType t) = EType (f t)

bindEType
 :: (forall k (h::k). Type src cs h -> EType src cs)
 -> EType src cs
 -> EType src cs
bindEType f (EType t) = f t

-- * Type 'KT'
-- | Existential for 'T' of a known 'Kind'.
data KT src cs k = forall (h::k). KT (Type src cs h)
type KType = KT
type KTerm src cs = KT src cs K.Type

instance Eq (KT src cs ki) where
	KT x == KT y = isJust $ eqTy x y
instance
 ( Show_TyConst cs
 , Fixity_TyConst cs
 ) => Show (KT src cs ki) where
	showsPrec p (KT t) = showTy p t
instance Inj_Source (EType src cs) src => Inj_Source (KT src cs k) src where
	inj_Source (KT t) = inj_Source $ EType t

unKT :: (forall (h::k). Type src cs h -> a k) -> KT src cs k -> a k
unKT f (KT t) = f t

ofKT :: (forall (h::k). Type src cs h -> a) -> KT src cs k -> a
ofKT f (KT t) = f t

mapKT
 :: (forall (h::k). Type src cs h -> Type src cs h)
 -> KT src cs k
 -> KT src cs k
mapKT f (KT t) = KT (f t)

bindKT
 :: (forall (h::k). Type src cs h -> KT src cs k)
 -> KT src cs k
 -> KT src cs k
bindKT f (KT t) = f t

bind2KT
 :: (forall (x::ka) (y::kb). Type src cs x -> Type src cs y -> KT src cs kxy)
 -> KT src cs ka
 -> KT src cs kb
 -> KT src cs kxy
bind2KT f (KT x) (KT y) = f x y

-- * Type 'Types'
data Types src cs (hs::[K.Type]) where
	TypesZ :: Types src cs '[]
	TypesS :: Type  src cs h
	       -> Types src cs hs
	       -> Types src cs (Proxy h ': hs)
infixr 5 `TypesS`

showTys
 :: forall cs src hs.
    Show_TyConst cs
 => Fixity_TyConst cs
 => Types src cs hs -> ShowS
showTys ts = showString "[" . go ts . showString "]"
	where
	go :: forall xs. Types src cs xs -> ShowS
	go TypesZ = showString ""
	go (TypesS h0 (TypesS h1 hs)) = showTy 10 h0 . showString ", " . showTy 10 h1 . go hs
	go (TypesS h hs) = showTy 10 h . go hs

instance
 ( Show_TyConst cs
 , Fixity_TyConst cs
 ) => Show (Types src cs hs) where
	showsPrec _p = showTys

eTypes :: Types src cs hs -> [EType src cs]
eTypes TypesZ = []
eTypes (TypesS t ts) = EType t : eTypes ts

{-
-- * Type 'UnProxy'
type family UnProxy (x::K.Type) :: k where
	UnProxy (Proxy x) = x
-}