module Control.Monad.Writer.Interface

import Control.Monad.Maybe
import Control.Monad.Error.Either
import Control.Monad.Reader.Reader
import Control.Monad.State.State
import Control.Monad.RWS.CPS as RWS
import Control.Monad.Trans
import Control.Monad.Writer.CPS as Writer

||| MonadWriter interface
|||
||| tell is like tell on the MUD's it shouts to monad
||| what you want to be heard. The monad carries this 'packet'
||| upwards, merging it if needed (hence the Monoid requirement).
|||
||| listen listens to a monad acting, and returns what the monad "said".
|||
||| pass lets you provide a writer transformer which changes internals of
||| the written object.
public export
interface (Monoid w, Monad m) => MonadWriter w m | m where
  ||| `writer (a,w)` embeds a simple writer action.
  writer : (a,w) -> m a
  writer (a, w) = tell w $> a

  ||| `tell w` is an action that produces the output `w`.
  tell : w -> m ()
  tell w = writer ((),w)

  ||| `listen m` is an action that executes the action `m` and adds
  ||| its output to the value of the computation.
  listen : m a -> m (a, w)

  ||| `pass m` is an action that executes the action `m`, which
  ||| returns a value and a function, and returns the value, applying
  ||| the function to the output.
  pass : m (a, w -> w) -> m a

||| `listens f m` is an action that executes the action `m` and adds
||| the result of applying @f@ to the output to the value of the computation.
public export
listens : MonadWriter w m => (w -> b) -> m a -> m (a, b)
listens f m = do (a, w) <- listen m
                 pure (a, f w)

||| `censor f m` is an action that executes the action `m` and
||| applies the function `f` to its output, leaving the return value
||| unchanged.
public export
censor : MonadWriter w m => (w -> w) -> m a -> m a
censor f m = pass $ do a <- m
                       pure (a, f)

public export %inline
(Monoid w, Monad m) => MonadWriter w (WriterT w m) where
  writer = Writer.writer
  tell   = Writer.tell
  listen = Writer.listen
  pass   = Writer.pass

public export %inline
(Monoid w, Monad m) => MonadWriter w (RWST r w s m) where
  writer = RWS.writer
  tell   = RWS.tell
  listen = RWS.listen
  pass   = RWS.pass

public export %inline
MonadWriter w m => MonadWriter w (EitherT e m) where
  writer = lift . writer
  tell   = lift . tell
  listen = mapEitherT \m => do (e,w) <- listen m
                               pure $ map (\a => (a,w)) e

  pass   = mapEitherT \m => pass $ do Right (r,f) <- m
                                        | Left l => pure $ (Left l, id)
                                      pure (Right r, f)

public export %inline
MonadWriter w m => MonadWriter w (MaybeT m) where
  writer = lift . writer
  tell   = lift . tell
  listen = mapMaybeT \m => do (e,w) <- listen m
                              pure $ map (\a => (a,w)) e

  pass   = mapMaybeT \m => pass $ do Just (r,f) <- m
                                       | Nothing => pure $ (Nothing, id)
                                     pure (Just r, f)
public export %inline
MonadWriter w m => MonadWriter w (ReaderT r m) where
  writer = lift . writer
  tell   = lift . tell
  listen = mapReaderT listen
  pass   = mapReaderT pass

public export %inline
MonadWriter w m => MonadWriter w (StateT s m) where
  writer = lift . writer
  tell   = lift . tell
  listen (ST m) = ST \s => do ((s',a),w) <- listen (m s)
                              pure (s',(a,w))

  pass   (ST m) = ST \s => pass $ do (s',(a,f)) <- m s
                                     pure ((s',a),f)