{-# 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 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 Gyehoek.Scheme.Syntax (Name (..), Prim(..), primDatumIso, Lit(..))
import Gyehoek.Sexp qualified as S
import Control.Category
import Prelude hiding ((.), id)
import qualified Data.HashSet as HS
import Data.Monoid (Endo)
import Data.Functor.Foldable.TH
import Data.Data.Lens (uniplate)
import Gyehoek.Prelude hiding (op)
import Gyehoek.Sexp (Datum)
import Gyehoek.Sexp (G, (:-)(..))
import qualified Data.InvertibleGrammar.Base as IG

-- 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 Val (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 (ValLabel "halt") xs

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

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

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

makePrisms ''Kappa
makePrisms ''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


-- DatumIso instances

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

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

instance S.DatumIso Imm where
  datumIso = S.match
    $ S.With (. S.int)
    $ S.With (. S.datumIso)
    $ S.With (. labelName)
    $ S.End

labelName :: S.DatumGrammar Name
labelName = S.coproduct
  [ S.decorate S.SynConstant >>> S.datumIso @Name >>> S.prismIso
        (S.expected "label")
        (prefixed @Name "$")
  , S.list $ S.el (S.sym "$") >>> S.el (S.datumIso @Name)
  ]

instance S.DatumIso Hob where
  datumIso = S.match
    $ S.With (. closure)
    $ S.End
    where
      -- closures can be printed, but not parsed.
      closure :: G (Datum :- t) (List Obj :- Name :- t)
      closure = IG.Flip $ IG.PartialIso
        (\(env:-code:-t) -> [S.sx|(<closure> #{code} ##{env})|] :- t)
        (const . Left $ mempty)

instance S.DatumIso Lambda where
  datumIso = S.with (lam >>>)
    where
      lam :: forall t. G (Datum :- t) (Exp :- Name :- List Name :- t)
      lam = S.lambdaLike
        S.lambdaKeyword
        binders
        (S.el $ S.datumIso @Exp)
      binders :: forall t. G (Datum :- t) (Name :- List Name :- t)
      binders =
        S.list (S.rest $ S.datumIso @Name)
        >>> S.flipped S.snoced
        >>> S.swap

instance S.DatumIso Kappa where
  datumIso = S.with \g ->
    S.lambdaLike S.kappaKeyword
      (S.datumIso @(List Name))
      (S.el $ S.datumIso @Exp)
    >>> g

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

instance S.DatumIso Exp where
  datumIso = S.match
    $ S.With (. prim)
    $ S.With (. letrec)
    $ S.With (. continue)
    $ S.With (. if_)
    $ S.With (. app)
    $ S.End
    where
      continue = S.list $
        S.el (S.decorate S.SynBuiltin >>> S.sym "continue")
        >>> S.el (S.decorate S.SynProcedure >>> S.datumIso)
        >>> S.rest S.datumIso
      letrec = S.letLike "letrec" S.datumIso S.datumIso S.datumIso
      if_ = S.ifLike "if"
        S.datumIso S.datumIso S.datumIso
      app :: forall t.
        G (Datum :- t) (Name :- ([Val] :- (Val :- t)))
      app = S.list $ S.el (S.datumIso @Val)
        -- >>> S.flipped Gyehoek.Datum.nonEmptyGrammar
        >>> S.rest (S.datumIso @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"
                      ))
      -- prim = S.headTagged2 "prim"
      --   (primDatumIso id (S.datumIso @Val))
      --   (S.datumIso @Kappa)
      prim = S.list $
        S.el (S.decorate S.SynBuiltin >>> S.sym "prim")
        >>> S.el (primDatumIso id (S.datumIso @Val))
        >>> S.el S.datumIso

instance S.DatumIso Program where
  datumIso = S.with \prog -> S.datumIso @Lambda >>> prog


-- quasiquoters

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

instance CPS Exp where toCPS = S.fromDatumUnsafe S.datumIso
instance CPS Val where toCPS = S.fromDatumUnsafe S.datumIso
instance CPS Kappa where toCPS = S.fromDatumUnsafe S.datumIso
instance CPS Lambda where toCPS = S.fromDatumUnsafe S.datumIso
instance CPS Abs where toCPS = S.fromDatumUnsafe S.datumIso
instance CPS Program where toCPS = S.fromDatumUnsafe S.datumIso

cps :: S.QuasiQuoter
cps = S.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