{-# LANGUAGE AllowAmbiguousTypes #-}
{-# LANGUAGE QuantifiedConstraints #-}
{-# LANGUAGE UndecidableInstances #-}

-- 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 Control.Hefty where

import Control.Applicative (Alternative, empty, (<|>))
import Control.Effect (SendIns (..), SendSig (..), type (~>))
import Control.Effect.Key (ByKey (ByKey), SendInsBy, SendSigBy, key, sendInsBy, sendSigBy)
import Control.Freer (Freer (liftIns), InjectIns, InjectInsBy, StateKey, injectIns, injectInsBy)
import Control.Monad (MonadPlus)
import Control.Monad.Base (MonadBase)
import Control.Monad.Fix (MonadFix, mfix)
import Control.Monad.IO.Class (MonadIO, liftIO)
import Control.Monad.RWS.Class (MonadRWS)
import Control.Monad.Reader.Class (MonadReader, ask, local)
import Control.Monad.State.Class (MonadState, get, put)
import Control.Monad.Writer.Class (MonadWriter, listen, pass, tell)
import Data.Effect (InsClass, SigClass)
import Data.Effect.Fail (Fail)
import Data.Effect.Fail qualified as E
import Data.Effect.Fix (Fix)
import Data.Effect.Fix qualified as E
import Data.Effect.NonDet (ChooseH, Empty, chooseH)
import Data.Effect.NonDet qualified as NonDet
import Data.Effect.Reader (Ask, Local, ask'', local'')
import Data.Effect.State (State, get'', put'')
import Data.Effect.Unlift (UnliftIO, pattern WithRunInIO)
import Data.Effect.Writer (Tell, WriterH, listen'', tell'')
import Data.Function ((&))
import Data.Kind (Type)
import Data.Tuple (swap)
import UnliftIO (MonadUnliftIO, withRunInIO)

newtype
    Hefty
        (f :: InsClass -> Type -> Type)
        (e :: SigClass)
        (a :: Type) = Hefty
    {forall (f :: InsClass -> InsClass) (e :: InsClass -> InsClass) a.
Hefty f e a -> f (e (Hefty f e)) a
unHefty :: f (e (Hefty f e)) a}

deriving newtype instance Functor (f (e (Hefty f e))) => Functor (Hefty f e)
deriving newtype instance Applicative (f (e (Hefty f e))) => Applicative (Hefty f e)
deriving newtype instance Monad (f (e (Hefty f e))) => Monad (Hefty f e)
deriving newtype instance (MonadBase b (f (e (Hefty f e))), Monad b) => MonadBase b (Hefty f e)

deriving newtype instance Foldable (f (e (Hefty f e))) => Foldable (Hefty f e)
deriving stock instance Traversable (f (e (Hefty f e))) => Traversable (Hefty f e)
deriving newtype instance Eq (f (e (Hefty f e)) a) => Eq (Hefty f e a)
deriving newtype instance Ord (f (e (Hefty f e)) a) => Ord (Hefty f e a)
deriving newtype instance Read (f (e (Hefty f e)) a) => Read (Hefty f e a)
deriving newtype instance Show (f (e (Hefty f e)) a) => Show (Hefty f e a)

overHefty ::
    (f (e (Hefty f e)) a -> f' (e' (Hefty f' e')) b) ->
    Hefty f e a ->
    Hefty f' e' b
overHefty :: forall (f :: InsClass -> InsClass) (e :: InsClass -> InsClass) a
       (f' :: InsClass -> InsClass) (e' :: InsClass -> InsClass) b.
(f (e (Hefty f e)) a -> f' (e' (Hefty f' e')) b)
-> Hefty f e a -> Hefty f' e' b
overHefty f (e (Hefty f e)) a -> f' (e' (Hefty f' e')) b
f = forall (f :: InsClass -> InsClass) (e :: InsClass -> InsClass) a.
f (e (Hefty f e)) a -> Hefty f e a
Hefty forall b c a. (b -> c) -> (a -> b) -> a -> c
. f (e (Hefty f e)) a -> f' (e' (Hefty f' e')) b
f forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall (f :: InsClass -> InsClass) (e :: InsClass -> InsClass) a.
Hefty f e a -> f (e (Hefty f e)) a
unHefty
{-# INLINE overHefty #-}

instance (Freer c fr, InjectIns e (e' (Hefty fr e'))) => SendIns e (Hefty fr e') where
    sendIns :: forall a. e a -> Hefty fr e' a
sendIns = forall (f :: InsClass -> InsClass) (e :: InsClass -> InsClass) a.
f (e (Hefty f e)) a -> Hefty f e a
Hefty forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall (c :: InsClass -> Constraint) (f :: InsClass -> InsClass)
       (e :: InsClass) a.
Freer c f =>
e a -> f e a
liftIns forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall (e :: InsClass) (e' :: InsClass). InjectIns e e' => e ~> e'
injectIns
    {-# INLINE sendIns #-}

instance (Freer c fr, InjectSig e e') => SendSig e (Hefty fr e') where
    sendSig :: forall a. e (Hefty fr e') a -> Hefty fr e' a
sendSig = forall (f :: InsClass -> InsClass) (e :: InsClass -> InsClass) a.
f (e (Hefty f e)) a -> Hefty f e a
Hefty forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall (c :: InsClass -> Constraint) (f :: InsClass -> InsClass)
       (e :: InsClass) a.
Freer c f =>
e a -> f e a
liftIns forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall (e :: InsClass -> InsClass) (e' :: InsClass -> InsClass)
       (f :: InsClass).
InjectSig e e' =>
e f ~> e' f
injectSig
    {-# INLINE sendSig #-}

class InjectSig e (e' :: SigClass) where
    injectSig :: e f ~> e' f

instance (Freer c fr, InjectInsBy key e (e' (Hefty fr e'))) => SendInsBy key e (Hefty fr e') where
    sendInsBy :: forall a. e a -> Hefty fr e' a
sendInsBy = forall (f :: InsClass -> InsClass) (e :: InsClass -> InsClass) a.
f (e (Hefty f e)) a -> Hefty f e a
Hefty forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall (c :: InsClass -> Constraint) (f :: InsClass -> InsClass)
       (e :: InsClass) a.
Freer c f =>
e a -> f e a
liftIns forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall {k} (key :: k) (e :: InsClass) (e' :: InsClass).
InjectInsBy key e e' =>
e ~> e'
injectInsBy @key
    {-# INLINE sendInsBy #-}

instance (Freer c fr, InjectSigBy key e e') => SendSigBy key e (Hefty fr e') where
    sendSigBy :: forall a. e (Hefty fr e') a -> Hefty fr e' a
sendSigBy = forall (f :: InsClass -> InsClass) (e :: InsClass -> InsClass) a.
f (e (Hefty f e)) a -> Hefty f e a
Hefty forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall (c :: InsClass -> Constraint) (f :: InsClass -> InsClass)
       (e :: InsClass) a.
Freer c f =>
e a -> f e a
liftIns forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall {k} (key :: k) (e :: InsClass -> InsClass)
       (e' :: InsClass -> InsClass) (f :: InsClass).
InjectSigBy key e e' =>
e f ~> e' f
injectSigBy @key
    {-# INLINE sendSigBy #-}

class InjectSigBy key e (e' :: SigClass) | key e' -> e where
    injectSigBy :: e f ~> e' f

instance
    ( Freer c fr
    , InjectInsBy ReaderKey (Ask r) (e (Hefty fr e))
    , InjectSigBy ReaderKey (Local r) e
    , Monad (fr (e (Hefty fr e)))
    ) =>
    MonadReader r (Hefty fr e)
    where
    ask :: Hefty fr e r
ask = forall {k} (key :: k) r (f :: InsClass).
SendInsBy key (Ask r) f =>
f r
ask'' @ReaderKey
    local :: forall a. (r -> r) -> Hefty fr e a -> Hefty fr e a
local = forall {k} (key :: k) r a (f :: InsClass).
SendSigBy key (Local r) f =>
(r -> r) -> f a -> f a
local'' @ReaderKey
    {-# INLINE ask #-}
    {-# INLINE local #-}

data ReaderKey

instance
    ( Freer c fr
    , InjectInsBy WriterKey (Tell w) (e (Hefty fr e))
    , InjectSigBy WriterKey (WriterH w) e
    , Monoid w
    , Monad (fr (e (Hefty fr e)))
    ) =>
    MonadWriter w (Hefty fr e)
    where
    tell :: w -> Hefty fr e ()
tell = forall {k} (key :: k) w (f :: InsClass).
SendInsBy key (Tell w) f =>
w -> f ()
tell'' @WriterKey
    listen :: forall a. Hefty fr e a -> Hefty fr e (a, w)
listen = forall (f :: InsClass) a b. Functor f => (a -> b) -> f a -> f b
fmap forall a b. (a, b) -> (b, a)
swap forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall {k} (key :: k) a w (f :: InsClass).
SendSigBy key (WriterH w) f =>
f a -> f (w, a)
listen'' @WriterKey
    pass :: forall a. Hefty fr e (a, w -> w) -> Hefty fr e a
pass Hefty fr e (a, w -> w)
m = forall w (m :: InsClass) a. MonadWriter w m => m (a, w -> w) -> m a
pass (forall {k} (key :: k) (f :: InsClass) a. f a -> ByKey key f a
ByKey Hefty fr e (a, w -> w)
m) forall a b. a -> (a -> b) -> b
& forall {k} (key :: k) (f :: InsClass) a. ByKey key f a -> f a
key @WriterKey
    {-# INLINE tell #-}
    {-# INLINE listen #-}

data WriterKey

instance
    (Freer c fr, InjectInsBy StateKey (State s) (e (Hefty fr e)), Monad (fr (e (Hefty fr e)))) =>
    MonadState s (Hefty fr e)
    where
    get :: Hefty fr e s
get = forall {k} (key :: k) s (f :: InsClass).
SendInsBy key (State s) f =>
f s
get'' @StateKey
    put :: s -> Hefty fr e ()
put = forall {k} (key :: k) s (f :: InsClass).
SendInsBy key (State s) f =>
s -> f ()
put'' @StateKey
    {-# INLINE get #-}
    {-# INLINE put #-}

instance
    ( Freer c fr
    , InjectInsBy ReaderKey (Ask r) (e (Hefty fr e))
    , InjectSigBy ReaderKey (Local r) e
    , InjectInsBy WriterKey (Tell w) (e (Hefty fr e))
    , InjectSigBy WriterKey (WriterH w) e
    , InjectInsBy StateKey (State s) (e (Hefty fr e))
    , Monoid w
    , Monad (fr (e (Hefty fr e)))
    ) =>
    MonadRWS r w s (Hefty fr e)

instance
    ( Freer c fr
    , InjectIns Empty (e (Hefty fr e))
    , InjectSig ChooseH e
    , Applicative (fr (e (Hefty fr e)))
    ) =>
    Alternative (Hefty fr e)
    where
    empty :: forall a. Hefty fr e a
empty = forall a (f :: InsClass). SendIns Empty f => f a
NonDet.empty
    Hefty fr e a
a <|> :: forall a. Hefty fr e a -> Hefty fr e a -> Hefty fr e a
<|> Hefty fr e a
b = forall a (f :: InsClass). SendSig ChooseH f => f a -> f a -> f a
chooseH Hefty fr e a
a Hefty fr e a
b
    {-# INLINE empty #-}
    {-# INLINE (<|>) #-}

instance
    ( Freer c fr
    , InjectIns Empty (e (Hefty fr e))
    , InjectSig ChooseH e
    , Monad (fr (e (Hefty fr e)))
    ) =>
    MonadPlus (Hefty fr e)

instance (Freer c fr, InjectIns IO (e (Hefty fr e)), Monad (fr (e (Hefty fr e)))) => MonadIO (Hefty fr e) where
    liftIO :: forall a. IO a -> Hefty fr e a
liftIO = forall (ins :: InsClass) (f :: InsClass) a.
SendIns ins f =>
ins a -> f a
sendIns
    {-# INLINE liftIO #-}

instance (Freer c fr, InjectIns Fail (e (Hefty fr e)), Monad (fr (e (Hefty fr e)))) => MonadFail (Hefty fr e) where
    fail :: forall a. String -> Hefty fr e a
fail = forall a (f :: InsClass). SendIns Fail f => String -> f a
E.fail
    {-# INLINE fail #-}

instance (Freer c fr, InjectSig Fix e, Monad (fr (e (Hefty fr e)))) => MonadFix (Hefty fr e) where
    mfix :: forall a. (a -> Hefty fr e a) -> Hefty fr e a
mfix = forall a (f :: InsClass). SendSig Fix f => (a -> f a) -> f a
E.mfix
    {-# INLINE mfix #-}

instance
    ( Freer c fr
    , InjectIns IO (e (Hefty fr e))
    , InjectSig UnliftIO e
    , Monad (fr (e (Hefty fr e)))
    ) =>
    MonadUnliftIO (Hefty fr e)
    where
    withRunInIO :: forall b.
((forall a. Hefty fr e a -> IO a) -> IO b) -> Hefty fr e b
withRunInIO (forall a. Hefty fr e a -> IO a) -> IO b
f = forall (f :: InsClass -> InsClass) (e :: InsClass -> InsClass) a.
f (e (Hefty f e)) a -> Hefty f e a
Hefty forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall (c :: InsClass -> Constraint) (f :: InsClass -> InsClass)
       (e :: InsClass) a.
Freer c f =>
e a -> f e a
liftIns forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall (e :: InsClass -> InsClass) (e' :: InsClass -> InsClass)
       (f :: InsClass).
InjectSig e e' =>
e f ~> e' f
injectSig forall a b. (a -> b) -> a -> b
$ forall (f :: InsClass) a. ((f ~> IO) -> IO a) -> UnliftIO f a
WithRunInIO (forall a. Hefty fr e a -> IO a) -> IO b
f
    {-# INLINE withRunInIO #-}