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

module Chatty.Printer where

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

-- Class Definition for MonadPrinter
class Monad m => MonadPrinter m where
  mprint :: String -> m ()
  mnoecho :: String -> m ()
  mnoecho = mprint
  mflush :: m ()
  mflush = return ()

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

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

-- Definition of DeafT + instances
-- 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
-- OutRedirT redirects all output to a given handle (much like >filename in the shell)
newtype OutRedirT m a = OutRedir { runOutRedirT' :: Handle -> m (a,Handle) }
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')

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

runOutRedir :: OutRedir a -> Handle -> IO a
runOutRedir = runOutRedirT

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."

runOutRedirF :: OutRedir a -> FilePath -> IOMode -> IO a
runOutRedirF = runOutRedirFT

-- Definition of RecorderT + instances
-- RecorderT catches all output (much like VAR=$(...) in the shell)
newtype RecorderT m a = Recorder { runRecorderT' :: [String] -> m (a,[String]) }
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
newtype Replayable = Replayable [String]
instance Show Replayable where show r = show ((\(Replayable x) -> length x) r) ++ ":" ++ show (replay r)

replayM :: Monad m => m Replayable -> m String
replayM r = do (Replayable r') <- r; return (concat $ reverse r')

replay :: Replayable -> String
replay (Replayable r) = concat $ reverse r

replay_ :: Monad m => RecorderT m String
replay_ = Recorder $ \s -> return (concat $ reverse s,s)

runRecorder :: Recorder a -> (a,Replayable)
runRecorder = second Replayable . runIdentity . flip runRecorderT' []

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
data DiscardO = DiscardO
data RecordO = RecordO
class RedirectionTarget t mt a r | t -> mt, t a -> r where
  (.>.) :: (Functor m,MonadIO m,MonadPrinter (mt m)) => mt m a -> t -> m r
  (.>>.) :: (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