content: {-# LANGUAGE BangPatterns      #-}
  {-# LANGUAGE OverloadedStrings #-}
  
  module Test.Common.TestHandler (testHandler) where
  
  ------------------------------------------------------------------------------
  import           Control.Concurrent            (threadDelay)
  import           Control.Exception             (throwIO)
  import           Control.Monad                 (liftM)
  import           Control.Monad.IO.Class        (MonadIO (liftIO))
  import           Data.ByteString.Builder       (Builder, byteString)
  import           Data.ByteString.Builder.Extra (flush)
  import qualified Data.ByteString.Char8         as S
  import qualified Data.ByteString.Lazy.Char8    as L
  import           Data.List                     (sort)
  import qualified Data.Map                      as Map
  import           Data.Maybe                    (fromMaybe)
  import           Data.Monoid                   (Monoid (mappend, mconcat, mempty))
  ------------------------------------------------------------------------------
  import           Snap.Core                     (Request (rqParams, rqURI), Snap, getParam, getRequest, logError, modifyResponse, redirect, route, rqClientAddr, rqClientPort, setContentLength, setContentType, setHeader, setResponseBody, setResponseCode, setTimeout, transformRequestBody, writeBS, writeBuilder, writeLBS)
  import           Snap.Internal.Debug           ()
  import           Snap.Util.FileServe           (serveDirectory)
  import           Snap.Util.FileUploads         (PartInfo (partContentType, partFileName), allowWithMaximumSize, defaultUploadPolicy, disallow, handleFileUploads)
  import           Snap.Util.GZip                (noCompression, withCompression)
  import           System.Directory              (createDirectoryIfMissing)
  import           System.IO.Streams             (OutputStream)
  import qualified System.IO.Streams             as Streams
  import           Test.Common.Rot13             (rot13)
  
  
  ------------------------------------------------------------------------------
  -- timeout handling
  ------------------------------------------------------------------------------
  timeoutTickleHandler :: Snap ()
  timeoutTickleHandler = do
      noCompression   -- FIXME: remove this when zlib-bindings and
                      -- zlib-enumerator support gzip stream flushing
      modifyResponse $ setResponseBody (trickleOutput 10)
                     . setContentType "text/plain"
      setTimeout 2
  
  
  badTimeoutTickleHandler :: Snap ()
  badTimeoutTickleHandler = do
      noCompression   -- FIXME: remove this when zlib-bindings and
                      -- zlib-enumerator support gzip stream flushing
      modifyResponse $ setResponseBody (trickleOutput 10)
                     . setContentType "text/plain"
      setTimeout 2
  
  
  trickleOutput :: Int -> OutputStream Builder -> IO (OutputStream Builder)
  trickleOutput n os = do
      Streams.fromList dots >>= Streams.mapM f >>= Streams.supplyTo os
      return os
    where
      dots = replicate n ".\n"
      f x  = threadDelay 1000000 >> return (byteString x `mappend` flush)
  
  
  ------------------------------------------------------------------------------
  pongHandler :: Snap ()
  pongHandler = modifyResponse $
                setResponseBody body .
                setContentType "text/plain" .
                setContentLength 4
    where
      body os = do Streams.write (Just $ byteString "PONG") os
                   return os
  
  echoUriHandler :: Snap ()
  echoUriHandler = do
      req <- getRequest
      writeBS $ rqURI req
  
  echoHandler :: Snap ()
  echoHandler = transformRequestBody return
  
  rot13Handler :: Snap ()
  rot13Handler = transformRequestBody (Streams.map rot13)
  
  bigResponseHandler :: Snap ()
  bigResponseHandler = do
      let sz = 4000000
      let s = L.take sz $ L.cycle $ L.fromChunks [S.replicate 400000 '.']
      modifyResponse $ setContentLength $ fromIntegral sz
      writeLBS s
  
  
  responseHandler :: Snap ()
  responseHandler = do
      !code <- liftM (read . S.unpack . fromMaybe "503") $ getParam "code"
      modifyResponse $ setResponseCode code
      writeBS $ S.pack $ show code
  
  
  uploadForm :: Snap ()
  uploadForm = do
      modifyResponse $ setContentType "text/html"
      writeBS form
  
    where
      form = S.concat [ "<html><head><title>Upload a file</title></head><body>\n"
                      , "<p>Upload some <code>text/plain</code> files:</p>\n"
                      , "<form method='post' "
                      , "enctype='multipart/form-data' "
                      , "action='/upload/handle'>\n"
                      , "<input type='file' "
                      , "accept='text/plain' "
                      , "multiple='true' "
                      , "name='file'></input>\n"
                      , "<input type='submit' name='Submit'></input>\n"
                      , "</form></body></html>" ]
  
  
  uploadHandler :: Snap ()
  uploadHandler = do
      logError "uploadHandler"
      liftIO $ createDirectoryIfMissing True tmpdir
      files <- handleFileUploads tmpdir defaultUploadPolicy partPolicy hndl
      let m = sort files
  
      params <- liftM (Prelude.map (\(a,b) -> (a,S.concat b)) .
                       Map.toAscList .
                       rqParams) getRequest
  
      modifyResponse $ setContentType "text/plain"
      writeBuilder $ buildRqParams params `mappend` buildFiles m
  
    where
      f p = fromMaybe "-" $ partFileName p
  
      hndl _ (Left e)          = throwIO e
      hndl partInfo (Right fp) = do
          !c <- liftIO $ S.readFile fp
          return $! (f partInfo, c)
  
      builder _ [] = mempty
      builder ty ((k,v):xs) =
          mconcat [ byteString ty
                  , byteString ":\n"
                  , byteString k
                  , byteString "\nValue:\n"
                  , byteString v
                  , byteString "\n\n"
                  , builder ty xs ]
  
  
      buildRqParams = builder "Param"
      buildFiles = builder "File"
  
      tmpdir = "dist/filetmp"
      partPolicy partInfo = if partContentType partInfo == "text/plain"
                              then allowWithMaximumSize 200000
                              else disallow
  
  serverHeaderHandler :: Snap ()
  serverHeaderHandler = modifyResponse $ setHeader "Server" "foo"
  
  
  chunkedResponse :: Snap ()
  chunkedResponse = writeBS "chunked"
  
  
  remoteAddrPort :: Snap ()
  remoteAddrPort = do
      rq <- getRequest
      let addr = rqClientAddr rq
      let port = rqClientPort rq
      let out = S.concat [ addr, ":", S.pack (show port) ]
      modifyResponse $ setContentLength $ fromIntegral $ S.length out
      writeBS out
  
  
  testHandler :: Snap ()
  testHandler = withCompression $
      route [ ("pong"              , pongHandler                       )
            , ("redirect"          , redirect "/pong"                  )
            , ("echo"              , echoHandler                       )
            , ("rot13"             , rot13Handler                      )
            , ("echoUri"           , echoUriHandler                    )
            , ("remoteAddrPort"    , remoteAddrPort                    )
            , ("fileserve"         , noCompression >>
                                     serveDirectory "testserver/static")
            , ("bigresponse"       , bigResponseHandler                )
            , ("respcode/:code"    , responseHandler                   )
            , ("upload/form"       , uploadForm                        )
            , ("upload/handle"     , uploadHandler                     )
            , ("timeout/tickle"    , timeoutTickleHandler              )
            , ("timeout/badtickle" , badTimeoutTickleHandler           )
            , ("server-header"     , serverHeaderHandler               )
            , ("chunked"           , chunkedResponse                   )
            ]
  