{-# LANGUAGE ConstraintKinds #-}
{-# LANGUAGE InstanceSigs #-}
{-# LANGUAGE QuantifiedConstraints #-}
{-# LANGUAGE RankNTypes #-}
{-# LANGUAGE UndecidableInstances #-}

module Symantic.Compiler.Term where

import Control.Applicative (Applicative (..))
import Control.Monad (Monad (..))
import Data.Function ((.))
import Data.Function qualified as Fun
import Data.Functor (Functor (..))
import Data.Int (Int)
import Data.String (String)
import GHC.Types (Constraint, Type)
import Unsafe.Coerce (unsafeCoerce)

import Symantic.Typer.Type (Const, Ty)

type Semantic = Type -> Type
type Syntax = Semantic -> Constraint

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

-- * Type 'CtxTe'

-- | GADT for an /interpreting context/:
-- accumulating at each /lambda abstraction/
-- the term of the introduced variable.
data CtxTe (sem :: Type -> Type) (as :: [Type]) where
  CtxTeZ :: CtxTe sem '[]
  CtxTeS :: sem a -> CtxTe sem as -> CtxTe sem (a ': as)

infixr 5 `CtxTeS`

class Addable sem where
  lit :: Int -> sem Int
  neg :: sem (Int -> Int)
  add :: sem (Int -> Int -> Int)

class Mulable sem where
  mul :: sem Int -> sem Int -> sem Int

class Abstractable meta sem where
  -- type AbstractableLam sem (a::Type) :: Constraint
  -- type AbstractableLam sem a = (()::Constraint)

  -- | Lambda term abstraction, in HOAS (Higher-Order Abstract Syntax) style.
  lam :: Ty meta '[] a -> (sem a -> sem b) -> sem (a -> b)

class UnAbstractable sem where
  -- | Application, aka. unabstract.
  (.@) :: sem (a -> b) -> sem a -> sem b

instance (forall sem. syn sem => Constable sem) => Constable (ForallSem syn) where
  constI = ForallSem constI
  constK = ForallSem constK
  constS = ForallSem constS
  constB = ForallSem constB
  constC = ForallSem constC
instance (forall sem. syn sem => Addable sem) => Addable (ForallSem syn) where
  lit n = ForallSem (lit n)
  neg = ForallSem neg
  add = ForallSem add
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 meta sem
  -- , forall sem a. syn sem => AbstractableLam sem a
  -- , forall sem. syn sem => AbstractableLam sem a
  -- , forall sem. syn sem => Typeable sem -- user instance not accepted
  -- , forall s1 s2. (syn s1, syn s2) => s1 ~ s2 -- crazy...
  ) =>
  Abstractable meta (ForallSem syn)
  where
  lam aTy f = ForallSem (lam aTy (\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)
instance
  (forall sem. syn sem => UnAbstractable sem) =>
  UnAbstractable (ForallSem syn)
  where
  ForallSem a2b .@ ForallSem a = ForallSem (a2b .@ a)

data TTerm (syn :: Syntax) (vs :: [Type]) (a :: Type) where
  -- TVar :: Var vs a -> TTerm syn vs a
  TConst :: ForallSem syn a -> TTerm syn vs a
  TLam :: TTerm syn (a ': vs) b -> TTerm syn vs (a -> b)
  TApp :: TTerm syn vs (a -> b) -> TTerm syn vs a -> TTerm syn vs b

data OTerm (syn :: Syntax) (vs :: [Type]) (a :: Type) where
  -- E contains embedded closed (i.e. already abstracted) terms.
  E :: ForallSem syn a -> OTerm syn vs a
  -- V represents a reference to the innermost/top environment
  -- variable, i.e. Z
  V :: OTerm syn (a ': vs) a
  -- N represents internalizing the innermost bound variable as a
  -- function argument. In other words, we can represent an open
  -- term referring to a certain variable as a function which
  -- takes that variable as an argument.
  N :: OTerm syn vs (a -> b) -> OTerm syn (a ': vs) b
  -- For efficiency, there is also a special variant of N for the
  -- case where the term does not refer to the topmost variable at all.
  W :: OTerm syn vs b -> OTerm syn (a ': vs) b
instance
  ( forall sem. syn sem => Abstractable meta sem
  , syn (ForallSem syn)
  ) =>
  Abstractable meta (OTerm syn '[])
  where
  lam aTy f = E (lam aTy (runOTerm . f . E))
instance
  ( forall sem. syn sem => Constable sem
  , syn (ForallSem syn)
  ) =>
  Constable (OTerm syn vs)
  where
  constI = E constI
  constK = E constK
  constS = E constS
  constB = E constB
  constC = E constC

newtype R a = R {unR :: a}
instance Abstractable meta R where
  lam _aTy f = R (unR . f . R)
instance UnAbstractable R where
  R f .@ R x = R (f x)
instance Constable R where
  constI = R Fun.id
  constK = R Fun.const
  constS = R (<*>)
  constB = R (.)
  constC = R Fun.flip

runOTerm :: OTerm syn '[] a -> ForallSem syn a
runOTerm t = case t of E t' -> t'

eval :: syn R => OTerm syn '[] a -> a
eval t
  | ForallSem sem <- runOTerm t =
      unR sem

class Constable sem where
  constI :: sem (a -> a)
  constK :: sem (a -> b -> a)
  constS :: sem ((a -> b -> c) -> (a -> b) -> a -> c)
  constB :: sem ((b -> c) -> (a -> b) -> a -> c)
  constC :: sem ((a -> b -> c) -> b -> a -> c)

instance
  ( forall sem. syn sem => Constable sem
  , forall sem. syn sem => UnAbstractable sem
  , syn (ForallSem syn)
  -- , forall a as. syn (OTerm syn (as)) => syn (OTerm syn (a ': as))
  ) =>
  UnAbstractable (OTerm syn vs)
  where
  (.@) = appOTerm

appOTerm ::
  forall syn as a b.
  ( forall sem. syn sem => Constable sem
  , forall sem. syn sem => UnAbstractable sem
  , syn (ForallSem syn)
  ) =>
  OTerm syn as (a -> b) ->
  OTerm syn as a ->
  OTerm syn as b
W e1 `appOTerm` W e2 = W (e1 `appOTerm` e2)
W e `appOTerm` E d = W (e `appOTerm` E d)
E d `appOTerm` W e = W (E d `appOTerm` e)
W e `appOTerm` V = N e
V `appOTerm` W e = N (E (constC .@ constI) `appOTerm` e)
W e1 `appOTerm` N e2 = N (E constB `appOTerm` e1 `appOTerm` e2)
N e1 `appOTerm` W e2 = N (E constC `appOTerm` e1 `appOTerm` e2)
N e1 `appOTerm` N e2 = N (E constS `appOTerm` e1 `appOTerm` e2)
N e `appOTerm` V = N (E constS `appOTerm` e `appOTerm` E constI)
V `appOTerm` N e = N (E (constS .@ constI) `appOTerm` e)
E d `appOTerm` N e = N (E (constB .@ d) `appOTerm` e)
E d `appOTerm` V = N (E d)
V `appOTerm` E d = N (E (constC .@ constI .@ d))
N e `appOTerm` E d = N (E (constC .@ constC .@ d) `appOTerm` e)
E d1 `appOTerm` E d2 = E (d1 .@ d2)

data CONST syn a where
  -- CInt :: Int -> Const Int
  -- CIf :: Const (Bool -> a -> a -> a)
  -- CAdd :: Const (Int -> Int -> Int)
  -- CGt :: Ord a => Const (a -> a -> Bool)
  F :: ForallSem syn a -> CONST syn a
  I :: CONST syn (a -> a)
  K :: CONST syn (a -> b -> a)
  S :: CONST syn ((a -> b -> c) -> (a -> b) -> a -> c)
  B :: CONST syn ((b -> c) -> (a -> b) -> a -> c)
  C :: CONST syn ((a -> b -> c) -> b -> a -> c)

-- conv :: Term vs α -> OTerm vs α
-- conv = \case
--  TVar VZ -> V
--  TVar (VS x) -> W (conv (TVar x))
--  TLam t -> case conv t of
--    V -> E (embed I)
--    E d -> E (K .$$ d)
--    N e -> e
--    W e -> K .$$ e
--  TApp t1 t2 -> conv t1 $$ conv t2
--  TConst c -> embed c

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
  (>>) = (*>)