{-# 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.Applicative (Applicative (..), liftA3)
import Control.Monad (Monad (..), sequence)
import Control.Monad.Classes as MC ()
import Control.Monad.Classes.Run as MC ()
import Control.Monad.Trans.Except qualified as MT
import Control.Monad.Trans.Identity (IdentityT (..))
import Control.Monad.Trans.Reader qualified as MT
import Control.Monad.Trans.State qualified as MT
import Control.Monad.Trans.Writer.CPS qualified as MT
import Data.Bool (Bool, otherwise)
import Data.Either (Either (..))
import Data.Eq (Eq (..))
import Data.Function (id, ($), (.))
import Data.Functor (Functor (..), (<$>))
import Data.Functor.Constant (Constant (..))
import Data.Int (Int)
import Data.Kind (Constraint, Type)
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.Set qualified as Set
import Data.String (String)
import GHC.Types
import Text.Read (Read (..), reads)
import Text.Show (Show (..))
import Type.Reflection (TypeRep, Typeable, eqTypeRep, typeRep, typeRepKind, withTypeable, (:~~:) (..))
import Unsafe.Coerce (unsafeCoerce)
import Prelude (error)
import Prelude qualified

import Symantic.Parser.Source
import Symantic.Typer.Eq
import Symantic.Typer.List

data Const meta a = Const
  { constType :: TypeRep a
  , constMeta :: meta
  }
  deriving (Show)
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`

-- ** Class 'Tyable'
class Tyable a where
  tyOf :: meta -> Len vs -> Ty meta vs a
instance (Tyable a, Tyable b) => Tyable (a b) where
  tyOf meta vs = TyApp (tyOf @a meta vs) (tyOf @b meta vs)
instance {-# OVERLAPPABLE #-} Typeable a => Tyable a where
  tyOf meta vs = TyConst{tyConst = Const{constType = typeRep, constMeta = meta}, tyConstLen = vs}

typeOf :: Monoid meta => Tyable a => Ty meta '[] a
typeOf = tyOf mempty LenZ

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 = tyOf @(->) mempty (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 TyConst{tyConst = x} TyConst{tyConst = y} = eqTy x y
  eqTy TyVar{tyVar = x} TyVar{tyVar = y} = eqTy x y
  eqTy TyApp{tyAppFun = xF, tyAppArg = xA} TyApp{tyAppFun = yF, tyAppArg = yA}
    | Just HRefl <- eqTy xF yF
    , Just HRefl <- eqTy xA yA =
        Just HRefl
  eqTy _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

-- HRefl <- eqTy ku kc -- $ kindInjP @(Ty_of_Type kc) ()
-- HRefl :: u:~~:c <- eqT
-- Just HRefl

-- ** 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 ::
    -- Provenanceable (ProvenanceKindOf t) m =>
    -- MC.MonadWriter (ProvenanceKindOf t) m =>
    Show meta =>
    ty meta vs (a :: aK) ->
    Kind meta aK

instance KindOf Var where
  kindOf Var{varKind} = varKind
instance KindOf Ty where
  kindOf TyConst{tyConst = Const{constType, constMeta} :: Const meta (a :: aK)} =
    withTypeable (typeRepKind constType) $
      tyOf @aK constMeta LenZ
  kindOf TyApp{tyAppFun = f} = case kindOf f of
    FUN _ kf -> kf
    t -> error $ "here: " <> show t
  kindOf TyVar{tyVar = Var{varKind, varMeta}} = varKind

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