{-# LANGUAGE OverloadedStrings #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE TupleSections #-}

module Network.Wai.Handler.Warp.Response (
    sendResponse
  ) where

import Blaze.ByteString.Builder (fromByteString, Builder, toByteStringIO, flush)
import Blaze.ByteString.Builder.HTTP (chunkedTransferEncoding, chunkedTransferTerminator)
import Control.Applicative
import Control.Exception
import Control.Monad.IO.Class (MonadIO, liftIO)
import Data.ByteString (ByteString)
import qualified Data.ByteString.Char8 as B (pack)
import qualified Data.CaseInsensitive as CI
import Data.Conduit
import Data.Conduit.Blaze (builderToByteString)
import qualified Data.Conduit.List as CL
import Data.Maybe (isJust, listToMaybe)
import Data.Monoid (mappend)
import qualified Network.HTTP.Types as H
import Network.Wai
import Network.Wai.Handler.Warp.ReadInt
import Network.Wai.Handler.Warp.Settings
import qualified Network.Wai.Handler.Warp.ResponseHeader as RH
import qualified Network.Wai.Handler.Warp.Timeout as T
import Network.Wai.Handler.Warp.Types
import qualified System.PosixCompat.Files as P
import Network.HTTP.Attoparsec (parseByteRanges)
import Numeric (showInt)

----------------------------------------------------------------
----------------------------------------------------------------

sendResponse :: Settings
             -> Cleaner -> Request -> Connection -> Response
             -> ResourceT IO Bool

----------------------------------------------------------------

sendResponse settings cleaner req conn (ResponseFile s0 hs0 path mpart0) =
    headerAndLength >>= sendResponse'
  where
    hs = addAccept hs0
    th = threadHandle cleaner
    headerAndLength = liftIO . try $ do
        (hadLength, cl) <-
            case readInt <$> checkLength hs of
                Just cl -> return (True, cl)
                Nothing ->
                    case mpart0 of
                        Just part -> return (False, fromIntegral $ filePartByteCount part)
                        Nothing -> (False, ) . fromIntegral . P.fileSize
                               <$> P.getFileStatus path
        let (s, addRange, beg, end) =
                case mpart0 of
                    Just part -> (s0, id, filePartOffset part, filePartByteCount part)
                    Nothing ->
                        case lookup H.hRange (requestHeaders req) >>= parseByteRanges >>= listToMaybe of
                            Just range | s0 == H.status200 ->
                                case range of
                                    H.ByteRangeFrom from -> rangeRes cl from (cl - 1)
                                    H.ByteRangeFromTo from to -> rangeRes cl from to
                                    H.ByteRangeSuffix count -> rangeRes cl (cl - count) (cl - 1)
                            _ -> (s0, id, 0, cl)
            hs'
                | hadLength = hs
                | otherwise = addLength end hs
        return (s, addRange hs', beg, end)

    rangeRes cl from to = (H.status206, (("Content-Range", rangeHeader cl from to):), from, to - from + 1)

    rangeHeader total from to = B.pack
      $ 'b' : 'y': 't' : 'e' : 's' : ' '
      : showInt from
      ( '-'
      : showInt to
      ( '/'
      : showInt total ""))

    sendResponse' (Right (s, lengthyHeaders, beg, end))
      | hasBody s req = liftIO $ do
          lheader <- composeHeader settings version s lengthyHeaders
          connSendFile conn path beg end (T.tickle th) [lheader] cleaner
          T.tickle th
          return isPersist
      | otherwise = liftIO $ do
          composeHeader settings version s hs >>= connSendAll conn
          T.tickle th
          return isPersist -- FIXME isKeepAlive?
      where
        version = httpVersion req
        (isPersist,_) = infoFromRequest req

    sendResponse' (Left (_ :: SomeException)) =
        sendResponse settings cleaner req conn notFound
      where
        notFound = responseLBS H.status404 [(H.hContentType, "text/plain")] "File not found"

----------------------------------------------------------------

sendResponse settings cleaner req conn (ResponseBuilder s hs b)
  | hasBody s req = liftIO $ do
      header <- composeHeaderBuilder settings version s hs needsChunked
      let body
            | needsChunked = header `mappend` chunkedTransferEncoding b
                                    `mappend` chunkedTransferTerminator
            | otherwise    = header `mappend` b
      flip toByteStringIO body $ \bs -> do
          connSendAll conn bs
          T.tickle th
      return isKeepAlive
  | otherwise = liftIO $ do
      composeHeader settings version s hs >>= connSendAll conn
      T.tickle th
      return isPersist
  where
    th = threadHandle cleaner
    version = httpVersion req
    reqinfo@(isPersist,_) = infoFromRequest req
    (isKeepAlive, needsChunked) = infoFromResponse hs reqinfo

----------------------------------------------------------------

sendResponse settings cleaner req conn (ResponseSource s hs bodyFlush)
  | hasBody s req = do
      header <- liftIO $ composeHeaderBuilder settings version s hs needsChunked
      let src = CL.sourceList [header] `mappend` cbody
      src $$ builderToByteString =$ connSink conn th
      return isKeepAlive
  | otherwise = liftIO $ do
      composeHeader settings version s hs >>= connSendAll conn
      T.tickle th
      return isPersist
  where
    th = threadHandle cleaner
    body =
           mapOutput (\x -> case x of
                    Flush -> flush
                    Chunk builder -> builder)
           bodyFlush
    cbody = if needsChunked then body $= chunk else body
    -- FIXME perhaps alloca a buffer per thread and reuse that in all
    -- functions below. Should lessen greatly the GC burden (I hope)
    chunk :: Conduit Builder (ResourceT IO) Builder
    chunk = await >>= maybe (yield chunkedTransferTerminator) (\x -> yield (chunkedTransferEncoding x) >> chunk)
    version = httpVersion req
    reqinfo@(isPersist,_) = infoFromRequest req
    (isKeepAlive, needsChunked) = infoFromResponse hs reqinfo

----------------------------------------------------------------
----------------------------------------------------------------

-- | Use 'connSendAll' to send this data while respecting timeout rules.
connSink :: Connection -> T.Handle -> Sink ByteString (ResourceT IO) ()
connSink Connection { connSendAll = send } th =
    sink
  where
    sink = await >>= maybe close push
    close = liftIO (T.resume th)
    push x = do
        liftIO $ T.resume th
        liftIO $ send x
        liftIO $ T.pause th
        sink
    -- We pause timeouts before passing control back to user code. This ensures
    -- that a timeout will only ever be executed when Warp is in control. We
    -- also make sure to resume the timeout after the completion of user code
    -- so that we can kill idle connections.

----------------------------------------------------------------

infoFromRequest :: Request -> (Bool,Bool)
infoFromRequest req = (checkPersist req, checkChunk req)

checkPersist :: Request -> Bool
checkPersist req
    | ver == H.http11 = checkPersist11 conn
    | otherwise       = checkPersist10 conn
  where
    ver = httpVersion req
    conn = lookup H.hConnection $ requestHeaders req
    checkPersist11 (Just x)
        | CI.foldCase x == "close"      = False
    checkPersist11 _                    = True
    checkPersist10 (Just x)
        | CI.foldCase x == "keep-alive" = True
    checkPersist10 _                    = False

checkChunk :: Request -> Bool
checkChunk req = httpVersion req == H.http11

----------------------------------------------------------------

infoFromResponse :: H.ResponseHeaders -> (Bool,Bool) -> (Bool,Bool)
infoFromResponse hs (isPersist,isChunked) = (isKeepAlive, needsChunked)
  where
    needsChunked = isChunked && not hasLength
    isKeepAlive = isPersist && (isChunked || hasLength)
    hasLength = isJust $ checkLength hs

checkLength :: H.ResponseHeaders -> Maybe ByteString
checkLength = lookup H.hContentLength

----------------------------------------------------------------

hasBody :: H.Status -> Request -> Bool
hasBody s req = sc /= 204
             && sc /= 304
             && sc >= 200
             && method /= H.methodHead
  where
    sc = H.statusCode s
    method = requestMethod req

----------------------------------------------------------------

addLength :: Integer -> H.ResponseHeaders -> H.ResponseHeaders
addLength cl hdrs = (H.hContentLength, B.pack $ show cl) : hdrs

addAccept :: H.ResponseHeaders -> H.ResponseHeaders
addAccept = (("Accept-Ranges", "bytes"):)

addEncodingHeader :: H.ResponseHeaders -> H.ResponseHeaders
addEncodingHeader hdrs = (hTransferEncoding, "chunked") : hdrs

addServerHeader :: Settings -> H.ResponseHeaders -> H.ResponseHeaders
addServerHeader settings hdrs = case lookup hServer hdrs of
    Nothing -> warpVersionHeader settings : hdrs
    Just _  -> hdrs

warpVersionHeader :: Settings -> H.Header
warpVersionHeader settings = (hServer, settingsServerName settings)

----------------------------------------------------------------

composeHeader :: Settings
              -> H.HttpVersion
              -> H.Status
              -> H.ResponseHeaders
              -> IO ByteString
composeHeader settings version s hs =
    RH.composeHeader version s
  $ addServerHeader settings hs

composeHeaderBuilder :: Settings -> H.HttpVersion -> H.Status -> H.ResponseHeaders -> Bool -> IO Builder
composeHeaderBuilder settings ver s hs True =
    fromByteString <$> composeHeader settings ver s (addEncodingHeader hs)
composeHeaderBuilder settings ver s hs False =
    fromByteString <$> composeHeader settings ver s hs