-- {-# LANGUAGE OverloadedStrings #-}
-- |
-- Module     : Network.Wai.Middleware.Time
-- License    : WTFPL2
--
-- Maintainer : tv@destroy.retiolum
--
-- WAI X-Response-Time middleware.

module Network.Wai.Middleware.Time
    ( time
    ) where

import Control.Monad.Trans    (liftIO)
import Data.ByteString.Char8  (ByteString, pack)
import Data.CaseInsensitive       (mk)
import Network.HTTP.Types.Header  (HeaderName)
import Network.HTTP.Types         (ResponseHeaders)
import Network.Wai.Internal   (Response (..))
import Network.Wai            (Middleware)
import System.Clock           (getTime, Clock(Monotonic), sec, nsec, TimeSpec)
import Text.Printf            (printf)

getTimeMonotonic :: IO TimeSpec
getTimeMonotonic = liftIO $ getTime Monotonic

-- | Use time to append an X-Response-Time header to the response.
time :: Middleware
time app rq = do
  t0 <- getTimeMonotonic
  rs <- app rq
  t1 <- getTimeMonotonic
  return $ addResponseTime (timeSpecsToByteString t0 t1) rs

addResponseTime :: ByteString -> Response -> Response
addResponseTime xrt rs = do
  case rs of
    (ResponseFile st hs path part) ->
      ResponseFile st (responseTime xrt hs) path part
    (ResponseBuilder st hs b) ->
      ResponseBuilder st (responseTime xrt hs) b
    (ResponseSource st hs s) ->
      ResponseSource st (responseTime xrt hs) s
    (ResponseRaw rawapp r) ->
      ResponseRaw rawapp $ addResponseTime xrt r

hResponseTime :: HeaderName
hResponseTime = mk . pack $ "X-Response-Time"

responseTime :: ByteString -> ResponseHeaders -> ResponseHeaders
responseTime xrt =
  (:) (hResponseTime, xrt)

timeSpecsToByteString :: TimeSpec -> TimeSpec -> ByteString
timeSpecsToByteString t0 t1 =
    pack $ printf "%d.%09ds" s ns
  where
    s = (sec t1) - (sec t0)
    ns = (nsec t1) - (nsec t0)