{-# 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 (SendInsBy, SendSigBy)
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')
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,
 )

makeKeyedEffect :: [Name] -> [Name] -> Q [Dec]
makeKeyedEffect :: [Name] -> [Name] -> Q [Dec]
makeKeyedEffect =
    MakeEffectConf
-> (EffectOrder
    -> Info -> DataInfo -> EffClsInfo -> EffectClassConf -> Q [Dec])
-> [Name]
-> [Name]
-> Q [Dec]
makeEffect'
        (forall a. Default a => a
def 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 forall a b. (a -> b) -> a -> b
$ Lens' EffectConf (Maybe SenderFunctionConf)
normalSenderGenConf forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a b. Prism (Maybe a) (Maybe b) a b
_Just forall b c a. (b -> c) -> (a -> b) -> a -> c
. Lens' SenderFunctionConf [Char]
senderFnName forall s t a b. ASetter s t a b -> (a -> b) -> s -> t
%~ (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 :: [EffConInfo]
ecParamVars :: [TyVarBndr ()]
ecName :: Name
ecParamVars :: EffClsInfo -> [TyVarBndr ()]
ecName :: EffClsInfo -> Name
ecEffs :: EffClsInfo -> [EffConInfo]
..} EffectClassConf{Name -> EffectConf
_confByEffect :: Name -> EffectConf
_confByEffect :: EffectClassConf -> Name -> EffectConf
..} = 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 = forall a. TyVarBndr a -> Name
tyVarName 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)
            forall a b. a -> (a -> b) -> b
& forall b a. b -> (a -> b) -> Maybe a -> b
maybe
                ( forall (m :: * -> *) a. MonadFail m => [Char] -> m a
fail forall b c a. (b -> c) -> (a -> b) -> a -> c
. Text -> [Char]
T.unpack forall a b. (a -> b) -> a -> b
$
                    forall a. Format Text a -> a
sformat
                        (Format ([Char] -> Text) ([Char] -> Text)
"No last apostrophe on the effect class ‘" forall r a r'. Format r a -> Format r' r -> Format r' a
% forall r. Format r ([Char] -> r)
string forall r a r'. Format r a -> Format r' r -> Format r' a
% Format Text Text
"’.")
                        (Name -> [Char]
nameBase Name
ecName)
                )
                forall (f :: * -> *) a. Applicative f => a -> f a
pure

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

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

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

    forall (t :: * -> *) (m :: * -> *) a b.
(Foldable t, Monad m) =>
t a -> (a -> m b) -> m ()
forM_ [EffConInfo]
ecEffs \con :: EffConInfo
con@EffConInfo{Name
effName :: Name
effName :: EffConInfo -> Name
..} -> do
        let EffectConf{Maybe SenderFunctionConf
_keyedSenderGenConf :: Maybe SenderFunctionConf
_keyedSenderGenConf :: EffectConf -> Maybe SenderFunctionConf
..} = Name -> EffectConf
_confByEffect Name
effName
        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 :: [Char]
_senderFnName :: SenderFunctionConf -> [Char]
..} -> do
            let sendCxt :: Type -> Type -> Type
sendCxt Type
effDataType Type
carrier = case EffectOrder
order of
                    EffectOrder
FirstOrder -> Name -> Type
ConT ''SendInsBy Type -> Type -> Type
`AppT` Type
key Type -> Type -> Type
`AppT` Type
effDataType Type -> Type -> Type
`AppT` Type
carrier
                    EffectOrder
HigherOrder -> Name -> Type
ConT ''SendSigBy 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 forall a. a -> a
id SenderFunctionConf
conf{_senderFnName :: [Char]
_senderFnName = Name -> [Char]
nameBase Name
effName forall a b. a -> (a -> b) -> b
& forall s a. Cons s s a a => Traversal' s a
_head forall s t a b. ASetter s t a b -> (a -> b) -> s -> t
%~ Char -> Char
toLower} EffConInfo
con \Type
_f ->
                forall (f :: * -> *) a. Applicative f => a -> f a
pure forall a b. (a -> b) -> a -> b
$ [Pat] -> Body -> [Dec] -> Clause
Clause [] (Exp -> Body
NormalB 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 = forall a. Eq a => [a] -> [a] -> Maybe [a]
stripSuffix [Char]
"'"