{-# 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(..)
  , Sexp(..)
  , Program(..)
  , CommandOrDef(..)
  , primSexpIso
  , pattern Void
  , free
  , subst
  , getName
  , scm
  , readExp
  , readProgram
  , free'
  , freeWithBound'
  , freeO
  , encodeProgram
  )
  where

import Data.List (intersperse)
import Language.SexpGrammar
  ( SexpIso(..), list, el, rest, sym, symbol )
import Language.SexpGrammar qualified as Sexp
import Language.SexpGrammar.Generic
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.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
  deriving stock (Show, Generic, Functor, Foldable, Traversable, Data, Eq)
  deriving anyclass (NFData)

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

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

pattern Void :: Lit
pattern Void = LitNil

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 (List 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 Sexp
  = SexpCons Sexp Sexp
  | SexpSymbol Text
  | SexpLit Lit
  deriving stock (Show, Generic, Data, Eq)
  deriving anyclass (NFData)

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

data 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 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 (. ht2 "+")
    $ With (. ht2 "-")
    $ With (. ht2 "*")
    $ With (. ht2 "/")
    $ With (. ht2 "cons")
    $ With (. ht1 "car")
    $ With (. ht1 "cdr")
    $ With (. ht1 "immediate?")
    $ With (. ht1 "cons?")
    $ With (. ht1 "integer?")
    $ With (. ht1 "write")
    $ With (. ht1 "zero?")
    $ With (. nullop "newline")
    $ With (. ht1' "make-closure")
    $ With (. GS.headTagged2 (namefn "env-ref") a Sexp.int)
    $ With (. ht1 "env-code")
    $ With (. ht1 "call/cc")
    $ End
    where
      idn s = el (sym (namefn s))
      nullop s = list $ idn s
      ht1 s = GS.headTagged1 (namefn s) a
      ht2 s = GS.headTagged2 (namefn s) a a
      ht1' s = GS.headTagged1' (namefn s) a 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 (. GS.schemeBool)
    $ With (. sexpIso)
    $ With (. GS.prefixSugar "quote" Sexp.Quote sexpIso)
    $ End

instance SexpIso Sexp where
  sexpIso = match
    $ With (\conss -> conss . GS.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 (. GS.let_ "let" sexpIso sexpIso sexpIso)
    $ With (. GS.let_ "letrec" 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 GS.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 = GS.makeSx [|| GS.fromSexp @Exp ||]

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



fileName :: FilePath -> FilePath
fileName "-" = "<interactive>"
fileName e = e

hGetContents :: FS.FileSystem :> es => FS.Handle -> Eff es Text
hGetContents h = T.decodeUtf8 <$> FB.hGetContents h

readProgram :: IOE :> es => FilePath -> Eff es Program
readProgram fp = runFileSystem $
  FS.withFile fp FS.ReadMode $ \h ->
    GS.parseSexps @CommandOrDef (fileName fp) <$> hGetContents h
     >>= either error (pure . MkProgram)

readExp :: IOE :> es => FilePath -> Eff es Exp
readExp fp = readProgram fp <&> (^?! #commandsAndDefs . _head . #Command)

encodeProgram :: Program -> Text
encodeProgram p = p.commandsAndDefs
  & fmap ((^?! _Right) . GS.encodePretty)
  & intersperse "\n\n"
  & mconcat