{-# LANGUAGE BangPatterns              #-}
{-# LANGUAGE CPP                       #-}
{-# LANGUAGE DeriveDataTypeable        #-}
{-# LANGUAGE ExistentialQuantification #-}
{-# LANGUAGE OverloadedStrings         #-}
{-# LANGUAGE ScopedTypeVariables       #-}

--------------------------------------------------------------------------------
-- | Helpers for running a 'Snap' web handler with compression.

module Snap.Util.GZip
  ( withCompression
  , withCompression'
  , noCompression
  , BadAcceptEncodingException
  , compressibleMimeTypes
  ) where

import           Control.Applicative              (Alternative ((<|>), many), Applicative ((*>), (<*), pure), (<$>))
import           Control.Exception                (Exception, throwIO)
import           Control.Monad                    (Functor (fmap), Monad ((>>), (>>=), return), MonadPlus (mplus), void, when)
import           Control.Monad.IO.Class           (MonadIO (liftIO))
import           Data.Attoparsec.ByteString.Char8 (Parser, char, endOfInput, isAlpha_ascii, isDigit, skipSpace, string, takeWhile, takeWhile1)
import           Data.ByteString.Builder          (Builder)
import           Data.ByteString.Char8            (ByteString)
import qualified Data.ByteString.Char8            as S (takeWhile)
import qualified Data.Char                        as Char (isSpace)
import           Data.Maybe                       (Maybe (Just, Nothing), fromMaybe, isJust, maybe)
import           Data.Set                         (Set)
import qualified Data.Set                         as Set (fromList, member)
import           Data.Typeable                    (Typeable)
import           Prelude                          (Either (..), Eq (..), IO, Show (show), id, not, ($), ($!), (&&), (++), (||))
import           Snap.Core                        (MonadSnap, clearContentLength, finishWith, getHeader, getRequest, getResponse, modifyResponse, modifyResponseBody, setHeader)
import           Snap.Internal.Debug              (debug)
import           Snap.Internal.Parsing            (fullyParse)
import           System.IO.Streams                (OutputStream)
import qualified System.IO.Streams                as Streams (compressBuilder, gzipBuilder)

------------------------------------------------------------------------------
-- | Runs a 'Snap' web handler with compression if available.
--
-- If the client has indicated support for @gzip@ or @deflate@ in its
-- @Accept-Encoding@ header, and the @Content-Type@ in the response is one of
-- the following types:
--
--   * @application/x-javascript@
--
--   * @application/json@
--
--   * @text/css@
--
--   * @text/html@
--
--   * @text/javascript@
--
--   * @text/plain@
--
--   * @text/xml@
--
--   * @application/x-font-truetype@
--
-- Then the given handler's output stream will be compressed,
-- @Content-Encoding@ will be set in the output headers, and the
-- @Content-Length@ will be cleared if it was set. (We can't process the
-- stream in O(1) space if the length is known beforehand.)
--
-- The wrapped handler will be run to completion, and then the 'Response'
-- that's contained within the 'Snap' monad state will be passed to
-- 'finishWith' to prevent further processing.
--
--
-- Example:
--
-- @
-- ghci> :set -XOverloadedStrings
-- ghci> import qualified "Data.Map" as M
-- ghci> import qualified "Snap.Test" as T
-- ghci> let r = T.get \"/\" M.empty >> T.addHeader \"Accept-Encoding\" \"gzip,deflate\"
-- ghci> let h = 'Snap.Core.modifyResponse' ('Snap.Core.setContentType' \"text\/plain\") >> 'Snap.Core.writeBS' \"some text\"
-- ghci> T.runHandler r h
-- HTTP\/1.1 200 OK
-- content-type: text\/plain
-- server: Snap\/test
-- date: Fri, 08 Aug 2014 15:40:45 GMT
--
-- some text
-- ghci> T.runHandler r ('withCompression' h)
-- HTTP\/1.1 200 OK
-- content-type: text\/plain
-- vary: Accept-Encoding
-- content-encoding: gzip
-- server: Snap\/test
-- date: Fri, 08 Aug 2014 15:40:10 GMT
--
--
-- @
withCompression :: MonadSnap m
                => m a   -- ^ the web handler to run
                -> m ()
withCompression = withCompression' compressibleMimeTypes


------------------------------------------------------------------------------
-- | The same as 'withCompression', with control over which MIME types to
-- compress.
withCompression' :: MonadSnap m
                 => Set ByteString
                    -- ^ set of compressible MIME types
                 -> m a
                    -- ^ the web handler to run
                 -> m ()
withCompression' mimeTable action = do
    _    <- action
    resp <- getResponse

    -- If a content-encoding is already set, do nothing. This prevents
    -- "withCompression $ withCompression m" from ruining your day.
    when (not $ isJust $ getHeader "Content-Encoding" resp) $ do
       let mbCt = fmap chop $ getHeader "Content-Type" resp

       debug $ "withCompression', content-type is " ++ show mbCt

       case mbCt of
         (Just ct) -> when (Set.member ct mimeTable) chkAcceptEncoding
         _         -> return $! ()


    getResponse >>= finishWith

  where
    chop = S.takeWhile (\c -> c /= ';' && not (Char.isSpace c))

    chkAcceptEncoding = do
        req <- getRequest
        debug $ "checking accept-encoding"
        let mbAcc = getHeader "Accept-Encoding" req
        debug $ "accept-encoding is " ++ show mbAcc
        let s = fromMaybe "" mbAcc

        types <- liftIO $ parseAcceptEncoding s

        chooseType Nothing types

    chooseType !m []               = maybe (return $! ()) id m
    chooseType !_ ("gzip":_)       = gzipCompression "gzip"
    chooseType !m ("deflate":xs)   =
        chooseType (m `mplus` Just (compressCompression "deflate")) xs

    chooseType !_ ("x-gzip":_)     = gzipCompression "x-gzip"
    chooseType !m ("x-deflate":xs) =
        chooseType (m `mplus` Just (compressCompression "x-deflate")) xs
    chooseType !m (_:xs)           = chooseType m xs


------------------------------------------------------------------------------
-- | Turn off compression by setting \"Content-Encoding: identity\" in the
-- response headers. 'withCompression' is a no-op when a content-encoding is
-- already set.
noCompression :: MonadSnap m => m ()
noCompression = modifyResponse $ setHeader "Content-Encoding" "identity"


------------------------------------------------------------------------------
-- private following
------------------------------------------------------------------------------


------------------------------------------------------------------------------
compressibleMimeTypes :: Set ByteString
compressibleMimeTypes = Set.fromList [ "application/x-font-truetype"
                                     , "application/x-javascript"
                                     , "application/json"
                                     , "text/css"
                                     , "text/html"
                                     , "text/javascript"
                                     , "text/plain"
                                     , "text/xml" ]




------------------------------------------------------------------------------
gzipCompression :: MonadSnap m => ByteString -> m ()
gzipCompression ce = modifyResponse f
  where
    f r = setHeader "Content-Encoding" ce    $
          setHeader "Vary" "Accept-Encoding" $
          clearContentLength                 $
          modifyResponseBody gcompress r


------------------------------------------------------------------------------
compressCompression :: MonadSnap m => ByteString -> m ()
compressCompression ce = modifyResponse f
  where
    f r = setHeader "Content-Encoding" ce    $
          setHeader "Vary" "Accept-Encoding" $
          clearContentLength                 $
          modifyResponseBody ccompress r


------------------------------------------------------------------------------
gcompress :: (OutputStream Builder -> IO (OutputStream Builder))
          -> OutputStream Builder
          -> IO (OutputStream Builder)
gcompress body stream = Streams.gzipBuilder 5 stream >>= body


------------------------------------------------------------------------------
ccompress :: (OutputStream Builder -> IO (OutputStream Builder))
          -> OutputStream Builder
          -> IO (OutputStream Builder)
ccompress body stream = Streams.compressBuilder 5 stream >>= body


------------------------------------------------------------------------------
-- We're not gonna bother with quality values; we'll do gzip or compress in
-- that order.
acceptParser :: Parser [ByteString]
acceptParser = do
    xs <- ((:[]) <$> encoding) <|> (return $! [])
    ys <- many (char ',' *> encoding)
    endOfInput
    return $! xs ++ ys
  where
    encoding = skipSpace *> c <* skipSpace

    c = do
        x <- coding
        qvalue <|> (return $! ())
        return x

    qvalue = do
        skipSpace
        void $! char ';'
        skipSpace
        void $! char 'q'
        skipSpace
        void $! char '='
        float
        return $! ()

    coding = string "*" <|> takeWhile1 isCodingChar

    isCodingChar ch = isDigit ch || isAlpha_ascii ch || ch == '-' || ch == '_'

    float = takeWhile isDigit >>
            (char '.' >> takeWhile isDigit >> (pure $! ())) <|> (pure $! ())


------------------------------------------------------------------------------
-- | Thrown when the 'Accept-Encoding' request header has invalid format.
data BadAcceptEncodingException = BadAcceptEncodingException
   deriving (Typeable)


------------------------------------------------------------------------------
instance Show BadAcceptEncodingException where
    show BadAcceptEncodingException = "bad 'accept-encoding' header"


------------------------------------------------------------------------------
instance Exception BadAcceptEncodingException


------------------------------------------------------------------------------
parseAcceptEncoding :: ByteString -> IO [ByteString]
parseAcceptEncoding s =
    case r of
      Left _ -> throwIO BadAcceptEncodingException
      Right x -> return x
  where
    r = fullyParse s acceptParser