{-# 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 GHC.Types (Constraint, Type)
import Symantic.Syntaxes.Classes (Unabstractable (..))
import Unsafe.Coerce (unsafeCoerce)

import Symantic.Typer.Type (Ty)

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

-- | Merge syntax 'Constraint's into a single 'Constraint'.
type family Syntaxes (syns :: [(Type -> Type) -> Constraint]) (sem :: Type -> Type) :: Constraint where
  Syntaxes '[] sem = (() :: Constraint)
  Syntaxes (syn ': syns) sem = ((syn sem, Syntaxes syns sem) :: Constraint)

-- * Type 'ForallSemantic'
data ForallSemantic (syns :: [Syntax]) (a :: Type)
  = ForallSemantic (forall sem. Syntaxes syns sem => sem a)

-- * Class 'Abstractable'
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 'Constable'
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. Syntaxes syns sem => Constable sem) => Constable (ForallSemantic syns) where
  constI = ForallSemantic constI
  constK = ForallSemantic constK
  constS = ForallSemantic constS
  constB = ForallSemantic constB
  constC = ForallSemantic constC
instance
  ( forall sem. Syntaxes syns 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 (ForallSemantic syns)
  where
  lam aTy f = ForallSemantic (lam aTy (\a -> let ForallSemantic b = f (forallSem a) in b))
    where
      -- Safe here because (a :: sem a) and (b :: sem b) share the same 'sem'.
      forallSem :: sem a -> ForallSemantic syns a
      forallSem a = ForallSemantic (unsafeCoerce a)
instance
  (forall sem. Syntaxes syns sem => Unabstractable sem) =>
  Unabstractable (ForallSemantic syns)
  where
  ForallSemantic a2b .@ ForallSemantic a = ForallSemantic (a2b .@ a)

-- * Type 'OpenTerm'
data OpenTerm (syns :: [Syntax]) (vs :: [Type]) (a :: Type) where
  -- E contains embedded closed (i.e. already abstracted) terms.
  E :: ForallSemantic syns a -> OpenTerm syns vs a
  -- V represents a reference to the innermost/top environment
  -- variable, i.e. Z
  V :: OpenTerm syns (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 :: OpenTerm syns vs (a -> b) -> OpenTerm syns (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 :: OpenTerm syns vs b -> OpenTerm syns (a ': vs) b
instance
  ( forall sem. Syntaxes syns sem => Abstractable meta sem
  , Syntaxes syns (ForallSemantic syns)
  ) =>
  Abstractable meta (OpenTerm syns '[])
  where
  lam aTy f = E (lam aTy (runOpenTerm . f . E))
instance
  ( forall sem. Syntaxes syns sem => Constable sem
  , Syntaxes syns (ForallSemantic syns)
  ) =>
  Constable (OpenTerm syns vs)
  where
  constI = E constI
  constK = E constK
  constS = E constS
  constB = E constB
  constC = E constC

-- * Type 'R'
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

runOpenTerm :: OpenTerm syns '[] a -> ForallSemantic syns a
runOpenTerm t = case t of E t' -> t'

eval :: Syntaxes syns R => OpenTerm syns '[] a -> a
eval t
  | ForallSemantic sem <- runOpenTerm t =
      unR sem

instance
  ( forall sem. Syntaxes syns sem => Constable sem
  , forall sem. Syntaxes syns sem => Unabstractable sem
  , Syntaxes syns (ForallSemantic syns)
  -- , forall a as. syns (OpenTerm syns (as)) => syns (OpenTerm syns (a ': as))
  ) =>
  Unabstractable (OpenTerm syns vs)
  where
  (.@) = appOpenTerm

appOpenTerm ::
  forall syns as a b.
  ( forall sem. Syntaxes syns sem => Constable sem
  , forall sem. Syntaxes syns sem => Unabstractable sem
  , Syntaxes syns (ForallSemantic syns)
  ) =>
  OpenTerm syns as (a -> b) ->
  OpenTerm syns as a ->
  OpenTerm syns as b
W e1 `appOpenTerm` W e2 = W (e1 `appOpenTerm` e2)
W e `appOpenTerm` E d = W (e `appOpenTerm` E d)
E d `appOpenTerm` W e = W (E d `appOpenTerm` e)
W e `appOpenTerm` V = N e
V `appOpenTerm` W e = N (E (constC .@ constI) `appOpenTerm` e)
W e1 `appOpenTerm` N e2 = N (E constB `appOpenTerm` e1 `appOpenTerm` e2)
N e1 `appOpenTerm` W e2 = N (E constC `appOpenTerm` e1 `appOpenTerm` e2)
N e1 `appOpenTerm` N e2 = N (E constS `appOpenTerm` e1 `appOpenTerm` e2)
N e `appOpenTerm` V = N (E constS `appOpenTerm` e `appOpenTerm` E constI)
V `appOpenTerm` N e = N (E (constS .@ constI) `appOpenTerm` e)
E d `appOpenTerm` N e = N (E (constB .@ d) `appOpenTerm` e)
E d `appOpenTerm` V = N (E d)
V `appOpenTerm` E d = N (E (constC .@ constI .@ d))
N e `appOpenTerm` E d = N (E (constC .@ constC .@ d) `appOpenTerm` e)
E d1 `appOpenTerm` E d2 = E (d1 .@ d2)

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