{-# LANGUAGE DeriveGeneric #-}
{-# LANGUAGE OverloadedLists #-}
{-# LANGUAGE OverloadedStrings #-}
{-# LANGUAGE OverloadedLabels #-}
{-# LANGUAGE TypeFamilies #-}
{-# LANGUAGE TypeOperators #-}
{-# LANGUAGE TemplateHaskell #-}
{-# LANGUAGE PartialTypeSignatures #-}
{-# LANGUAGE DerivingStrategies #-}
{-# LANGUAGE OrPatterns #-}
{-# LANGUAGE PatternSynonyms #-}
{-# LANGUAGE DeriveAnyClass #-}
module Gyehoek.Scheme.Syntax
  ( Name(..)
  , Prim(..)
  , Lit(..)
  , Def(..)
  , Exp(..)
  , ExpF(..)
  , Program(..)
  , CommandOrDef(..)
  , primDatumIso
  , free
  , subst
  , getName
  , scm
  , free'
  , freeWithBound'
  , freeO
  )
  where

import Data.List (intersperse)
import Effectful
import Prelude hiding ((.), id)
import Control.Category
import Gyehoek.Sexp qualified as GS
import Gyehoek.GenSym (Gen)
import Control.Lens
import Data.Generics.Labels ()
import Data.String (IsString)
import Data.Functor.Foldable.TH (makeBaseFunctor)
import Data.Functor.Foldable hiding (fold)
import qualified Data.HashSet as HS
import Data.Foldable (fold, toList)
import Language.Haskell.TH.Quote (QuasiQuoter)
import Effectful.FileSystem (runFileSystem)
import qualified Effectful.FileSystem.IO as FS
import qualified Data.Text.Encoding as T
import qualified Effectful.FileSystem.IO.ByteString as FB
import qualified Data.Set.Ordered as O
import Gyehoek.Sexp.Grammar qualified as Sexp
import Gyehoek.Sexp.Grammar qualified as S
import Gyehoek.Sexp.Grammar (DatumIso, G, DataIso, (:-)((:-)))
import Gyehoek.Prelude


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

instance Prefixed Name where
  prefixed (MkName s) = _Wrapped' . prefixed @Text s . from _Wrapped'

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
  | PrimMakeClosure { code :: e, env :: List e }
  | PrimEnvRef e Int
  | PrimEnvCode e
  | PrimCallCC e
  | PrimValues (List e)
  | PrimCallWithValues e e
  deriving stock (Show, Generic, Functor, Foldable, Traversable, Data, Eq)
  deriving anyclass (NFData)

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

data Lit
  = LitInt Int
  | LitBool Bool
  | LitString Text
  deriving stock (Show, Generic, Data, Eq)
  deriving anyclass (NFData)

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

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

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

newtype Program = MkProgram
  { commandsAndDefs :: List CommandOrDef
  }
  deriving stock (Show, Generic, Data)
  deriving anyclass (NFData)

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 DatumIso Name where
  datumIso = S.decorate S.SynVariable
    >>> S.symbol
    >>> S.iso coerce coerce

primDatumIso
  :: (Text -> Text)
  -> S.DatumGrammar a -> S.DatumGrammar (Prim a)
primDatumIso namefn a = S.match
    $ S.With (. ht2 "+")
    $ S.With (. ht2 "-")
    $ S.With (. ht2 "*")
    $ S.With (. ht2 "/")
    $ S.With (. ht2 "cons")
    $ S.With (. ht1 "car")
    $ S.With (. ht1 "cdr")
    $ S.With (. ht1 "immediate?")
    $ S.With (. ht1 "cons?")
    $ S.With (. ht1 "integer?")
    $ S.With (. ht1 "write")
    $ S.With (. ht1 "zero?")
    $ S.With (. nullop "newline")
    $ S.With (. ht1' "make-closure")
    $ S.With (. S.headTagged2 (namefn "env-ref") a S.int)
    $ S.With (. ht1 "env-code")
    $ S.With (. ht1 "call/cc")
    $ S.With (. ht0' "values")
    $ S.With (. ht2 "call-with-values")
    $ S.End
    where
      idn = S.el . S.sym . namefn
      nullop s = S.list $ idn s
      ht1 s = S.headTagged1 (namefn s) a
      ht2 s = S.headTagged2 (namefn s) a a
      ht1' s = S.headTagged1' (namefn s) a a
      ht0' s = S.headTagged0' (namefn s) a

instance DatumIso a => DatumIso (Prim a) where
  datumIso = primDatumIso id S.datumIso

instance DatumIso Lit where
  datumIso = S.match
    $ S.With (. S.int)
    $ S.With (. S.boolean)
    $ S.With (. S.string)
    $ S.End

instance DatumIso Def where
  datumIso = S.match
    $ S.With (. defconst)
    $ S.With (. defun)
    $ S.End
    where
      defconst = S.list $ S.el (S.sym "define")
        >>> S.el S.datumIso >>> S.el S.datumIso
      defun = S.list $ S.el (S.sym "define")
        >>> S.el args >>> S.rest S.datumIso
      args = S.list $ S.el S.datumIso >>> S.rest S.datumIso
  
instance DatumIso Exp where
  datumIso = S.match
    $ S.With (. S.letLike "let" S.datumIso S.datumIso S.datumIso)
    $ S.With (. S.letLike "letrec" S.datumIso S.datumIso S.datumIso)
    $ S.With (. S.datumIso)
    $ S.With (. begin)
    $ S.With (. S.ifLike "if" S.datumIso S.datumIso S.datumIso)
    $ S.With (. S.datumIso)
    $ S.With (. lam)
    $ S.With (. S.datumIso)
    $ S.With (\app -> app . S.list (S.el S.datumIso >>> S.rest S.datumIso))
    $ S.End
    where
      lam = S.lambdaLike S.lambdaKeyword S.datumIso (S.el S.datumIso)
      begin :: forall t. G (S.Datum :- t) (NonEmpty Exp :- t)
      begin = S.beginLike "begin" $
        S.el (S.datumIso @Exp) >>> S.rest (S.datumIso @Exp)
        >>> S.onTail (S.Iso
          (\(xs:-x:-t) -> (x:|xs):-t)
          (\((x:|xs):-t) -> xs:-x:-t))

instance DatumIso CommandOrDef where
  datumIso = S.match
    $ S.With (\_Command -> _Command . S.datumIso)
    $ S.With (\_Definition -> _Definition . S.datumIso)
    $ S.With (\_Begin -> _Begin . S.beginLike "begin" (S.rest S.datumIso))
    $ S.End

instance DataIso Program where
  dataIso = S.dataIso @(List CommandOrDef) >>> S.iso coerce coerce


-- utilities

scm :: QuasiQuoter
scm = GS.makeSx [|| S.fromDatumUnsafe @Exp S.datumIso ||]

freeWithBound' :: Foldable f => f Name -> Exp -> List Name
freeWithBound' bound = filter (`elem` bound) . free'

freeO :: Exp -> O.OSet Name
freeO = O.unbiased . cata \case
  ExpVarF x -> O.Bias @O.L $ O.singleton x
  ExpLetF bs e ->
    foldOf (each . _2) bs
    <> (e & coerced %~ deleteFromO (bs ^.. each . _1))
  ExpLetRecF bs e ->
    (foldOf (each . _2) bs & coerced %~ deleteFromO binds)
    <> (e & coerced %~ deleteFromO binds)
    where binds = bs ^.. each . _1
  ExpLambdaF bs e -> e & coerced %~ deleteFromO bs
  e -> fold e

free' :: Exp -> List Name
free' = toList @O.OSet . O.unbiased . cata \case
  ExpVarF x -> O.Bias @O.L $ O.singleton x
  ExpLetF bs e ->
    foldOf (each . _2) bs
    <> (e & coerced %~ deleteFromO (bs ^.. each . _1))
  ExpLetRecF bs e ->
    (foldOf (each . _2) bs & coerced %~ deleteFromO binds)
    <> (e & coerced %~ deleteFromO binds)
    where binds = bs ^.. each . _1
  ExpLambdaF bs e -> e & coerced %~ deleteFromO bs
  e -> fold e

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

deleteFromO :: (Foldable f, Ord a) => f a -> O.OSet a -> O.OSet a
deleteFromO = flip $ foldr O.delete

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