{-# LANGUAGE ConstraintKinds #-}
{-# LANGUAGE DataKinds #-}
{-# LANGUAGE DeriveFunctor #-}
{-# LANGUAGE AllowAmbiguousTypes #-}
{-# LANGUAGE NoImplicitPrelude #-}
{-# LANGUAGE NoMonomorphismRestriction #-}
{-# LANGUAGE QuantifiedConstraints #-}
{-# LANGUAGE RankNTypes #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE StandaloneDeriving #-}
{-# LANGUAGE TypeFamilies #-}
{-# LANGUAGE UndecidableInstances #-}
{-# LANGUAGE UndecidableSuperClasses #-}
{-# LANGUAGE ImportQualifiedPost #-}
{-# LANGUAGE PatternSynonyms #-}
{-# LANGUAGE StandaloneKindSignatures #-}
{-# LANGUAGE PolyKinds #-}
{-# LANGUAGE DataKinds #-}
{-# LANGUAGE ViewPatterns #-}

module ForallSem where

import Data.Functor (Functor(..), (<$>))
import Data.Functor.Constant (Constant(..))
import Control.Applicative (Applicative(..))
import Control.Monad (Monad(..))
import Data.Kind (Type, Constraint)
import Data.Function (($), (.), id)
import Data.String (String)
import Data.Semigroup (Semigroup(..))
import Data.Maybe (Maybe(..), isJust)
import Text.Read (Read(..), reads)
import Data.Bool (otherwise)
import Text.Show (Show(..))
import Data.Eq (Eq(..))
import Data.Either (Either(..))
import Data.Int (Int)
import Prelude (error)
import Control.Monad.Trans.Except qualified as MT
import Control.Monad.Trans.Reader qualified as MT
import Control.Monad.Trans.State qualified as MT
import Type.Reflection
import Unsafe.Coerce (unsafeCoerce)
import Data.Proxy (Proxy(..))
import Prelude qualified
import GHC.Types

type Semantic = Type -> Type
type Syntax = Semantic -> Constraint
class Syn0 (sem::Semantic) where
  syn0 :: sem ()
class Syn1 (sem::Semantic) where
  syn1 :: sem ()
class Syn2 (sem::Semantic) where
  syn2 :: sem () -> sem () -> sem ()

p0 :: (forall sem. syn sem => Syn0 sem) => ForallSem syn ()
p0 = ForallSem syn0

p1 :: (forall sem. syn sem => Syn1 sem) => ForallSem syn ()
p1 = ForallSem syn1

p3 :: forall syn.
      ( forall sem. syn sem => Syn0 sem
      , forall sem. syn sem => Syn1 sem
      , forall sem. syn sem => Syn2 sem
      ) => ForallSem syn ()
p3 = ForallSem $ syn2 (unForallSem @syn p0) (unForallSem @syn p1)







data ForallSem (syn::Syntax) (a::Type)
   = ForallSem { unForallSem :: forall sem. syn sem => sem a }

instance ( forall sem. syn sem => Functor sem
         ) => Functor (ForallSem syn) where
  fmap f (ForallSem sem) = ForallSem (fmap f sem)
  a <$ (ForallSem sem) = ForallSem (a <$ sem)
instance ( forall sem. syn sem => Applicative sem
         , Functor (ForallSem syn)
         ) => Applicative (ForallSem syn) where
  pure a = ForallSem (pure a)
  liftA2 f (ForallSem a) (ForallSem b) = ForallSem (liftA2 f a b)
  (<*>) (ForallSem f) (ForallSem a) = ForallSem ((<*>) f a)
  (<*) (ForallSem f) (ForallSem a) = ForallSem ((<*) f a)
  (*>) (ForallSem f) (ForallSem a) = ForallSem ((*>) f a)
instance ( forall sem. syn sem => Monad sem
         , Applicative (ForallSem syn)
         ) => Monad (ForallSem syn) where
  (>>=) (ForallSem sa) f = ForallSem (sa >>= \a -> case f a of ForallSem sb -> sb)
  return = pure
  (>>) = (*>)

-- * Interpreter 'Parser'
data TermVT syn = forall vs a. TermVT
  { typeTermVT :: Ty vs a
  , unTermVT  :: ForallSem syn a
  }

-- instance Eq (TyVT) where
--   TyVT x == TyVT y =
--     let (x', y') = appendVars x y in
--     isJust $ eqTypeKi x' y'

-- * Type family @(++)@
type family (++) xs ys where
  '[] ++ ys  = ys
  -- xs  ++ '[] = xs
  (x ': xs) ++ ys = x ': xs ++ ys
infixr 5 ++

-- * Type 'Len'
data Len (xs::[k]) where
  LenZ :: Len '[]
  LenS :: Len xs -> Len (x ': xs)

instance Show (Len vs) where
  showsPrec _p = showsPrec 10 . intLen

addLen :: Len a -> Len b -> Len (a ++ b)
addLen LenZ     b = b
addLen (LenS a) b = LenS $ addLen a b

shiftLen :: forall t b a.  Len a -> Len (a ++      b) -> Len (a ++ t ': b)
shiftLen LenZ b = LenS b
shiftLen (LenS a) (LenS b) = LenS $ shiftLen @t @b a b

intLen :: Len xs -> Int
intLen = go 0
  where
  go :: Int -> Len xs -> Int
  go i LenZ     = i
  go i (LenS l) = go (1 Prelude.+ i) l

-- ** Class 'LenInj'
class LenInj (vs::[k]) where
  lenInj :: Len vs
instance LenInj '[] where
  lenInj = LenZ
instance LenInj as => LenInj (a ': as) where
  lenInj = LenS lenInj

type Name = String
-- 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] -> aK -> Type where
  TyConst :: Len vs -> TypeRep a -> Ty vs a
  TyVar :: Name -> Var vs a -> Ty vs a
  TyApp :: Ty vs f -> Ty vs a -> Ty vs (f a)
infixl 9 `TyApp`
deriving instance Show (Ty vs a)
instance AllocVars Ty where
  allocVarsL len (TyConst l c)  = TyConst (addLen len l) c
  allocVarsL len (TyApp f a)    = TyApp (allocVarsL len f) (allocVarsL len a)
  allocVarsL len (TyVar n v)    = TyVar n (allocVarsL len v)
  --allocVarsL len (TyFam l f as) = TyFam (addLen len l) f $ allocVarsL len `mapTys` as
  
  allocVarsR len (TyConst l c)  = TyConst (addLen l len) c
  allocVarsR len (TyApp f a)    = TyApp (allocVarsR len f) (allocVarsR len a)
  allocVarsR len (TyVar n v)    = TyVar n (allocVarsR len v)
  --allocVarsR len (TyFam l f as) = TyFam (addLen l len) f $ allocVarsR len `mapTys` as

data TyVT = forall vs a. TyVT (Ty vs a)
deriving instance Show TyVT
instance Eq TyVT where
  TyVT x == TyVT y = isJust (eqTyKi x' y')
    where (x', y') = appendVars x y

-- | /Heterogeneous type equality/:
-- return two proofs when two 'Type's are equals:
-- one for the type and one for the kind.
eqTyKi :: Ty vs (x::xK) -> Ty vs (y::yK) -> Maybe (x:~~:y)
eqTyKi (TyConst _lx x) (TyConst _ly y) = eqTypeRep x y
eqTyKi (TyVar   _nx x) (TyVar   _ny y) = eqVarKi   x y
eqTyKi (TyApp fx ax) (TyApp fy ay)
 | Just HRefl <- eqTyKi fx fy
 , Just HRefl <- eqTyKi ax ay
 = Just HRefl
-- eqTyKi (TyFam _ _lx fx ax) (TyFam _ _ly fy ay)
--  | Just HRefl <- eqTypeRep fx fy
--  , Just HRefl <- eqTypes ax ay
--  = Just HRefl
-- -- NOTE: 'TyFam' is expanded as much as possible.
-- eqTyKi x@TyFam{} y = eqTyKi (expandFam x) y
-- eqTyKi x y@TyFam{} = eqTyKi x (expandFam y)
eqTyKi _x _y = Nothing

-- data Some :: (Type -> Type) -> Type where
--   Some :: TypeRep a -> t a -> Some t
-- 
-- mapSome :: (forall a. s a -> t a) -> Some s -> Some t
-- mapSome f (Some a t) = Some a (f t)
-- 
-- type SomeTy = Some (Constant ())

-- pattern SomeTy :: TypeRep a -> SomeTy
-- pattern SomeTy a = Some a (Constant ())
-- {-# COMPLETE SomeTy #-}

-- someType :: TypeRep a -> SomeTy
-- someType a = Some a (Constant ())

--infer :: Vars vs -> TermVT syn -> Maybe (TermVT syn)
--infer ctx te = case te of
--Lit i -> return $ Some TTyInt (embed (CInt i))


data Parser syn = Parser
  { unParser ::
      --MT.State (TokenTerm a)
        (Either ErrMsg (TermVT syn))
  }

{-
instance ( forall sem. syn sem => Functor sem
         ) => Functor (Parser syn) where
  fmap f (Parser esa) = Parser $
    case esa of
      Left e -> Left e
      Right (ForallSem sem) -> Right (ForallSem (f <$> sem))
  a <$ Parser esa = Parser $
    case esa of
      Left e -> Left e
      Right (ForallSem sem) -> Right (ForallSem (a <$ sem))
instance ( forall sem. syn sem => Applicative sem
         , Applicative (ForallSem syn)
         , forall err. syn (Either err)
         , syn (Parser syn) -- FIXME: what constraint is still missing to still require that?
         ) => Applicative (Parser syn) where
  pure a = Parser (Right (ForallSem (pure a)))
  liftA2 f (Parser a) (Parser b) = Parser (liftA2 (liftA2 f) a b)
  (<*>) (Parser f) (Parser a) = Parser (liftA2 (<*>) f a)
  (*>) (Parser f) (Parser a) = Parser (liftA2 (*>) f a)
  (<*) (Parser f) (Parser a) = Parser (liftA2 (<*) f a)
instance ( forall sem. syn sem => Monad sem
         , forall err. syn (Either err)
         , syn (Parser syn) -- FIXME: what constraint is still missing to still require that?
         ) => Monad (Parser syn) where
  (>>=) (Parser efsa) f = Parser (efsa >>= \(ForallSem sa) -> sa >>= unParser . f)
-}

-- * Type 'BinTree'
-- | /Binary Tree/.
data BinTree a
 =   BinTree0 a
 |   BinTree2 (BinTree a) (BinTree a)
 deriving (Eq, Show)

type TermAST = BinTree TokenTerm
data TokenTerm
   = TokenTermAtom String
   | TokenTermAbst String TyVT TermAST
   deriving (Show)

tree0 = add (lit 8) (neg (add (lit 1) (lit 2)))
tree1 = mul (add (lit 8) (neg (add (lit 1) (lit 2)))) (lit 2)
tree2 = lam (\(x::sem Int) -> mul (lit 0) (add (lit 1) (lit 2)))
tree0Print = print tree0
tree1Print = print tree1
tree2Print = print tree2

type ErrMsg = String

safeRead :: Read a => String -> Either ErrMsg a
safeRead s = case reads s of
               [(x,"")] -> Right x
               _        -> Left $ "Read error: " <> s

instance (forall sem. syn sem => Addable sem) => Addable (ForallSem syn) where
  lit n = ForallSem (lit n)
  neg (ForallSem a) = ForallSem (neg a)
  add (ForallSem a) (ForallSem b) = ForallSem (add a b)
instance (forall sem. syn sem => Mulable sem) => Mulable (ForallSem syn) where
  mul (ForallSem a) (ForallSem b) = ForallSem (mul a b)
instance ( forall sem. syn sem => Abstractable sem
         --, forall sem. syn sem => Typeable sem -- user instance not accepted
         --, forall s1 s2. (syn s1, syn s2) => s1 ~ s2 -- crazy...
         ) => Abstractable (ForallSem syn) where
  ForallSem a2b .@ ForallSem a = ForallSem (a2b .@ a)
  lam f = ForallSem (lam (\a -> let ForallSem b = f (forallSem a) in b))
    where
    -- Safe here because a :: sem a and b :: sem b share the same 'sem'.
    forallSem :: sem a -> ForallSem syn a
    forallSem a = ForallSem (unsafeCoerce a)



parseAddable ::
 ( forall sem. syn sem => Addable sem
 , syn (ForallSem syn)
 ) =>
 (TermAST -> Parser syn) -> TermAST -> Parser syn
parseAddable _final (BinTree0 (TokenTermAtom s))
  | Right (i::Int) <- safeRead s
  = Parser $ Right $ TermVT (TyConst LenZ typeRep) $ ForallSem @_ @Int $ lit i
parseAddable final (BinTree2 (BinTree0 (TokenTermAtom "Neg")) aT) = Parser $
  case final aT of
    Parser aE -> do
      TermVT aTy (ForallSem a :: ForallSem syn a) <- aE
      case eqTyKi aTy (TyConst (lenVars aTy) (typeRep @Int)) of
        Nothing -> Left "TypeError"
        Just HRefl -> do
          Right $ TermVT aTy $ neg a
parseAddable final (BinTree2 (BinTree2 (BinTree0 (TokenTermAtom "Add")) aT) bT) =
  Parser $ do
    TermVT aTy (ForallSem a :: ForallSem syn a) <- unParser (final aT)
    case eqTyKi aTy (TyConst (lenVars aTy) (typeRep @Int)) of
      Nothing -> Left "TypeError"
      Just HRefl -> do
        TermVT bTy (ForallSem b :: ForallSem syn b) <- unParser (final bT)
        case eqTyKi bTy (TyConst (lenVars bTy) (typeRep @Int)) of
          Nothing -> Left "TypeError"
          Just HRefl -> do
            Right $ TermVT aTy $ ForallSem @_ @Int $ add a b
parseAddable _final e = Parser $ Left $ "Invalid tree: " <> show e

--tree0Parser :: (forall sem. syn sem => Addable sem, syn (ForallSem syn)) => Either ErrMsg (ForallSem syn Int)
--tree0Parser = unParser $ parse tree0Print

class    (Addable sem, Mulable sem, Abstractable sem) => FinalSyntaxes sem
instance (Addable sem, Mulable sem, Abstractable sem) => FinalSyntaxes sem

parseMulable ::
  ( forall sem. syn sem => Mulable sem
  , forall sem. syn sem => Addable sem
  , syn (ForallSem syn)
  ) =>
  (TermAST -> Parser syn) -> TermAST -> Parser syn
parseMulable final (BinTree2 (BinTree2 (BinTree0 (TokenTermAtom "Mul")) aT) bT) =
  Parser $ do
    case (final aT, final bT) of
      (Parser aE, Parser bE) -> do
        TermVT aTy (ForallSem a :: ForallSem syn a) <- aE
        case eqTyKi aTy (TyConst (lenVars aTy) (typeRep @Int)) of
          Nothing -> Left "TypeError"
          Just HRefl -> do
            TermVT bTy (ForallSem b :: ForallSem syn b) <- bE
            case eqTyKi bTy (TyConst (lenVars bTy) (typeRep @Int)) of
              Nothing -> Left "TypeError"
              Just HRefl -> do
                    Right $ TermVT aTy (ForallSem @_ @a $ mul a b)
parseMulable final t = parseAddable final t

parseAbstractable ::
  ( forall sem. syn sem => Mulable sem
  , forall sem. syn sem => Addable sem
  , forall sem. syn sem => Abstractable sem
  , syn (ForallSem syn)
  ) =>
  a ~ Int =>
  (TermAST -> Parser syn) -> TermAST -> Parser syn
-- TODO: parse variables
parseAbstractable final (BinTree0 (TokenTermAbst n (TyVT (aTy::Ty vs a)) bT)) =
  Parser $
    case final bT of
      Parser bE -> do
            TermVT abTy (abF :: ForallSem syn ab) <- bE
            case eqTyKi (TyConst (lenVars aTy) (typeRep @Type)) aTy of
              Nothing -> Left "TypeError"
              Just HRefl -> do
                case abTy of
                  FUN iTy bTy ->
                    case eqTyKi (TyConst (lenVars bTy) (typeRep @Type)) bTy of
                      Nothing -> Left "TypeError"
                      Just HRefl -> do
                        -- FIXME: appendVars
                        case eqTyKi aTy iTy of
                          Nothing -> Left "TypeError"
                          Just HRefl -> do
                            case abF of
                              ForallSem ab ->
                                Right $ TermVT abTy $ ForallSem @_ @ab $
                                  -- FIXME: Tyable
                                  lam (ab .@)
                  _ -> Left "TypeError"
parseAbstractable final t = parseMulable final t

pattern FUN :: forall k (fun :: k) vs. ()
            => forall (r1 :: RuntimeRep) (r2 :: RuntimeRep)
                      (arg :: TYPE r1) (res :: TYPE r2).
               (k ~ Type, fun ~~ (arg -> res))
            => Ty vs arg
            -> Ty vs res
            -> Ty vs fun
pattern FUN arg res <- (TyConst _len (eqTypeRep (typeRep @(->)) -> Just HRefl) `TyApp` arg) `TyApp` res
  where FUN arg res = (TyConst (lenVars arg) (typeRep @(->)) `TyApp` arg) `TyApp` res

-- * Class 'LenVars'
-- | Return the 'Len' of the 'Var' context.
class LenVars a where
  lenVars :: a -> Len (VarsOf a)

type instance VarsOf (Ty vs a) = vs
instance LenVars (Ty vs a) where
  lenVars (TyConst len _c) = len
  lenVars (TyApp f _a)     = lenVars f
  lenVars (TyVar _n v)     = lenVars v
  -- lenVars (TyFam l _f _as) = l

-- * Type family 'VarsOf'
-- | Return the 'Var' context of something.
type family   VarsOf a :: [Type]
type instance VarsOf (Var vs v) = vs
--type instance VarsOf (UsedVars src vs vs') = vs'



-- * Class 'AllocVars'
-- | Allocate 'Var's in a 'Var' context,
-- either to the left or to the right.
class AllocVars (a::[Type] -> k -> Type) where
  allocVarsL :: Len len -> a vs x -> a (len ++ vs) x
  allocVarsR :: Len len -> a vs x -> a (vs ++ len) x

appendVars ::
 AllocVars a =>
 AllocVars b =>
 LenVars (a vs_x x) =>
 LenVars (b vs_y y) =>
 VarsOf (a vs_x x) ~ vs_x =>
 VarsOf (b vs_y y) ~ vs_y =>
 a vs_x (x::kx) ->
 b vs_y (y::ky) ->
 ( a (vs_x ++ vs_y) x
 , b (vs_x ++ vs_y) y )
appendVars x y =
  ( allocVarsR (lenVars y) x
  , allocVarsL (lenVars x) y
  )



-- Var
data Var (vs::[Type]) (v::vK) :: Type where
  VarZ :: TypeRep vk -> Len (Proxy (v::vK) ': vs) -> Var (Proxy (v::vK) ': vs) v
  VarS :: Var vs v -> Var (not_v ': vs) v
instance AllocVars Var where
  allocVarsL LenZ v       = v
  allocVarsL (LenS len) v = VarS (allocVarsL len v)

  allocVarsR len (VarZ k l) = VarZ k (addLen l len)
  allocVarsR len (VarS v)   = VarS (allocVarsR len v)

data VarV vs = forall v. VarV (Var vs v)
instance LenVars (Var vs v) where
  lenVars (VarZ _k l) = l
  lenVars (VarS v)    = LenS (lenVars v)

deriving instance Show (Var vs v)

eqVarKi :: Var vs x -> Var vs y -> Maybe ((x::kx) :~~: (y::ky))
eqVarKi VarZ{} VarZ{}     = Just HRefl
eqVarKi (VarS x) (VarS y) = eqVarKi x y
eqVarKi _ _               = Nothing

-- Vars
data Vars (vs::[Type]) :: Type where
  -- VarsZ represents an empty context.
  VarsZ :: Vars '[]
  -- A cons stores a variable name and its type,
  -- and then the rest of the context.
  VarsS :: String -> TypeRep vK -> Vars vs -> Vars (Proxy (v::vK) ': vs)
type instance VarsOf (Vars vs) = vs
instance LenVars (Vars vs) where
  lenVars VarsZ = LenZ
  lenVars (VarsS _n _kv vs) = LenS $ lenVars vs


-- Vars
data Env :: [Type] -> Type where
  ENil :: Env '[]
  ECons :: v -> Env vs -> Env (v ': vs)

-- (!) :: Env vs -> Var vs v -> v
-- (ECons x _) ! VarZ = x
-- (ECons _ e) ! (VarS x) = e ! x
-- ENil ! _ = error "GHC can't tell this is impossible"

lookupVars :: String -> Vars vs -> Maybe (VarV vs)
lookupVars _ VarsZ = Nothing
lookupVars n (VarsS vN vK vs)
  | n == vN = Just (VarV (VarZ vK (LenS (lenVars vs))))
  | otherwise = (\(VarV v) -> VarV (VarS v)) <$> lookupVars n vs


fix :: (a -> a) -> a
fix f = f (fix f)

parse ::
  ( forall sem. syn sem => Addable sem
  , forall sem. syn sem => Mulable sem
  , forall sem. syn sem => Abstractable sem
  , syn (ForallSem syn)
  ) =>
  TermAST -> Parser syn
parse = fix parseAbstractable

tree0ParsePrint :: TermAST
tree0ParsePrint = case parse @FinalSyntaxes tree0Print of
  Parser (Left e) -> error e
  Parser (Right (TermVT _ty (ForallSem sem))) -> print sem

tree1ParsePrint :: TermAST
tree1ParsePrint = case parse @FinalSyntaxes tree1Print of
  Parser (Left e) -> error e
  Parser (Right (TermVT _ty (ForallSem sem))) -> print sem

tree2ParsePrint :: TermAST
tree2ParsePrint = case parse @FinalSyntaxes tree2Print of
  Parser (Left e) -> error e
  Parser (Right (TermVT _ty (ForallSem sem))) -> print sem

class Addable sem where
  lit :: Int -> sem Int
  neg :: sem Int -> sem Int
  add :: sem Int -> sem Int -> sem Int
class Mulable sem where
  mul :: sem Int -> sem Int -> sem Int
class Abstractable sem where
  -- | Lambda term abstraction, in HOAS (Higher-Order Abstract Syntax) style.
  lam :: Typeable a => (sem a -> sem b) -> sem (a -> b)
  -- | Application, aka. unabstract.
  (.@) :: sem (a -> b) -> sem a -> sem b


-- * Interpreter 'Printer'
newtype Printer a = Printer { unPrinter :: TermAST }
print :: Printer a -> TermAST
print = unPrinter
print2 :: String -> Printer a1 -> Printer a2 -> Printer a3
print2 n (Printer aT) (Printer bT) = Printer $ BinTree2 (BinTree2 (BinTree0 (TokenTermAtom n)) aT) bT
instance Addable Printer where
  lit n = Printer $ BinTree0 (TokenTermAtom (show n))
  neg (Printer aT) = Printer $ BinTree2 (BinTree0 (TokenTermAtom "Neg")) aT
  add = print2 "Add"
instance Mulable Printer where
  mul = print2 "Mul"
instance Abstractable Printer where
  Printer f .@ Printer a = Printer $ BinTree2 f a
  lam (f::Printer a -> Printer b) = Printer $
                                  -- FIXME: Tyable? instead of Typeable?
      BinTree0 (TokenTermAbst "x" (TyVT (TyConst (lenInj @_ @'[]) (typeRep @a)))
                                  (unPrinter (f (Printer (BinTree0 (TokenTermAtom "x"))))))