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

module Symantic.Syntaxes where

import Control.Applicative (Applicative (..))
import Control.Monad (Monad (..))
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 Data.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.Maybe (Maybe (..), isJust)
import Data.Proxy (Proxy (..))
import Data.Semigroup (Semigroup (..))
import Data.String (String)
import GHC.Types
import Text.Read (Read (..), reads)
import Text.Show (Show (..))
import Type.Reflection (TypeRep, Typeable, eqTypeRep, typeRep, (:~~:) (..))
import Unsafe.Coerce (unsafeCoerce)
import Prelude (error)
import Prelude qualified

import Symantic
import Symantic.Typer

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

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)

--
-- Parsing
--

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{tyConst = Const{constType = typeRep, constMeta = ()}, tyConstLen = LenZ} $ 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 eqTy aTy TyConst{tyConst = Const{constType = typeRep @Int, constMeta = ()}, tyConstLen = lenVar aTy} 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 eqTy aTy TyConst{tyConst = Const{constType = typeRep @Int, constMeta = ()}, tyConstLen = lenVar aTy} of
      Nothing -> Left "TypeError"
      Just HRefl -> do
        TermVT bTy (ForallSem b :: ForallSem syn b) <- unParser (final bT)
        case eqTy bTy TyConst{tyConst = Const{constType = typeRep @Int, constMeta = ()}, tyConstLen = lenVar bTy} 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

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 eqTy aTy TyConst{tyConst = Const{constType = typeRep @Int, constMeta = ()}, tyConstLen = lenVar aTy} of
          Nothing -> Left "TypeError"
          Just HRefl -> do
            TermVT bTy (ForallSem b :: ForallSem syn b) <- bE
            case eqTy bTy TyConst{tyConst = Const{constType = typeRep @Int, constMeta = ()}, tyConstLen = lenVar bTy} 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 meta vs a)) bT)) =
  Parser $
    case final bT of
      Parser bE -> do
        TermVT abTy (abF :: ForallSem syn ab) <- bE
        case eqTy TyConst{tyConst = Const{constType = typeRep @Type, constMeta = ()}, tyConstLen = lenVar aTy} aTy of
          Nothing -> Left "TypeError"
          Just HRefl -> do
            case abTy of
              FUN iTy bTy ->
                case eqTy TyConst{tyConst = Const{constType = typeRep @Type, constMeta = ()}, tyConstLen = lenVar bTy} bTy of
                  Nothing -> Left "TypeError"
                  Just HRefl -> do
                    let (aTy', iTy') = appendVars aTy iTy
                    case eqTy 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

--
-- Printing
--

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{tyConst = Const{constType = typeRep @a, constMeta = ()}, tyConstLen = LenZ})
            (unPrinter (f (Printer (BinTree0 (TokenTermAtom "x")))))
        )

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

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

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

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