{-# LANGUAGE FlexibleInstances, MultiParamTypeClasses, FunctionalDependencies #-}

-- | Provides a typeclass for all monads that may print text.
module Text.Chatty.Printer where

import Control.Arrow
import Control.Monad
import Control.Monad.State
import Control.Monad.Identity
import System.IO

-- | A typeclass for all monads that may output strings.
class Monad m => MonadPrinter m where
  -- | Just print it!
  mprint :: String -> m ()
  -- | Print it, except you are IO.
  mnoecho :: String -> m ()
  mnoecho = mprint
  -- | Flush the buffer.
  mflush :: m ()
  mflush = return ()

instance MonadPrinter IO where
  mprint = putStr
  mnoecho _ = return ()
  mflush = hFlush stdout

instance Monad m => MonadPrinter (StateT String m) where
  mprint s = modify (++s)

-- | DeafT discards all output (much like >\/dev\/null in the shell)
newtype DeafT m a = Deaf { runDeafT :: m a }

instance Monad m => Monad (DeafT m) where
  return = Deaf . return
  (Deaf d) >>= f = Deaf $ do d' <- d; runDeafT (f d')

instance MonadTrans DeafT where
  lift = Deaf

instance Functor m => Functor (DeafT m) where
  fmap f (Deaf a) = Deaf $ fmap f a

instance MonadIO m => MonadIO (DeafT m) where
  liftIO = lift . liftIO

instance Monad m => MonadPrinter (DeafT m) where
  mprint _ = return ()

-- Definition of OutRedirT + instances
-- | Redirects all output to a given handle (much like >filename in the shell)
newtype OutRedirT m a = OutRedir { runOutRedirT' :: Handle -> m (a,Handle) }
-- | 'OutRedirT' on a blank 'IO' monad
type OutRedir = OutRedirT IO

instance Monad m => Monad (OutRedirT m) where
  return a = OutRedir $ \h -> return (a,h)
  (OutRedir r) >>= f = OutRedir $ \h -> do (a,h') <- r h; runOutRedirT' (f a) h'

instance MonadTrans OutRedirT where
  lift m = OutRedir $ \h -> do a <- m; return (a,h)

instance MonadIO m => MonadIO (OutRedirT m) where
  liftIO = lift . liftIO

instance MonadIO m => MonadPrinter (OutRedirT m) where
  mprint s = OutRedir $ \h -> do liftIO $ hPutStr h s; return ((),h)
  mflush = OutRedir $ \h -> do liftIO $ hFlush h; return ((),h)

instance Monad m => Functor (OutRedirT m) where
  fmap f a = OutRedir $ \h -> do (a',h') <- runOutRedirT' a h; return (f a',h')

-- | Run 'OutRedirT' with a 'Handle'
runOutRedirT :: Functor m => OutRedirT m a -> Handle -> m a
runOutRedirT m h = fmap fst $ runOutRedirT' m h

-- | Run 'OutRedir' with a 'Handle'
runOutRedir :: OutRedir a -> Handle -> IO a
runOutRedir = runOutRedirT

-- | Run 'OutRedirT' with a 'FilePath'
runOutRedirFT :: (Functor m,MonadIO m) => OutRedirT m a -> FilePath -> IOMode -> m a
runOutRedirFT m fp md
  | md `elem` [AppendMode,WriteMode] = do
    h <- liftIO $ openFile fp md
    a <- runOutRedirT m h
    liftIO $ hClose h
    return a
  | otherwise = error "runOutRedirFT does only accept AppendMode or WriteMode."

-- | Run 'OutRedir' with a 'FilePath'
runOutRedirF :: OutRedir a -> FilePath -> IOMode -> IO a
runOutRedirF = runOutRedirFT

-- Definition of RecorderT + instances
-- | Catches all output (much like VAR=$(...) in the shell)
newtype RecorderT m a = Recorder { runRecorderT' :: [String] -> m (a,[String]) }
-- | 'RecorderT' on the 'Identity'
type Recorder = RecorderT Identity

instance Monad m => Monad (RecorderT m) where
  return a = Recorder $ \s -> return (a,s)
  (Recorder r) >>= f = Recorder $ \s -> do (a,s') <- r s; runRecorderT' (f a) s'

instance MonadTrans RecorderT where
  lift m = Recorder $ \s -> do a <- m; return (a,s)

instance Monad m => MonadPrinter (RecorderT m) where
  mprint s = Recorder $ \s' -> return ((),s:s')

instance Monad m => Functor (RecorderT m) where
  fmap f a = Recorder $ \s -> do (a',s') <- runRecorderT' a s; return (f a',s')

instance MonadIO m => MonadIO (RecorderT m) where
  liftIO = lift . liftIO

-- Helper methods for RecorderT
-- | The recorder state. Use this together with 'replay', 'replayM' or 'replay_'.
newtype Replayable = Replayable [String]
instance Show Replayable where show r = show ((\(Replayable x) -> length x) r) ++ ":" ++ show (replay r)

-- | Replay a recorder state inside a 'Monad'.
replayM :: Monad m => m Replayable -> m String
replayM r = do (Replayable r') <- r; return (concat $ reverse r')

-- | Replay a recorder state in a pure context.
replay :: Replayable -> String
replay (Replayable r) = concat $ reverse r

-- | Replay the current recorder state without leaving the recorder.
replay_ :: Monad m => RecorderT m String
replay_ = Recorder $ \s -> return (concat $ reverse s,s)

-- | Run 'Recorder' and also return its state.
runRecorder :: Recorder a -> (a,Replayable)
runRecorder = second Replayable . runIdentity . flip runRecorderT' []

-- | Run 'RecorderT' and also return its state.
runRecorderT :: (Functor m,Monad m) => RecorderT m a -> m (a,Replayable)
runRecorderT = fmap (second Replayable) . flip runRecorderT' []

-- | Line-terminating alternative to 'mprint'
mprintLn :: MonadPrinter m => String -> m ()
mprintLn = mprint . (++"\n")

-- Shell-like syntax
-- | Redirection target that discards input.
data DiscardO = DiscardO
-- | Redirection target that records input.
data RecordO = RecordO
-- | Class for all redirection targets.
class RedirectionTarget t mt a r | t -> mt, t a -> r where
  -- | Overwriting redirection.
  (.>.) :: (Functor m,MonadIO m,MonadPrinter (mt m)) => mt m a -> t -> m r
  -- | Appending redirection.
  (.>>.) :: (Functor m,MonadIO m,MonadPrinter (mt m)) => mt m a -> t -> m r
  (.>>.) = (.>.)
instance RedirectionTarget DiscardO DeafT a a where
  m .>. _ = runDeafT m
instance RedirectionTarget RecordO RecorderT a (a,Replayable) where
  m .>. _ = runRecorderT m
instance RedirectionTarget FilePath OutRedirT a a where
  m .>. fp = runOutRedirFT m fp WriteMode
  m .>>. fp = runOutRedirFT m fp AppendMode