{-# 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.Stack (HasCallStack)
import GHC.Types
import Text.Show (Show (..))
import Type.Reflection (TypeRep, Typeable, eqTypeRep, typeRep, typeRepKind, (:~~:) (..), pattern App)
import Prelude (error)
import Prelude qualified

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

data Const meta (a :: aK) =
  -- InferTyable aK =>
  Const
  { constType :: TypeRep a
  , constMeta :: meta
  }

-- deriving (Show)
deriving instance Show meta => Show (Const meta a)
instance EqTy (Const meta) (Const meta) where
  eqTy Const{constType = x} Const{constType = y} = eqTypeRep x y

type Name = String

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

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

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

instance EqTy (Var meta vs) (Var meta 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 meta vs = forall v. VarV (Var meta vs v)

class LenVar var where
  lenVar :: var meta 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 meta vs x -> var meta (len ++ vs) x
  allocVarsR :: Len len -> var meta vs x -> var meta (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 meta xVS (x :: xK) ->
  yVar meta yVS (y :: yK) ->
  ( xVar meta (xVS ++ yVS) x
  , yVar meta (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 meta vs v -> Int
indexVar VarZ{} = 0
indexVar (VarS v) = 1 Prelude.+ indexVar v

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

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

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

insertVar :: Name -> VarV meta vs -> Vars meta vs -> Vars meta 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 meta newVS a where
  UsedVarsZ :: UsedVars oldVS meta '[] a
  UsedVarsS ::
    Var meta oldVS (v :: vK) ->
    UsedVars oldVS meta newVS a ->
    UsedVars oldVS meta (Proxy (v :: vK) ': newVS) a

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

lookupUsedVars ::
  Var meta oldVS (v :: vK) ->
  UsedVars oldVS meta newVS a ->
  Maybe (Var meta 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, varMeta} ->
          Just
            VarZ
              { varZNext = lenVar us
              , varZKind = varKind
              , varZMeta = varMeta
              }
    Nothing -> VarS Data.Functor.<$> lookupUsedVars v us

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

-- * Class 'VarOccursIn'

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

-- * Class 'UsedVarsOf'
class UsedVarsOf t where
  usedVarsOf ::
    UsedVars vs meta oldVS a ->
    t meta vs b ->
    (forall newVS. UsedVars vs meta 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. Type -> [Type] -> aK -> Type where
  TyConst :: {tyConst :: Const meta a, tyConstLen :: Len vs} -> Ty meta vs a
  TyVar :: {tyVar :: Var meta vs a, tyVarName :: Name} -> Ty meta vs a
  TyApp :: {tyAppFun :: Ty meta vs f, tyAppArg :: Ty meta vs a} -> Ty meta vs (f a)
infixl 9 `TyApp`

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

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

pattern FUN ::
  -- CReq
  forall k (fun :: k) meta vs.
  () =>
  {-Monoid meta-}
  -- 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 meta vs arg ->
  Ty meta vs res ->
  Ty meta vs fun
pattern FUN arg res <-
  TyApp
    { tyAppFun =
      TyApp
        { tyAppFun = TyConst{tyConst = Const{constType = eqTypeRep (typeRep @(->)) -> Just HRefl}}
        , tyAppArg = arg
        }
    , tyAppArg = res
    }

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

deriving instance Show meta => Show (Ty meta 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

-- 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 (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 (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 (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 (TyFam _src _len _fam as) k = usedVarsOf vs as k

-- ** Type 'TyT'

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

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

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

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

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

deriving instance Show meta => Show (TyVT meta)
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 meta vs) (Ty meta vs) where
  eqTy = go
    where
      go :: forall xK xT yK yT. Ty meta vs (xT :: xK) -> Ty meta 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 _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]) meta vs (a :: aK) where
  TysZ :: Tys '[] meta vs a
  TysS :: Ty meta vs a -> Tys as meta vs b -> Tys (Proxy a ': as) meta vs b
infixr 5 `TysS`

-- type instance SourceOf (Tys vs src ts) = src
-- instance ConstsOf (Tys meta 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 meta.
  Typeable c =>
  -- (kc ~ Type_of_Ty (Ty_of_Type kc)) =>
  -- KindInjP (Ty_of_Type kc) =>
  Const meta u ->
  Maybe (c :~~: u)
proj_ConstKi = eqTypeRep (typeRep @c) . constType

-- ** Type @(#>)@

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

-- | Qualify a 'Type'.
(#>) :: Monoid meta => Ty meta vs q -> Ty meta vs t -> Ty meta vs (q #> t)
(#>) q t =
  -- (tyConstLen @(K (#>)) @(#>) (lenVars q) `tyApp` q) `tyApp` t
  TyApp
    { tyAppFun =
        TyApp
          { tyAppFun = TyConst{tyConst = Const{constType = typeRep @(#>), constMeta = mempty}, 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 meta => Ty meta vs qx -> Ty meta vs qy -> Ty meta vs (qx # qy)
(#) a b =
  -- (tyConstLen @(K (#)) @(#) (lenVars a) `tyApp` a) `tyApp` b
  TyApp
    { tyAppFun =
        TyApp
          { tyAppFun = TyConst{tyConst = Const{constType = typeRep @(#), constMeta = mempty}, tyConstLen = lenVar a}
          , tyAppArg = a
          }
    , tyAppArg = b
    }

infixr 2 #

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

instance KindOf Var where
  kindOf Var{varKind} = varKind
instance KindOf Ty where
  kindOf TyConst{tyConst = Const{..} :: Const meta (a :: aK)} =
    tyOfTypeRep constMeta LenZ (typeRepKind constType)
  kindOf TyApp{tyAppFun = f} = case kindOf f of
    TyApp
      { tyAppFun =
        TyApp
          { tyAppFun = TyConst{tyConst = Const{constType = eqTypeRep (typeRep @(->)) -> Just HRefl}}
          }
      , tyAppArg = k
      } -> k
    t -> error $ "here: " <> show t
  kindOf TyVar{tyVar = Var{varKind, varMeta}} = varKind

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