{-# LANGUAGE OverloadedLabels #-}
{-# LANGUAGE TemplateHaskell #-}
{-# LANGUAGE PatternSynonyms #-}
{-# LANGUAGE ViewPatterns #-}
{-# LANGUAGE TypeFamilies #-}
{-# LANGUAGE FunctionalDependencies #-}
{-# LANGUAGE DeriveAnyClass #-}
module Gyehoek.CPS.Syntax
  ( Val(..)
  , Kappa(..)
  , Lambda(..)
  , Exp(..)
  , ExpF(..)
  , Name(..)
  , Prim(..)
  , Program(..)
  , Lit(..)
  , Imm(..)
  , Obj(..)
  , Hob(..)
  , pattern Void
  , pattern Halt
  , pattern Halt1
  , _MkKappa
  , _ExpPrim
  , _ExpLetRec
  , _ExpApply
  , _AbsLambda'
  , binders
  , body
  , op
  , args
  , cont
  , cps
  , pattern AbsLambda'
  , pattern AbsKappa'
  , Abs(..)
  , Free(..)
  , pattern ValLabel
  , labelName -- don't like that this is part of the api
  )
  where

import Language.SexpGrammar qualified as S
import Gyehoek.Sexp qualified
import Gyehoek.Scheme.Syntax (Name (..), Prim(..), primSexpIso, Lit(..), pattern Void)
import Language.SexpGrammar.Generic
import Control.Category
import Prelude hiding ((.), id)
import Language.Haskell.TH.Quote (QuasiQuoter)
import Language.Sexp.Located (Sexp)
import qualified Data.InvertibleGrammar.Base as IG
import Data.InvertibleGrammar.Base (type (:-)((:-)))
import qualified Data.HashSet as HS
import Data.Monoid (Endo)
import Data.Functor.Foldable.TH
import qualified Gyehoek.Sexp as GS
import qualified Language.Sexp.Located as SL
import Data.Data.Lens (uniplate)
import Gyehoek.Prelude hiding (op)

-- Data types

data Val
  = ValImm Imm
  | ValVar Name
  deriving (Show, Generic, Data, Eq)

pattern ValLabel :: Name -> Val
pattern ValLabel x = ValImm (ImmLabel x)

data Imm
  = ImmInt Int
  | ImmBool Bool
  | ImmLabel Name
  deriving stock (Show, Generic, Data, Eq)
  deriving anyclass (NFData)

data Obj
  = ObjImm Imm
  | ObjHob Hob
  deriving stock (Show, Generic, Data, Eq)
  deriving anyclass (NFData)

-- | a heap object.
data Hob
  = HobClosure { label :: Name, env :: List Obj }
  deriving stock (Show, Generic, Data, Eq)
  deriving anyclass (NFData)

data Kappa = MkKappa { binders :: List Name, body :: Exp }
  deriving (Show, Generic, Data, Eq)

data Lambda = MkLambda { binders :: List Name, ktail :: Name, body :: Exp }
  deriving (Show, Generic, Data, Eq)

data Abs
  = AbsKappa Kappa
  | AbsLambda Lambda
  deriving (Show, Generic, Data, Eq)

pattern AbsKappa' :: [Name] -> Exp -> Abs
pattern AbsKappa' xs e = AbsKappa (MkKappa xs e)

pattern AbsLambda' :: [Name] -> Name -> Exp -> Abs
pattern AbsLambda' xs e ktail = AbsLambda (MkLambda xs e ktail)

data Exp
  = ExpPrim (Prim Val) Kappa
  | ExpLetRec { binders :: List (Name, Abs), body :: Exp }
  | ExpContinue Name (List Val)
  | ExpIf Val Exp Exp
  | ExpApply
    { op :: Val
    , args :: List Val
    , cont :: Name
    }
  deriving (Show, Generic, Data, Eq)

pattern Halt :: List Val -> Exp
pattern Halt xs = ExpContinue "halt" xs

pattern Halt1 :: Val -> Exp
pattern Halt1 x = ExpContinue "halt" [x]

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

data Program = MkProgram
  { body :: Exp
  }
  deriving (Show, Generic, Data)

makePrisms ''Kappa
-- makeLenses ''Kappa
makePrisms ''Exp
-- makeLenses ''Exp
-- makeFieldsNoPrefix ''Exp
-- makeFieldsNoPrefix ''Kappa
-- makeLensesWith abbreviatedFields ''Exp
-- makeLensesFor [("binders", "_binders"), ("body", "_body")] ''Exp
makeFieldsId ''Exp
makeFieldsId ''Kappa
makeFieldsId ''Lambda
makeBaseFunctor ''Exp

instance HasBinders Abs (List Name) where
  binders k (AbsKappa kap) = AbsKappa <$> binders k kap
  binders k (AbsLambda lam) = AbsLambda <$> binders k lam

instance HasBody Abs Exp where
  body k (AbsKappa kap) = AbsKappa <$> body k kap
  body k (AbsLambda lam) = AbsLambda <$> body k lam

_AbsLambda' :: Prism' Abs (List Name, Name, Exp)
_AbsLambda' = prism'
  (\(bs,ktail,e) -> AbsLambda' bs ktail e)
  (\case AbsLambda' bs ktail e -> Just (bs,ktail,e)
         _ -> Nothing)

instance Plated Exp where
  plate = uniplate
  -- plate k = \case
  --   ExpPrim p kap -> ExpPrim p <$> body k kap
  --   ExpLetRec bs e -> ExpLetRec <$> (each . _2 . body) k bs <*> k e
  --   ExpContinue c xs -> pure $ ExpContinue c xs
  --   ExpIf c t f -> ExpIf c <$> k t <*> k f
  --   ExpApply f xs ktail -> pure $ ExpApply f xs ktail


-- SexpIso instances

instance S.SexpIso Val where
  sexpIso = match
    $ With (\imm -> imm . S.sexpIso)
    $ With (\var -> var . S.sexpIso)
    $ End

instance S.SexpIso Obj where
  sexpIso = match
    $ With (\imm -> imm . S.sexpIso)
    $ With (\hob -> hob . S.sexpIso)
    $ End

instance S.SexpIso Imm where
  sexpIso = match
    $ With (. S.int)
    $ With (. GS.schemeBool)
    $ With (. labelName)
    $ End

labelName :: S.SexpGrammar Name
labelName = S.coproduct
  [ S.sexpIso @Name >>> Gyehoek.Sexp.prismIso
        (S.expected "label")
        (prefixed @Name "$")
  , S.list $ S.el (S.sym "$") >>> S.el (S.sexpIso @Name)
  ]

instance S.SexpIso Hob where
  sexpIso = match
    $ With (. closure)
    $ End
    where
      -- closures can be printed, but not parsed.
      closure :: S.Grammar S.Position (Sexp :- t) (List Obj :- Name :- t)
      closure = IG.Flip $ IG.PartialIso
        (\(env:-code:-t) -> SL.Modified SL.Hash [GS.sx|(#{code} ##{env})|] :- t)
        (const . Left $ mempty)

instance S.SexpIso Lambda where
  sexpIso = match
    $ With (. lambda)
    $ End
    where
      lambda = S.list $
        S.el Gyehoek.Sexp.lambdaKeyword
        >>> S.el binders
        >>> S.el S.sexpIso
      binders :: forall t.
        IG.Grammar S.Position (Sexp :- t) (Name :- List Name :- t)
      binders = S.list $
        S.rest (S.sexpIso @Name)
        >>> S.onTail (S.flipped $ IG.PartialIso
                     (\(ktail:-args:-t) -> (args ++ [ktail]) :- t)
                     (\(args:-t) -> case args ^? _Snoc of
                         Just (args',ktail) -> Right $ ktail :- args' :- t
                         Nothing -> Left $ S.expected "cont param")
                     )

instance S.SexpIso Kappa where
  sexpIso = match
    $ With (. kappa)
    $ End
    where
      kappa = S.list $
        S.el Gyehoek.Sexp.kappaKeyword
        >>> S.el (S.list $ S.rest S.sexpIso)
        >>> S.el S.sexpIso

instance S.SexpIso Abs where
  sexpIso = match
    $ With (\lambda -> lambda . S.sexpIso)
    $ With (\kappa -> kappa . S.sexpIso)
    $ End

instance S.SexpIso Exp where
  sexpIso = match
    $ With (. prim)
    $ With (. letrec)
    $ With (. continue)
    $ With (. if_)
    $ With (. app)
    $ End
    where
      continue = S.list $
        S.el (S.sym "continue")
        >>> S.el S.sexpIso
        >>> S.rest S.sexpIso
      letrec = Gyehoek.Sexp.let_ "letrec" S.sexpIso S.sexpIso S.sexpIso
      if_ = S.list $ S.el (S.sym "if")
        >>> S.el S.sexpIso >>> S.el S.sexpIso >>> S.el S.sexpIso
      app :: forall t.
        IG.Grammar S.Position (Sexp :- t) (Name :- ([Val] :- (Val :- t)))
      app = S.list $ S.el (S.sexpIso @Val)
        -- >>> S.flipped Gyehoek.Sexp.nonEmptyGrammar
        >>> S.rest (S.sexpIso @Val)
        -- >>> _
        >>> S.onTail (S.flipped $ IG.PartialIso
                      (\(karg :- args :- op :- t) ->
                         (args ++ [ValVar karg]) :- op :- t)
                      (\(xs :- op :- t) -> case xs ^? _Snoc of
                          Just (args,preview #ValVar -> Just karg) ->
                            Right $ karg:- args :- op :- t
                          _ -> Left $ S.expected "continuation arg"
                      ))
        where
          _ = S.flipped $ Gyehoek.Sexp.nonEmptyGrammar @S.Position @Val
      prim = S.list $
        S.el (S.sym "prim")
        >>> S.el (primSexpIso id (S.sexpIso @Val))
        >>> S.el S.sexpIso

instance S.SexpIso Program where
  sexpIso = with \prog -> S.sexpIso @Exp >>> prog


-- quasiquoters

class Data a => CPS a where
  toCPS :: Sexp -> a

instance CPS Exp where toCPS = Gyehoek.Sexp.fromSexp
instance CPS Val where toCPS = Gyehoek.Sexp.fromSexp
instance CPS Kappa where toCPS = Gyehoek.Sexp.fromSexp
instance CPS Lambda where toCPS = Gyehoek.Sexp.fromSexp
instance CPS Abs where toCPS = Gyehoek.Sexp.fromSexp
instance CPS Program where toCPS = Gyehoek.Sexp.fromSexp

cps :: QuasiQuoter
cps = Gyehoek.Sexp.makeSx' [| toCPS |]



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

toHashSetOf :: Hashable a => Getting (Endo (HashSet a)) s a -> s -> HashSet a
toHashSetOf l = foldrOf l HS.insert mempty

class Free a where
  free :: a -> HashSet Name
  free = freeWithBound mempty

  freeWithBound :: HashSet Name -> a -> HashSet Name
  freeWithBound bound = HS.fromList . freeWithBound' bound

  -- | Free variables given in the order of their appearance.
  free' :: a -> List Name
  free' = freeWithBound' mempty

  freeWithBound' :: HashSet Name -> a -> List Name

instance Free Abs where
  freeWithBound' bound (AbsKappa kap) = freeWithBound' bound kap
  freeWithBound' bound (AbsLambda lam) = freeWithBound' bound lam

instance Free Exp where
  freeWithBound' bound = \case
    ExpPrim p k ->
      p & toListOf (folded . #ValVar . filtered (`notElem` bound))
        & (<> freeWithBound' bound k)
    ExpLetRec bs m ->
      foldMapOf (each . _2) (freeWithBound' bound') bs
      <> freeWithBound' bound' m
      where bound' = bound & insertFrom (bs ^.. each . _1)
    ExpContinue k xs -> filter (`notElem` bound) (k : xs ^.. each . #ValVar)
    ExpIf c t f ->
      (c ^.. #ValVar . filtered (`notElem` bound))
      <> freeWithBound' bound t <> freeWithBound' bound f
    ExpApply f xs k ->
      (f:xs) ^.. (each . #ValVar . filtered (`notElem` bound))
      <> (k ^.. filtered (`notElem` bound))

instance Free Kappa where
  freeWithBound' bound (MkKappa xs m) =
    freeWithBound' (bound & insertFrom xs) m

instance Free Lambda where
  freeWithBound' bound (MkLambda xs k m) =
    freeWithBound' (bound & insertFrom (k:xs)) m



class Vars a where
  -- | Traverse the immediate variables of an expression.
  vars :: Traversal' a Name

instance Vars Val where
  vars k (ValVar x) = ValVar <$> k x
  vars _ x = pure x

instance Vars a => Vars (Prim a) where
  vars k p = traverseOf (each . vars) k p

instance Vars Exp where
  vars k (ExpPrim p kap) = ExpPrim <$> vars k p <*> pure kap
  vars k (ExpContinue kname xs) = ExpContinue <$> k kname <*> pure xs
  vars _ e = pure e