{-# LANGUAGE ConstraintKinds #-}
{-# LANGUAGE PatternSynonyms #-}
{-# LANGUAGE PolyKinds #-}
{-# LANGUAGE RankNTypes #-}
{-# LANGUAGE StandaloneDeriving #-}
{-# LANGUAGE UndecidableInstances #-}
{-# LANGUAGE UndecidableSuperClasses #-}
{-# LANGUAGE ViewPatterns #-}
{-# OPTIONS_GHC -Wno-partial-fields #-}

module Symantic.Typer.Type where

import Control.Monad.Classes as MC ()
import Control.Monad.Classes.Run as MC ()
import Data.Eq (Eq (..))
import Data.Function (($))
import Data.Functor ((<$>))
import Data.Map.Strict qualified as Map
import Data.Maybe (Maybe (..), isJust)
import Data.Monoid (Monoid (..))
import Data.Proxy (Proxy (..))
import Data.Semigroup (Semigroup (..))
import Data.String (String)
import GHC.Types
import Text.Show (Show (..))
import Type.Reflection (TypeRep, Typeable, eqTypeRep, typeRep, typeRepKind, (:~~:) (..), pattern App)
import Prelude qualified

import Symantic.Typer.Eq
import Symantic.Typer.List

data WithProv prov a = WithProv {withProv :: prov, unWithProv :: a}

-- * Type 'Const'
type Const = TypeRep

-- * Type 'Name'
type Name = String

-- * Type 'Prov'
type Prov = Type

-- * Type 'Var'
data Var :: forall vK. Prov -> [Type] -> vK -> Type where
  VarZ ::
    { varZKind :: Kind prov vK
    , varZNext :: Len ws
    } ->
    Var prov (Proxy (v :: vK) ': ws) (v :: vK)
  VarS :: Var prov vs v -> Var prov (not_v ': vs) v
deriving instance Show prov => Show (Var prov vs v)

pattern Var :: () => () => Kind prov vK -> Len ws -> Var prov vs' (v :: vK)
pattern Var{varKind, varNext} <- (var -> VarVS VarZ{varZKind = varKind, varZNext = varNext})
{-# COMPLETE Var #-}

-- -- * Type 'VarVS'
data VarVS prov (v :: vK) = forall vs. VarVS (Var prov vs v)
var :: Var prov vs (v :: vK) -> VarVS prov (v :: vK)
var v@VarZ{} = VarVS v
var (VarS v) = var v

instance EqTy (Var prov vs) (Var prov vs) where
  eqTy VarZ{} VarZ{} = Just HRefl
  eqTy (VarS x) (VarS y) = eqTy x y
  eqTy _ _ = Nothing

type VarSem = Type -> [Type] -> Type -> Type

-- data VarV (var::VarSem) vs = forall vK (v::vK). VarVSem (var vs v)
data VarV prov vs = forall v. VarV (Var prov vs v)

class LenVar var where
  lenVar :: var prov vs a -> Len vs
instance LenVar Var where
  lenVar VarZ{varZNext} = LenS varZNext
  lenVar (VarS v) = LenS (lenVar v)

class LenVar var => AllocVars var where
  allocVarsL :: Len len -> var prov vs x -> var prov (len ++ vs) x
  allocVarsR :: Len len -> var prov vs x -> var prov (vs ++ len) x
instance AllocVars Var where
  allocVarsL LenZ v = v
  allocVarsL (LenS len) v = VarS (allocVarsL len v)

  allocVarsR len VarZ{..} = VarZ{varZNext = addLen varZNext len, ..}
  allocVarsR len (VarS v) = VarS (allocVarsR len v)

appendVars ::
  AllocVars xVar =>
  AllocVars yVar =>
  xVar prov xVS (x :: xK) ->
  yVar prov yVS (y :: yK) ->
  ( xVar prov (xVS ++ yVS) x
  , yVar prov (xVS ++ yVS) y
  )
appendVars x y =
  ( allocVarsR (lenVar y) x
  , allocVarsL (lenVar x) y
  )

-- ** Type 'IndexVar'

-- | Index of a 'Var'.
type IndexVar = Int

indexVar :: Var prov vs v -> Int
indexVar VarZ{} = 0
indexVar (VarS v) = 1 Prelude.+ indexVar v

-- instance KindOf (Var prov vs) where
--   kindOf = fmap go
--     where
--     go :: forall vs2 aK a. Var prov vs2 (Proxy (a::aK)) -> TypeRep aK
--     go VarZ{varKind} = varKind
--     go (VarS v) = go v

-- * Type 'Vars'
type Vars prov vs = Map.Map Name (VarV prov vs)

lookupVar :: Name -> Vars prov vs -> Maybe (VarV prov vs)
lookupVar = Map.lookup

insertVar :: Name -> VarV prov vs -> Vars prov vs -> Vars prov vs
insertVar = Map.insert

-- * Type 'UsedVars'

-- | List of 'Var's, used to change the context of a 'Var'
-- when removing unused 'Var's.
data UsedVars oldVS prov newVS a where
  UsedVarsZ :: UsedVars oldVS prov '[] a
  UsedVarsS ::
    Var prov oldVS (v :: vK) ->
    UsedVars oldVS prov newVS a ->
    UsedVars oldVS prov (Proxy (v :: vK) ': newVS) a

instance LenVar (UsedVars oldVS) where
  lenVar UsedVarsZ = LenZ
  lenVar (UsedVarsS _v oldVS) = LenS $ lenVar oldVS

lookupUsedVars ::
  Var prov oldVS (v :: vK) ->
  UsedVars oldVS prov newVS a ->
  Maybe (Var prov newVS (v :: vK))
lookupUsedVars _v UsedVarsZ = Nothing
lookupUsedVars v (UsedVarsS v' us) = do
  case v `eqTy` v' of
    Just HRefl ->
      case v' of
        Var{varKind} ->
          Just
            VarZ
              { varZNext = lenVar us
              , varZKind = varKind
              }
    Nothing -> VarS Data.Functor.<$> lookupUsedVars v us

insertUsedVars ::
  Var prov vs v ->
  UsedVars vs prov us a ->
  Maybe (UsedVars vs prov (Proxy v ': us) a)
insertUsedVars v vs =
  case lookupUsedVars v vs of
    Just{} -> Nothing
    Nothing -> Just (UsedVarsS v vs)

-- | Remove unused 'Var's from the context.
normalizeVarsTy ::
  Ty prov vs (a :: aK) ->
  (forall newVS. Ty prov newVS (a :: aK) -> ret) ->
  ret
normalizeVarsTy ty k =
  usedVarsOf UsedVarsZ ty $ \vs ->
    k $ goTy vs ty
  where
    goTy :: UsedVars oldVS prov newVS b -> Ty prov oldVS a -> Ty prov newVS a
    goTy vs TyConst{..} = TyConst{tyConstLen = lenVar vs, ..}
    goTy vs TyApp{..} = TyApp{tyAppFun = goTy vs tyAppFun, tyAppArg = goTy vs tyAppArg}
    goTy vs TyVar{..} =
      case tyVar `lookupUsedVars` vs of
        Nothing -> Prelude.error "[BUG] normalizeVarsTy: impossible lookup fail"
        Just v' -> TyVar{tyVar = v', ..}
    goTy vs TyProv{..} = TyProv{tyUnProv = goTy vs tyUnProv, ..}

-- goTy vs (TyFam src _len fam as) = TyFam src (lenVars vs) fam $ goTys vs as

-- goTys ::
-- UsedVars src vs vs' ->
-- Tys src vs ts ->
-- Tys src vs' ts
-- goTys _vs TysZ = TypesZ
-- goTys vs (TysS t ts) =
--  goTy vs t `TypesS`
--  goTys vs ts

-- * Class 'VarOccursIn'

-- | Test whether a given 'Var' occurs within @(t)@.
class VarOccursIn (t :: Type -> [Type] -> aK -> Type) where
  varOccursIn :: Var prov vs v -> t prov vs a -> Bool

-- * Class 'UsedVarsOf'
class UsedVarsOf t where
  usedVarsOf ::
    UsedVars vs prov oldVS a ->
    t prov vs b ->
    (forall newVS. UsedVars vs prov newVS a -> ret) ->
    ret

-- * Type 'Ty'

-- Use a CUSK, because it's a polymorphically recursive type,
-- See https://ryanglscott.github.io/2020/01/05/five-benefits-to-using-standalonekindsignatures/
data Ty :: forall aK. Prov -> [Type] -> aK -> Type where
  TyConst :: {tyConst :: Const a, tyConstLen :: Len vs} -> Ty prov vs a
  TyVar :: {tyVar :: Var prov vs a, tyVarName :: Name} -> Ty prov vs a
  TyApp :: {tyAppFun :: Ty prov vs f, tyAppArg :: Ty prov vs a} -> Ty prov vs (f a)
  TyProv :: {tyProv :: prov, tyUnProv :: Ty prov vs a} -> Ty prov vs a
infixl 9 `TyApp`

tyOfTypeRep :: Len vs -> TypeRep a -> Ty prov vs a
tyOfTypeRep len ty = case ty of
  App f x -> TyApp{tyAppFun = tyOfTypeRep len f, tyAppArg = tyOfTypeRep len x}
  constType -> TyConst{tyConst = constType, tyConstLen = len}

-- | Monomorphic 'Ty' from a 'Typeable'.
monoTy :: forall {aK} (a :: aK) prov. Typeable a => Ty prov '[] a
monoTy = tyOfTypeRep LenZ typeRep

aTy :: LenInj vs => Monoid prov => Ty prov (Proxy a ': vs) (a :: Type)
bTy :: LenInj vs => Monoid prov => Ty prov (Proxy a ': Proxy b ': vs) (b :: Type)
cTy :: LenInj vs => Monoid prov => Ty prov (Proxy a : Proxy b : Proxy c : vs) (c :: Type)
dTy :: LenInj vs => Monoid prov => Ty prov (Proxy a : Proxy b : Proxy c : Proxy d : vs) (d :: Type)
aTy = TyVar{tyVar = VarZ{varZKind = monoTy @Type, varZNext = lenInj}, tyVarName = "a"}
bTy = TyVar{tyVar = VarS $ VarZ{varZKind = monoTy @Type, varZNext = lenInj}, tyVarName = "b"}
cTy = TyVar{tyVar = VarS $ VarS $ VarZ{varZKind = monoTy @Type, varZNext = lenInj}, tyVarName = "c"}
dTy = TyVar{tyVar = VarS $ VarS $ VarS $ VarZ{varZKind = monoTy @Type, varZNext = lenInj}, tyVarName = "d"}

pattern SkipTyProv :: () => () => Ty prov vs a -> Ty prov vs a
pattern SkipTyProv ty <- (skipTyProv -> ty)
{-# COMPLETE SkipTyProv #-}

skipTyProv :: Ty prov vs a -> Ty prov vs a
skipTyProv TyProv{tyUnProv} = tyUnProv
skipTyProv x = x

pattern FUN ::
  -- CReq
  forall k (fun :: k) prov vs.
  () =>
  {-Monoid prov-}
  -- CProv
  forall
    (r1 :: RuntimeRep)
    (r2 :: RuntimeRep)
    (arg :: TYPE r1)
    (res :: TYPE r2).
  ( k ~ Type
  , fun ~~ (arg -> res)
  -- , r1 ~ LiftedRep
  -- , r2 ~ LiftedRep {- fixme: remove those 2 constraints -}
  ) =>
  Ty prov vs arg ->
  Ty prov vs res ->
  Ty prov vs fun
pattern FUN arg res <-
  SkipTyProv
    TyApp
      { tyAppFun =
        SkipTyProv
          TyApp
            { tyAppFun = SkipTyProv TyConst{tyConst = eqTypeRep (typeRep @(->)) -> Just HRefl}
            , tyAppArg = arg
            }
      , tyAppArg = res
      }

(~>) :: Monoid prov => Ty prov vs a -> Ty prov vs b -> Ty prov vs (a -> b)
(~>) arg res =
  TyApp
    { tyAppFun =
        TyApp
          { tyAppFun =
              TyConst{tyConst = typeRep, tyConstLen = (lenVar arg)}
          , tyAppArg = arg
          }
    , tyAppArg = res
    }
infixr 0 ~>

deriving instance Show prov => Show (Ty prov vs a)

-- type instance SourceOf (Ty vs src t) = src
instance LenVar Ty where
  lenVar TyConst{tyConstLen} = tyConstLen
  lenVar TyApp{tyAppFun} = lenVar tyAppFun
  lenVar TyVar{tyVar} = lenVar tyVar
  lenVar TyProv{tyUnProv} = lenVar tyUnProv

-- lenVars (TyFam l _f _as) = l
instance AllocVars Ty where
  allocVarsL len ty@TyConst{..} = ty{tyConstLen = addLen len tyConstLen}
  allocVarsL len TyApp{..} = TyApp{tyAppFun = allocVarsL len tyAppFun, tyAppArg = allocVarsL len tyAppArg}
  allocVarsL len ty@TyVar{..} = ty{tyVar = allocVarsL len tyVar}
  allocVarsL len TyProv{..} = TyProv{tyProv, tyUnProv = allocVarsL len tyUnProv}

  -- allocVarsL len (TyFam l f as) = TyFam (addLen len l) f $ allocVarsL len `mapTys` as

  allocVarsR len ty@TyConst{..} = ty{tyConstLen = addLen tyConstLen len}
  allocVarsR len TyApp{..} = TyApp{tyAppFun = allocVarsR len tyAppFun, tyAppArg = allocVarsR len tyAppArg, ..}
  allocVarsR len ty@TyVar{..} = ty{tyVar = allocVarsR len tyVar}
  allocVarsR len TyProv{..} = TyProv{tyProv, tyUnProv = allocVarsR len tyUnProv}

-- allocVarsR len (TyFam l f as) = TyFam (addLen l len) f $ allocVarsR len `mapTys` as

instance VarOccursIn Ty where
  varOccursIn v TyVar{tyVar = v'} | Just{} <- v `eqTy` v' = False
  varOccursIn _v TyVar{} = False
  varOccursIn _v TyConst{} = False
  varOccursIn v TyApp{tyAppFun = f, tyAppArg = a} = varOccursIn v f Prelude.|| varOccursIn v a
  varOccursIn v TyProv{..} = varOccursIn v tyUnProv

-- varOccursIn v (TyFam _src _len _fam as) = varOccursIn v as
instance UsedVarsOf Ty where
  usedVarsOf vs TyConst{} k = k vs
  usedVarsOf vs TyApp{tyAppFun = f, tyAppArg = a} k =
    usedVarsOf vs f $ \vs' ->
      usedVarsOf vs' a k
  usedVarsOf vs TyVar{tyVar = v} k =
    case insertUsedVars v vs of
      Nothing -> k vs
      Just vs' -> k vs'
  usedVarsOf vs TyProv{tyUnProv} k = usedVarsOf vs tyUnProv k

-- usedVarsOf vs (TyFam _src _len _fam as) k = usedVarsOf vs as k

-- ** Type 'TyT'

-- | Existential for 'Type'.
data TyT prov vs = forall a. TyT (Ty prov vs a)

-- type instance SourceOf (TyT vs src) = src
instance Eq (TyT prov vs) where
  TyT x == TyT y = isJust $ eqTy x y
instance Show prov => Show (TyT prov vs) where
  showsPrec p (TyT x) = showsPrec p x

-- * Type 'Kind'
type Kind prov a = Ty prov '[] a

-- ** Type 'KindK'
type KindK prov = TyT prov '[]

-- ** Type 'TyVT'
data TyVT prov = forall vs a. TyVT (Ty prov vs a)

deriving instance Show prov => Show (TyVT prov)
instance Eq (TyVT m) where
  TyVT x == TyVT y = isJust (eqTy x' y')
    where
      (x', y') = appendVars x y

-- instance (Monoid src) => EqTy (MT.Writer src) (MT.Writer src) where
--  eqTy x y = eqTy (unSourced x) (unSourced y)
instance EqTy (Ty prov vs) (Ty prov vs) where
  eqTy = go
    where
      go :: forall xK xT yK yT. Ty prov vs (xT :: xK) -> Ty prov vs (yT :: yK) -> Maybe (xT :~~: yT)
      go TyConst{tyConst = x} TyConst{tyConst = y} = eqTy x y
      go TyVar{tyVar = x} TyVar{tyVar = y} = eqTy x y
      go TyApp{tyAppFun = xF, tyAppArg = xA} TyApp{tyAppFun = yF, tyAppArg = yA}
        | Just HRefl <- go xF yF
        , Just HRefl <- go xA yA =
            Just HRefl
      go TyProv{tyUnProv = x} y = go x y
      go x TyProv{tyUnProv = y} = go x y
      go _x _y = Nothing

-- eqTy (TyFam _ _lx xF xA) (TyFam _ _ly yF yA)
--  | Just HRefl <- eqTypeRep xF yF
--  , Just HRefl <- eqTys xA yA
--  = Just HRefl
-- -- NOTE: 'TyFam' is expanded as much as possible.
-- eqTy x@TyFam{} y = eqTy (expandFam x) y
-- eqTy x y@TyFam{} = eqTy x (expandFam y)

-- * Type 'Tys'
data Tys (as :: [Type]) prov vs (a :: aK) where
  TysZ :: Tys '[] prov vs a
  TysS :: Ty prov vs a -> Tys as prov vs b -> Tys (Proxy a ': as) prov vs b
infixr 5 `TysS`

-- type instance SourceOf (Tys vs src ts) = src
-- instance ConstsOf (Tys prov vs ts) where
--  constsOf TysZ = Set.empty
--  constsOf (TysS t ts) = constsOf t `Set.union` constsOf ts
instance UsedVarsOf (Tys as) where
  usedVarsOf vs TysZ k = k vs
  usedVarsOf vs (TysS t ts) k =
    usedVarsOf vs t $ \vs' ->
      usedVarsOf vs' ts k
instance VarOccursIn (Tys as) where
  varOccursIn _v TysZ = False
  varOccursIn v (TysS t ts) = varOccursIn v t Prelude.|| varOccursIn v ts

-- ** Type family 'FunArg'
type family FunArg (fun :: Type) :: Type where
  FunArg (q #> f) = FunArg f
  FunArg (a -> b) = a

-- ** Type family 'FunRes'
type family FunRes (fun :: Type) :: Type where
  FunRes (q #> f) = FunRes f
  FunRes (a -> b) = b

-- -- | Return, when the given 'Type' is a function,
-- -- the argument of that function.
-- unTyFun ::
--  forall t src tys.
--  SourceInj (TyVT src) src =>
--  --Constable (->) =>
--  Type src tys t ->
--  Maybe ( Type src tys (FunArg t)
--        , Type src tys (FunRes t) )
-- unTyFun ty_ini = go ty_ini
--   where
--   go ::
--    Type src tys x ->
--    Maybe ( Type src tys (FunArg x)
--          , Type src tys (FunRes x) )
--   go (TyApp _ (TyApp _ (TyConst _ _ f) a) b)
--    | Just HRefl <- proj_ConstKi @(K (->)) @(->) f
--    = Just ((a `withSource` TyVT ty_ini), (b `withSource` TyVT ty_ini))
--   go (TyApp _ (TyApp _ (TyConst _ _ f) _a) b)
--    | Just HRefl <- proj_ConstKi @(K (#>)) @(#>) f
--    = go b
--   go TyApp{}   = Nothing
--   go TyConst{} = Nothing
--   go TyVar{}   = Nothing
--   go TyFam{}   = Nothing

proj_ConstKi ::
  forall kc (c :: kc) u.
  Typeable c =>
  -- (kc ~ Type_of_Ty (Ty_of_Type kc)) =>
  -- KindInjP (Ty_of_Type kc) =>
  Const u ->
  Maybe (c :~~: u)
proj_ConstKi = eqTypeRep (typeRep @c)

-- ** Type @(#>)@

-- | /Type constant/ to qualify a 'Type'.
newtype (#>) (q :: Constraint) (a :: Type) = Qual (q => a)

-- | Qualify a 'Type'.
(#>) :: Monoid prov => Ty prov vs q -> Ty prov vs t -> Ty prov vs (q #> t)
(#>) q t =
  -- (tyConstLen @(K (#>)) @(#>) (lenVars q) `tyApp` q) `tyApp` t
  TyApp
    { tyAppFun =
        TyApp
          { tyAppFun = TyConst{tyConst = typeRep @(#>), tyConstLen = lenVar q}
          , tyAppArg = q
          }
    , tyAppArg = t
    }

infixr 0 #>

-- NOTE: should actually be (-1)
-- to compose well with @infixr 0 (->)@
-- but this is not allowed by GHC.

-- ** Class @(#)@

-- | 'Constraint' union.
class ((x, y) :: Constraint) => x # y

instance ((x, y) :: Constraint) => x # y

-- | Unify two 'K.Constraint's.
(#) :: Monoid prov => Ty prov vs qx -> Ty prov vs qy -> Ty prov vs (qx # qy)
(#) a b =
  -- (tyConstLen @(K (#)) @(#) (lenVars a) `tyApp` a) `tyApp` b
  TyApp
    { tyAppFun =
        TyApp
          { tyAppFun = TyConst{tyConst = typeRep @(#), tyConstLen = lenVar a}
          , tyAppArg = a
          }
    , tyAppArg = b
    }

infixr 2 #

-- * Class 'KindOf'
class KindOf (ty :: Prov -> [Type] -> aK -> Type) where
  -- | Return the kind of the given 'Var'.
  kindOf ::
    -- HasCallStack =>
    ProvenanceKindOf ty prov =>
    Show prov =>
    ty prov vs (a :: aK) ->
    Kind prov (aK :: Type)

instance KindOf Var where
  kindOf v@Var{varKind} = TyProv (provenanceKindOf v) varKind
instance KindOf Ty where
  kindOf ty = TyProv (provenanceKindOf ty) (go ty)
    where
      go ::
        Show prov =>
        Ty prov vs (a :: aK) ->
        Kind prov (aK :: Type)
      go TyConst{tyConst = constType :: Const (a :: aK)} =
        tyOfTypeRep LenZ (typeRepKind constType)
      go TyApp{tyAppFun = f} = case go f of
        SkipTyProv
          TyApp
            { tyAppFun =
              SkipTyProv
                TyApp
                  { tyAppFun = SkipTyProv TyConst{tyConst = eqTypeRep (typeRep @(->)) -> Just HRefl}
                  }
            , tyAppArg = k
            } -> k
        t -> Prelude.error $ "here: " <> show t
      go TyVar{..} = varKind tyVar
      go TyProv{..} = go tyUnProv

-- ** Class 'ProvenanceKindOf'
class ProvenanceKindOf (ty :: Prov -> [Type] -> aK -> Type) prov where
  provenanceKindOf :: ty prov vs (a :: aK) -> prov
instance ProvenanceKindOf ty () where
  provenanceKindOf _ = ()
instance ProvenanceKindOf Ty String where
  provenanceKindOf = show
instance ProvenanceKindOf Var String where
  provenanceKindOf = show

-- kindOf (TyFam _src _len fam _as) = kindOfConst fam