    {-# LANGUAGE Rank2Types, ScopedTypeVariables #-}
    ---------------------------------------------------------
    -- |
    -- Module        : Network.Wai.Middleware.Gzip
    -- Copyright     : Michael Snoyman
    -- License       : BSD3
    --
    -- Maintainer    : Michael Snoyman <michael@snoyman.com>
    -- Stability     : Unstable
    -- Portability   : portable
    --
    -- Automatic gzip compression of responses.
    --
    ---------------------------------------------------------
    module Network.Wai.Middleware.Gzip
        ( gzip
        , GzipSettings
        , gzipFiles
        , GzipFiles (..)
        , gzipCheckMime
        , def
        , defaultCheckMime
        ) where

    import Network.Wai
    import Data.Maybe (fromMaybe, isJust)
    import qualified Data.ByteString.Char8 as S8
    import qualified Data.ByteString as S
    import Data.Default.Class
    import Network.HTTP.Types ( Status, Header, hContentEncoding, hUserAgent
                              , hContentType, hContentLength)
    import System.Directory (doesFileExist, createDirectoryIfMissing)
    import Blaze.ByteString.Builder (fromByteString)
    import Control.Exception (try, SomeException)
    import qualified Data.Set as Set
    import Network.Wai.Header
    import Network.Wai.Internal
    import qualified Data.Streaming.Blaze as B
    import qualified Data.Streaming.Zlib as Z
    import qualified Blaze.ByteString.Builder as Blaze
    import Control.Monad (unless)
    import Data.Function (fix)
    import Control.Exception (throwIO)
    import qualified System.IO as IO
    import Data.ByteString.Lazy.Internal (defaultChunkSize)
    import Data.Word8 (_semicolon)

    data GzipSettings = GzipSettings
        { gzipFiles :: GzipFiles
        , gzipCheckMime :: S.ByteString -> Bool
        }

    data GzipFiles = GzipIgnore | GzipCompress | GzipCacheFolder FilePath
        deriving (Show, Eq, Read)

    instance Default GzipSettings where
        def = GzipSettings GzipIgnore defaultCheckMime

    defaultCheckMime :: S.ByteString -> Bool
    defaultCheckMime bs =
        S8.isPrefixOf "text/" bs || bs' `Set.member` toCompress
      where
        bs' = fst $ S.break (== _semicolon) bs
        toCompress = Set.fromList
            [ "application/json"
            , "application/javascript"
            , "application/ecmascript"
            ]

    -- | Use gzip to compress the body of the response.
    --
    -- Analyzes the \"Accept-Encoding\" header from the client to determine
    -- if gzip is supported.
    --
    -- Possible future enhancements:
    --
    -- * Only compress if the response is above a certain size.
    gzip :: GzipSettings -> Middleware
    gzip set app env sendResponse = app env $ \res ->
        case res of
            ResponseRaw{} -> sendResponse res
            ResponseFile{} | gzipFiles set == GzipIgnore -> sendResponse res
            _ -> if "gzip" `elem` enc && not isMSIE6 && not (isEncoded res) && (bigEnough res)
                    then
                        case (res, gzipFiles set) of
                            (ResponseFile s hs file Nothing, GzipCacheFolder cache) ->
                                case lookup hContentType hs of
                                    Just m
                                        | gzipCheckMime set m -> compressFile s hs file cache sendResponse
                                    _ -> sendResponse res
                            _ -> compressE set res sendResponse
                    else sendResponse res
      where
        enc = fromMaybe [] $ (splitCommas . S8.unpack)
                        `fmap` lookup "Accept-Encoding" (requestHeaders env)
        ua = fromMaybe "" $ lookup hUserAgent $ requestHeaders env
        isMSIE6 = "MSIE 6" `S.isInfixOf` ua
        isEncoded res = isJust $ lookup hContentEncoding $ responseHeaders res

        bigEnough rsp = case contentLength (responseHeaders rsp) of
          Nothing -> True -- This could be a streaming case
          Just len -> len >= minimumLength

        -- For a small enough response, gzipping will actually increase the size
        -- Potentially for anything less than 860 bytes gzipping could be a net loss
        -- The actual number is application specific though and may need to be adjusted
        -- http://webmasters.stackexchange.com/questions/31750/what-is-recommended-minimum-object-size-for-gzip-performance-benefits
        minimumLength = 860

    compressFile :: Status -> [Header] -> FilePath -> FilePath -> (Response -> IO a) -> IO a
    compressFile s hs file cache sendResponse = do
        e <- doesFileExist tmpfile
        if e
            then onSucc
            else do
                createDirectoryIfMissing True cache
                x <- try $
                     IO.withBinaryFile file IO.ReadMode $ \inH ->
                     IO.withBinaryFile tmpfile IO.WriteMode $ \outH -> do
                        deflate <- Z.initDeflate 7 $ Z.WindowBits 31
                        -- FIXME this code should write to a temporary file, then
                        -- rename to the final file
                        let goPopper popper = fix $ \loop -> do
                                res <- popper
                                case res of
                                    Z.PRDone -> return ()
                                    Z.PRNext bs -> do
                                        S.hPut outH bs
                                        loop
                                    Z.PRError ex -> throwIO ex
                        fix $ \loop -> do
                            bs <- S.hGetSome inH defaultChunkSize
                            unless (S.null bs) $ do
                                Z.feedDeflate deflate bs >>= goPopper
                                loop
                        goPopper $ Z.finishDeflate deflate
                either onErr (const onSucc) (x :: Either SomeException ()) -- FIXME bad! don't catch all exceptions like that!
      where
        onSucc = sendResponse $ responseFile s (fixHeaders hs) tmpfile Nothing

        onErr _ = sendResponse $ responseFile s hs file Nothing -- FIXME log the error message

        tmpfile = cache ++ '/' : map safe file
        safe c
            | 'A' <= c && c <= 'Z' = c
            | 'a' <= c && c <= 'z' = c
            | '0' <= c && c <= '9' = c
        safe '-' = '-'
        safe '_' = '_'
        safe _ = '_'

    compressE :: GzipSettings
              -> Response
              -> (Response -> IO a)
              -> IO a
    compressE set res sendResponse =
        case lookup hContentType hs of
            Just m | gzipCheckMime set m ->
                let hs' = fixHeaders hs
                 in wb $ \body -> sendResponse $ responseStream s hs' $ \sendChunk flush -> do
                        (blazeRecv, _) <- B.newBlazeRecv B.defaultStrategy
                        deflate <- Z.initDeflate 1 (Z.WindowBits 31)
                        let sendBuilder builder = do
                                popper <- blazeRecv builder
                                fix $ \loop -> do
                                    bs <- popper
                                    unless (S.null bs) $ do
                                        sendBS bs
                                        loop
                            sendBS bs = Z.feedDeflate deflate bs >>= deflatePopper
                            flushBuilder = do
                                sendBuilder Blaze.flush
                                deflatePopper $ Z.flushDeflate deflate
                                flush
                            deflatePopper popper = fix $ \loop -> do
                                result <- popper
                                case result of
                                    Z.PRDone -> return ()
                                    Z.PRNext bs' -> do
                                        sendChunk $ fromByteString bs'
                                        loop
                                    Z.PRError e -> throwIO e

                        body sendBuilder flushBuilder
                        sendBuilder Blaze.flush
                        deflatePopper $ Z.finishDeflate deflate
            _ -> sendResponse res
      where
        (s, hs, wb) = responseToStream res

    -- Remove Content-Length header, since we will certainly have a
    -- different length after gzip compression.
    fixHeaders :: [Header] -> [Header]
    fixHeaders =
        ((hContentEncoding, "gzip") :) . filter notLength
      where
        notLength (x, _) = x /= hContentLength

    splitCommas :: String -> [String]
    splitCommas [] = []
    splitCommas x =
        let (y, z) = break (== ',') x
         in y : splitCommas (dropWhile (== ' ') $ drop 1 z)