-- | A light-weight wrapper around @Network.Wai@ to provide easy conduit support.
module Network.Wai.Conduit
    ( -- * Request body
      sourceRequestBody
      -- * Response body
    , responseSource
    , responseRawSource
      -- * Re-export
    , module Network.Wai
    ) where

import Network.Wai
import Data.Conduit
import Control.Monad.IO.Class (MonadIO, liftIO)
import Data.ByteString (ByteString)
import qualified Data.ByteString as S
import Control.Monad (unless)
import Network.HTTP.Types
import Data.ByteString.Builder (Builder)
import Data.IORef
import qualified Data.Conduit.List as CL

-- | Stream the request body.
--
-- Since 3.0.0
sourceRequestBody :: MonadIO m => Request -> Source m ByteString
sourceRequestBody req =
    loop
  where
    go = liftIO (requestBody req)

    loop = do
        bs <- go
        unless (S.null bs) $ do
            yield bs
            loop

-- | Create an HTTP response out of a @Source@.
--
-- Since 3.0.0
responseSource :: Status -> ResponseHeaders -> Source IO (Flush Builder) -> Response
responseSource s hs src = responseStream s hs $ \send flush ->
    src $$ CL.mapM_ (\mbuilder ->
        case mbuilder of
            Chunk b -> send b
            Flush -> flush)

-- | Create a raw response using a @Source@ and @Sink@ to represent the input
-- and output, respectively.
--
-- Since 3.0.0
responseRawSource :: (MonadIO m, MonadIO n)
                  => (Source m ByteString -> Sink ByteString n () -> IO ())
                  -> Response
                  -> Response
responseRawSource app =
    responseRaw app'
  where
    app' recv send =
        app src sink
      where
        src = do
            bs <- liftIO recv
            unless (S.null bs) $ do
                yield bs
                src
        sink = CL.mapM_ $ liftIO . send