module Text.Chatty.Finalizer where

import Control.Monad
import Control.Monad.IO.Class
import Control.Monad.Trans.Class
import System.IO

-- | Class for all handle-finalizing monads. Required for file redirections.
class Monad m => MonadFinalizer m where
  -- | Enqueue handle
  mqfh :: Handle -> m ()
  -- | Enqueue list of handles
  mqfhs :: [Handle] -> m ()
  mqfhs = foldr ((>>) . mqfh) (return ())
  -- | Finalize all queued handles
  mfin :: m ()

-- | Handle-closing transformer
newtype HandleCloserT m a = HandleCloser { runHandleCloserT :: [Handle] -> m (a,[Handle]) }

instance Monad m => Monad (HandleCloserT m) where
  return a = HandleCloser $ \hs -> return (a,hs)
  (HandleCloser m) >>= f = HandleCloser $ \hs -> do (a,hs') <- m hs; runHandleCloserT (f a) hs'

instance MonadTrans HandleCloserT where
  lift m = HandleCloser $ \hs -> do a <- m; return (a,hs)

instance Monad m => Functor (HandleCloserT m) where
  fmap f a = HandleCloser $ \hs -> do (a',hs') <- runHandleCloserT a hs; return (f a',hs')

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

instance MonadIO m => MonadFinalizer (HandleCloserT m) where
  mqfh h = HandleCloser $ \hs -> return ((),h:hs)
  mfin = HandleCloser $ \hs -> do
    sequence_ $ fmap (liftIO.hClose) hs
    return ((),[])

-- | Run function with handle closer
withLazyIO :: (MonadIO m,Functor m) => HandleCloserT m a -> m a
withLazyIO m = fmap fst $ runHandleCloserT (do a <- m; mfin; return a) []