{-# LANGUAGE DeriveGeneric #-}
{-# LANGUAGE OverloadedStrings #-}
{-# LANGUAGE OverloadedLabels #-}
{-# LANGUAGE TypeFamilies #-}
{-# LANGUAGE TypeOperators #-}
{-# LANGUAGE TemplateHaskell #-}
{-# LANGUAGE PartialTypeSignatures #-}
{-# LANGUAGE DerivingStrategies #-}
{-# LANGUAGE OrPatterns #-}
{-# LANGUAGE PatternSynonyms #-}
module Gyehoek.Scheme.Syntax
  ( Name(..)
  , Prim(..)
  , Lit(..)
  , Def(..)
  , Exp(..)
  , Sexp(..)
  , Program(..)
  , CommandOrDef(..)
  , primSexpIso
  , pattern Void
  , free
  , subst
  , getName
  , scm
  )
  where

import Data.Text (Text)
import Data.List (List)
import Language.SexpGrammar
  ( SexpIso(..), list, el, rest, sym, symbol )
import Language.SexpGrammar qualified as Sexp
import Language.Sexp.Located qualified as S
import Language.SexpGrammar.Generic
import GHC.Generics
import Prelude hiding ((.), id)
import Control.Category
import Data.List.NonEmpty (NonEmpty)
import Gyehoek.Sexp qualified
import Gyehoek.GenSym (Gen)
import Control.Lens
import Data.String (IsString)
import Data.Hashable (Hashable)
import Data.Data (Data)
import Data.Functor.Foldable.TH (makeBaseFunctor)
import Data.Functor.Foldable hiding (fold)
import Data.HashSet (HashSet)
import qualified Data.HashSet as HS
import Data.Foldable (fold)
import Language.Haskell.TH.Quote (QuasiQuoter)


newtype Name = MkName { inner :: Text }
  deriving newtype (Show, Eq, Ord, IsString, Gen, Hashable)
  deriving stock (Generic, Data)

getName :: Name -> Text
getName (MkName x) = x

data Prim e
  = PrimAdd e e
  | PrimSub e e
  | PrimMul e e
  | PrimDiv e e
  | PrimCons e e
  | PrimCar e
  | PrimCdr e
  | PrimImmediateP e
  | PrimConsP e
  | PrimIntegerP e
  | PrimWrite e
  | PrimZeroP e
  | PrimNewline
  deriving (Show, Generic, Functor, Foldable, Traversable, Data, Eq)

instance Each (Prim e) (Prim e') e e'

data Lit
  = LitInt Int
  | LitNil
  | LitBool Bool
  | LitString Text
  | LitQuote Sexp
  deriving (Show, Generic, Data, Eq)

pattern Void :: Lit
pattern Void = LitNil

data Def
  = DefConstant Name Exp
  | DefProcedure Name (List Name) (List Exp)
  deriving (Show, Generic, Data)

data Exp
  = ExpLet (NonEmpty (Name, Exp)) Exp
  | ExpPrim (Prim Exp)
  | ExpBegin (List Exp)
  | ExpIf Exp Exp Exp
  | ExpLit Lit
  | ExpLambda (List Name) Exp
  | ExpVar Name
  | ExpApply Exp (List Exp)
  deriving (Show, Generic, Data)

data Sexp
  = SexpCons Sexp Sexp
  | SexpSymbol Text
  | SexpLit Lit
  deriving (Show, Generic, Data, Eq)

data CommandOrDef
  = Command Exp
  | Definition Def
  | Begin (List CommandOrDef)
  deriving (Show, Generic, Data)

data Program = MkProgram
  { commandsAndDefs :: List CommandOrDef
  }
  deriving (Show, Generic, Data)

instance Each Program Program (Either Exp Def) (Either Exp Def) where
  each = #commandsAndDefs . each . go
    where
      inj = either Command Definition
      go :: Traversal' CommandOrDef (Either Exp Def)
      go k (Command e) = inj <$> k (Left e)
      go k (Definition d) = inj <$> k (Right d)
      go k (Begin xs) = Begin <$> traverse (go k) xs

makeBaseFunctor ''Exp



instance SexpIso Name where
  sexpIso = symbol >>> Sexp.partialOsi f g
    where
      f = Right . MkName
      g (MkName s) = s

primSexpIso :: (Text -> Text) -> Sexp.SexpGrammar a -> Sexp.SexpGrammar (Prim a)
primSexpIso namefn a = match
    $ With (. binop "+")
    $ With (. binop "-")
    $ With (. binop "*")
    $ With (. binop "/")
    $ With (. binop "cons")
    $ With (. unop "car")
    $ With (. unop "cdr")
    $ With (. unop "immediate?")
    $ With (. unop "cons?")
    $ With (. unop "integer?")
    $ With (. unop "write")
    $ With (. unop "zero?")
    $ With (. nullop "newline")
    $ End
    where
      idn s = el (sym (namefn s))
      nullop s = list $ idn s
      unop s = list $ idn s >>> el a
      binop s = list $ idn s >>> el a >>> el a

instance SexpIso a => SexpIso (Prim a) where
  -- sexpIso = primSexpIso ("prim:"<>) sexpIso
  sexpIso = primSexpIso id sexpIso

instance SexpIso Lit where
  sexpIso = match
    $ With (. sexpIso)
    $ With (. sym "nil")
    $ With (. bool)
    $ With (. sexpIso)
    $ With (. Gyehoek.Sexp.prefixSugar "quote" Sexp.Quote sexpIso)
    $ End
    where
      bool :: Sexp.SexpGrammar Bool
      bool = Sexp.hashed $ Sexp.partialOsi f g
        where
          f (S.Symbol ("t";"true")) = Right True
          f (S.Symbol ("f";"false")) = Right False
          f _ = Left $ Sexp.expected "bool"
          g True = S.Symbol "true"
          g False = S.Symbol "false"

instance SexpIso Sexp where
  sexpIso = match
    $ With (\conss -> conss . Gyehoek.Sexp.todo)
    $ With (\s -> s . symbol)
    $ With (\lit -> lit . sexpIso)
    $ End

instance SexpIso Def where
  sexpIso = match
    $ With (. defconst)
    $ With (. defun)
    $ End
    where
      defconst = list $ el (sym "define") >>> el sexpIso >>> el sexpIso
      defun = list $ el (sym "define") >>> el args >>> rest sexpIso
      args = list $ el sexpIso >>> rest sexpIso
  
instance SexpIso Exp where
  sexpIso = match
    $ With (. Gyehoek.Sexp.let_ "let" sexpIso sexpIso sexpIso)
    $ With (. sexpIso)
    $ With (\bgn -> bgn . list (el (sym "begin") >>> rest sexpIso))
    $ With (. if_)
    $ With (. sexpIso)
    $ With (. lam)
    $ With (. sexpIso)
    $ With (\app -> app . list (el sexpIso >>> rest sexpIso))
    $ End
    where
      if_ = list $ el (sym "if") >>> el sexpIso >>> el sexpIso >>> el sexpIso
      lam = list
        ( el Gyehoek.Sexp.lambdaKeyword
          >>> el (sexpIso @(List Name))
          >>> el sexpIso )

instance SexpIso CommandOrDef where
  sexpIso = match
    $ With (\_Command -> _Command . sexpIso)
    $ With (\_Definition -> _Definition . sexpIso)
    $ With (\_Begin -> _Begin . bgn)
    $ End
    where
      bgn = list $ el (sym "begin") >>> rest sexpIso


-- utilities

scm :: QuasiQuoter
scm = Gyehoek.Sexp.makeSx [|| Gyehoek.Sexp.fromSexp @Exp ||]

free :: Exp -> HashSet Name
free = cata \case
  ExpVarF x -> HS.singleton x
  ExpLetF bs e -> error "todo lol"
  ExpLambdaF binders vs -> deleteFrom binders vs
  e -> fold e

deleteFrom :: (Foldable f, Hashable a) => f a -> HashSet a -> HashSet a
deleteFrom = flip $ foldr HS.delete

insertFrom :: (Foldable f, Hashable a) => f a -> HashSet a -> HashSet a
insertFrom = flip $ foldr HS.insert

subst :: (Name -> Maybe Exp) -> Exp -> Exp
subst f = \e -> cata go e mempty where
  go (ExpVarF x) bound
    | not (x `HS.member` bound), Just e' <- f x = e'
    | otherwise                         = ExpVar x
  go (ExpLetF _ _) _ = error "todo lol"
  go (ExpLambdaF bs e) bound = e $ insertFrom bs bound
  go e bound = embed $ fmap ($ bound) e