{-# LANGUAGE OverloadedStrings #-}
{-# LANGUAGE TemplateHaskellQuotes #-}

-- This Source Code Form is subject to the terms of the Mozilla Public
-- License, v. 2.0. If a copy of the MPL was not distributed with this
-- file, You can obtain one at https://mozilla.org/MPL/2.0/.

module Data.Effect.Key.TH where

import Control.Effect.Key (SendFOEBy, SendHOEBy)
import Control.Lens ((%~), (<&>), _Just, _head)
import Control.Monad (forM_)
import Control.Monad.Writer (execWriterT, tell)
import Data.Char (toLower)
import Data.Default (def)
import Data.Effect.Key (type (##>), type (#>))
import Data.Effect.TH (makeEffect', noDeriveHFunctor)
import Data.Effect.TH.Internal (
    DataInfo,
    EffClsInfo (EffClsInfo),
    EffConInfo (EffConInfo),
    EffectClassConf (EffectClassConf),
    EffectConf (EffectConf, _keyedSenderGenConf),
    EffectOrder (FirstOrder, HigherOrder),
    MakeEffectConf,
    SenderFunctionConf (SenderFunctionConf),
    alterEffectConf,
    ecEffs,
    ecName,
    ecParamVars,
    effName,
    genSenderArmor,
    normalSenderGenConf,
    senderFnName,
    tyVarName,
    _confByEffect,
    _keyedSenderGenConf,
    _senderFnName,
 )
import Data.Function ((&))
import Data.List.Extra (stripSuffix)
import Data.Text qualified as T
import Formatting (sformat, string, (%))
import Language.Haskell.TH (
    Body (NormalB),
    Clause (Clause),
    Dec (DataD, TySynD),
    Exp (AppTypeE, VarE),
    Info,
    Name,
    Q,
    TyVarBndr (PlainTV),
    Type (AppT, ConT, InfixT, VarT),
    mkName,
    nameBase,
 )
import Language.Haskell.TH.Datatype.TyVarBndr (pattern BndrReq)

makeKeyedEffect :: [Name] -> [Name] -> Q [Dec]
makeKeyedEffect :: [Name] -> [Name] -> Q [Dec]
makeKeyedEffect =
    MakeEffectConf
-> (EffectOrder
    -> Info -> DataInfo -> EffClsInfo -> EffectClassConf -> Q [Dec])
-> [Name]
-> [Name]
-> Q [Dec]
makeEffect'
        (MakeEffectConf
forall a. Default a => a
def MakeEffectConf
-> (MakeEffectConf -> MakeEffectConf) -> MakeEffectConf
forall a b. a -> (a -> b) -> b
& MakeEffectConf -> MakeEffectConf
changeNormalSenderFnNameFormat)
        EffectOrder
-> Info -> DataInfo -> EffClsInfo -> EffectClassConf -> Q [Dec]
genEffectKey
{-# INLINE makeKeyedEffect #-}

makeKeyedEffect_ :: [Name] -> [Name] -> Q [Dec]
makeKeyedEffect_ :: [Name] -> [Name] -> Q [Dec]
makeKeyedEffect_ =
    MakeEffectConf
-> (EffectOrder
    -> Info -> DataInfo -> EffClsInfo -> EffectClassConf -> Q [Dec])
-> [Name]
-> [Name]
-> Q [Dec]
makeEffect'
        (MakeEffectConf
forall a. Default a => a
def MakeEffectConf
-> (MakeEffectConf -> MakeEffectConf) -> MakeEffectConf
forall a b. a -> (a -> b) -> b
& MakeEffectConf -> MakeEffectConf
noDeriveHFunctor MakeEffectConf
-> (MakeEffectConf -> MakeEffectConf) -> MakeEffectConf
forall a b. a -> (a -> b) -> b
& MakeEffectConf -> MakeEffectConf
changeNormalSenderFnNameFormat)
        EffectOrder
-> Info -> DataInfo -> EffClsInfo -> EffectClassConf -> Q [Dec]
genEffectKey
{-# INLINE makeKeyedEffect_ #-}

changeNormalSenderFnNameFormat :: MakeEffectConf -> MakeEffectConf
changeNormalSenderFnNameFormat :: MakeEffectConf -> MakeEffectConf
changeNormalSenderFnNameFormat =
    (EffectConf -> EffectConf) -> MakeEffectConf -> MakeEffectConf
alterEffectConf ((EffectConf -> EffectConf) -> MakeEffectConf -> MakeEffectConf)
-> (EffectConf -> EffectConf) -> MakeEffectConf -> MakeEffectConf
forall a b. (a -> b) -> a -> b
$ (Maybe SenderFunctionConf -> Identity (Maybe SenderFunctionConf))
-> EffectConf -> Identity EffectConf
Lens' EffectConf (Maybe SenderFunctionConf)
normalSenderGenConf ((Maybe SenderFunctionConf -> Identity (Maybe SenderFunctionConf))
 -> EffectConf -> Identity EffectConf)
-> (([Char] -> Identity [Char])
    -> Maybe SenderFunctionConf -> Identity (Maybe SenderFunctionConf))
-> ([Char] -> Identity [Char])
-> EffectConf
-> Identity EffectConf
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (SenderFunctionConf -> Identity SenderFunctionConf)
-> Maybe SenderFunctionConf -> Identity (Maybe SenderFunctionConf)
forall a b (p :: * -> * -> *) (f :: * -> *).
(Choice p, Applicative f) =>
p a (f b) -> p (Maybe a) (f (Maybe b))
_Just ((SenderFunctionConf -> Identity SenderFunctionConf)
 -> Maybe SenderFunctionConf -> Identity (Maybe SenderFunctionConf))
-> (([Char] -> Identity [Char])
    -> SenderFunctionConf -> Identity SenderFunctionConf)
-> ([Char] -> Identity [Char])
-> Maybe SenderFunctionConf
-> Identity (Maybe SenderFunctionConf)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ([Char] -> Identity [Char])
-> SenderFunctionConf -> Identity SenderFunctionConf
Lens' SenderFunctionConf [Char]
senderFnName (([Char] -> Identity [Char]) -> EffectConf -> Identity EffectConf)
-> ([Char] -> [Char]) -> EffectConf -> EffectConf
forall s t a b. ASetter s t a b -> (a -> b) -> s -> t
%~ ([Char] -> [Char] -> [Char]
forall a. [a] -> [a] -> [a]
++ [Char]
"'_")
{-# INLINE changeNormalSenderFnNameFormat #-}

genEffectKey :: EffectOrder -> Info -> DataInfo -> EffClsInfo -> EffectClassConf -> Q [Dec]
genEffectKey :: EffectOrder
-> Info -> DataInfo -> EffClsInfo -> EffectClassConf -> Q [Dec]
genEffectKey EffectOrder
order Info
_ DataInfo
_ EffClsInfo{[TyVarBndr ()]
[EffConInfo]
Name
ecEffs :: EffClsInfo -> [EffConInfo]
ecName :: EffClsInfo -> Name
ecParamVars :: EffClsInfo -> [TyVarBndr ()]
ecName :: Name
ecParamVars :: [TyVarBndr ()]
ecEffs :: [EffConInfo]
..} EffectClassConf{Name -> EffectConf
_confByEffect :: EffectClassConf -> Name -> EffectConf
_confByEffect :: Name -> EffectConf
..} = WriterT [Dec] Q () -> Q [Dec]
forall (m :: * -> *) w a. Monad m => WriterT w m a -> m w
execWriterT do
    let keyedOp :: Name
keyedOp = case EffectOrder
order of
            EffectOrder
FirstOrder -> ''(#>)
            EffectOrder
HigherOrder -> ''(##>)

        pvs :: [Name]
pvs = TyVarBndr () -> Name
forall a. TyVarBndr a -> Name
tyVarName (TyVarBndr () -> Name) -> [TyVarBndr ()] -> [Name]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> [TyVarBndr ()]
ecParamVars

    [Char]
ecNamePlain <-
        [Char] -> Maybe [Char]
removeLastApostrophe (Name -> [Char]
nameBase Name
ecName)
            Maybe [Char]
-> (Maybe [Char] -> WriterT [Dec] Q [Char])
-> WriterT [Dec] Q [Char]
forall a b. a -> (a -> b) -> b
& WriterT [Dec] Q [Char]
-> ([Char] -> WriterT [Dec] Q [Char])
-> Maybe [Char]
-> WriterT [Dec] Q [Char]
forall b a. b -> (a -> b) -> Maybe a -> b
maybe
                ( [Char] -> WriterT [Dec] Q [Char]
forall a. [Char] -> WriterT [Dec] Q a
forall (m :: * -> *) a. MonadFail m => [Char] -> m a
fail ([Char] -> WriterT [Dec] Q [Char])
-> (Text -> [Char]) -> Text -> WriterT [Dec] Q [Char]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Text -> [Char]
T.unpack (Text -> WriterT [Dec] Q [Char]) -> Text -> WriterT [Dec] Q [Char]
forall a b. (a -> b) -> a -> b
$
                    Format Text ([Char] -> Text) -> [Char] -> Text
forall a. Format Text a -> a
sformat
                        (Format ([Char] -> Text) ([Char] -> Text)
"No last apostrophe on the effect class ‘" Format ([Char] -> Text) ([Char] -> Text)
-> Format Text ([Char] -> Text) -> Format Text ([Char] -> Text)
forall r a r'. Format r a -> Format r' r -> Format r' a
% Format Text ([Char] -> Text)
forall r. Format r ([Char] -> r)
string Format Text ([Char] -> Text)
-> Format Text Text -> Format Text ([Char] -> Text)
forall r a r'. Format r a -> Format r' r -> Format r' a
% Format Text Text
"’.")
                        (Name -> [Char]
nameBase Name
ecName)
                )
                [Char] -> WriterT [Dec] Q [Char]
forall a. a -> WriterT [Dec] Q a
forall (f :: * -> *) a. Applicative f => a -> f a
pure

    let keyDataName :: Name
keyDataName = [Char] -> Name
mkName ([Char] -> Name) -> [Char] -> Name
forall a b. (a -> b) -> a -> b
$ [Char]
ecNamePlain [Char] -> [Char] -> [Char]
forall a. [a] -> [a] -> [a]
++ [Char]
"Key"
        key :: Type
key = Name -> Type
ConT Name
keyDataName

    [Dec] -> WriterT [Dec] Q ()
forall w (m :: * -> *). MonadWriter w m => w -> m ()
tell [Cxt
-> Name
-> [TyVarBndr BndrVis]
-> Maybe Type
-> [Con]
-> [DerivClause]
-> Dec
DataD [] Name
keyDataName [] Maybe Type
forall a. Maybe a
Nothing [] []]

    [Dec] -> WriterT [Dec] Q ()
forall w (m :: * -> *). MonadWriter w m => w -> m ()
tell
        [ Name -> [TyVarBndr BndrVis] -> Type -> Dec
TySynD
            ([Char] -> Name
mkName [Char]
ecNamePlain)
            ([Name]
pvs [Name] -> (Name -> TyVarBndr BndrVis) -> [TyVarBndr BndrVis]
forall (f :: * -> *) a b. Functor f => f a -> (a -> b) -> f b
<&> (Name -> BndrVis -> TyVarBndr BndrVis
forall flag. Name -> flag -> TyVarBndr flag
`PlainTV` BndrVis
BndrReq))
            (Type -> Name -> Type -> Type
InfixT Type
key Name
keyedOp ((Type -> Type -> Type) -> Type -> Cxt -> Type
forall b a. (b -> a -> b) -> b -> [a] -> b
forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
foldl Type -> Type -> Type
AppT (Name -> Type
ConT Name
ecName) ((Name -> Type) -> [Name] -> Cxt
forall a b. (a -> b) -> [a] -> [b]
map Name -> Type
VarT [Name]
pvs)))
        ]

    [EffConInfo]
-> (EffConInfo -> WriterT [Dec] Q ()) -> WriterT [Dec] Q ()
forall (t :: * -> *) (m :: * -> *) a b.
(Foldable t, Monad m) =>
t a -> (a -> m b) -> m ()
forM_ [EffConInfo]
ecEffs \con :: EffConInfo
con@EffConInfo{Name
effName :: EffConInfo -> Name
effName :: Name
..} -> do
        let EffectConf{Maybe SenderFunctionConf
_keyedSenderGenConf :: EffectConf -> Maybe SenderFunctionConf
_keyedSenderGenConf :: Maybe SenderFunctionConf
..} = Name -> EffectConf
_confByEffect Name
effName
        Maybe SenderFunctionConf
-> (SenderFunctionConf -> WriterT [Dec] Q ()) -> WriterT [Dec] Q ()
forall (t :: * -> *) (m :: * -> *) a b.
(Foldable t, Monad m) =>
t a -> (a -> m b) -> m ()
forM_ Maybe SenderFunctionConf
_keyedSenderGenConf \conf :: SenderFunctionConf
conf@SenderFunctionConf{[Char]
_senderFnName :: SenderFunctionConf -> [Char]
_senderFnName :: [Char]
..} -> do
            let sendCxt :: Type -> Type -> Type
sendCxt Type
effDataType Type
carrier = case EffectOrder
order of
                    EffectOrder
FirstOrder -> Name -> Type
ConT ''SendFOEBy Type -> Type -> Type
`AppT` Type
key Type -> Type -> Type
`AppT` Type
effDataType Type -> Type -> Type
`AppT` Type
carrier
                    EffectOrder
HigherOrder -> Name -> Type
ConT ''SendHOEBy Type -> Type -> Type
`AppT` Type
key Type -> Type -> Type
`AppT` Type
effDataType Type -> Type -> Type
`AppT` Type
carrier

            (Type -> Type -> Type)
-> ([TyVarBndrSpec] -> [TyVarBndrSpec])
-> SenderFunctionConf
-> EffConInfo
-> (Type -> Q Clause)
-> WriterT [Dec] Q ()
genSenderArmor Type -> Type -> Type
sendCxt [TyVarBndrSpec] -> [TyVarBndrSpec]
forall a. a -> a
id SenderFunctionConf
conf{_senderFnName = nameBase effName & _head %~ toLower} EffConInfo
con \Type
_f ->
                Clause -> Q Clause
forall a. a -> Q a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Clause -> Q Clause) -> Clause -> Q Clause
forall a b. (a -> b) -> a -> b
$ [Pat] -> Body -> [Dec] -> Clause
Clause [] (Exp -> Body
NormalB (Exp -> Body) -> Exp -> Body
forall a b. (a -> b) -> a -> b
$ Name -> Exp
VarE ([Char] -> Name
mkName [Char]
_senderFnName) Exp -> Type -> Exp
`AppTypeE` Type
key) []

removeLastApostrophe :: String -> Maybe String
removeLastApostrophe :: [Char] -> Maybe [Char]
removeLastApostrophe = [Char] -> [Char] -> Maybe [Char]
forall a. Eq a => [a] -> [a] -> Maybe [a]
stripSuffix [Char]
"'"