    {-# LANGUAGE DeriveDataTypeable #-}
    {-# LANGUAGE RankNTypes #-}
    {-# LANGUAGE NamedFieldPuns #-}
    {-# LANGUAGE PatternGuards #-}
    {-# LANGUAGE CPP #-}

    module Network.Wai.Handler.Warp.HTTP2.Worker (
        Respond
      , response
      , worker
      ) where

    #if __GLASGOW_HASKELL__ < 709
    import Control.Applicative
    #endif
    import Control.Concurrent
    import Control.Concurrent.STM
    import Control.Exception (Exception, SomeException(..), AsyncException(..))
    import qualified Control.Exception as E
    import Control.Monad (void, when)
    import Data.Typeable
    import qualified Network.HTTP.Types as H
    import Network.HTTP2
    import Network.HTTP2.Priority
    import Network.Wai
    import Network.Wai.Handler.Warp.HTTP2.EncodeFrame
    import Network.Wai.Handler.Warp.HTTP2.Manager
    import Network.Wai.Handler.Warp.HTTP2.Types
    import Network.Wai.Handler.Warp.IORef
    import Network.Wai.HTTP2
        ( Chunk(..)
        , HTTP2Application
        , PushPromise
        , Responder(runResponder)
        , RespondFunc
        )
    import qualified Network.Wai.Handler.Warp.Settings as S
    import qualified Network.Wai.Handler.Warp.Timeout as T

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

    -- | An 'HTTP2Application' takes a function of status, headers, trailers, and
    -- body; this type implements that by currying some internal arguments.
    --
    -- The token type of the RespondFunc is set to be ().  This is a bit
    -- anti-climactic, but the real benefit of the token type is that the
    -- application is forced to call the responder, and making it a boring type
    -- doesn't break that property.
    --
    -- This is the argument to a 'Responder'.
    type Respond = IO () -> Stream -> RespondFunc ()

    -- | This function is passed to workers.  They also pass responses from
    -- 'HTTP2Application's to this function.  This function enqueues commands for
    -- the HTTP/2 sender.
    response :: Context -> Manager -> ThreadContinue -> Respond
    response ctx mgr tconf tickle strm s h strmbdy = do
        -- TODO(awpr) HEAD requests will still stream.

        -- We must not exit this WAI application.
        -- If the application exits, streaming would be also closed.
        -- So, this work occupies this thread.
        --
        -- We need to increase the number of workers.
        myThreadId >>= replaceWithAction mgr
        -- After this work, this thread stops to decrease the number of workers.
        setThreadContinue tconf False

        runStream ctx OResponse tickle strm s h strmbdy

    -- | Set up a waiter thread and run the stream body with functions to enqueue
    -- 'Sequence's on the stream's queue.
    runStream :: Context
              -> (Stream -> H.Status -> H.ResponseHeaders -> Aux -> Output)
              -> Respond
    runStream Context{outputQ} mkOutput tickle strm s h strmbdy = do
        -- Since 'Body' is loop, we cannot control it.
        -- So, let's serialize 'Builder' with a designated queue.
        sq <- newTBQueueIO 10 -- fixme: hard coding: 10
        tvar <- newTVarIO SyncNone
        let out = mkOutput strm s h (Persist sq tvar)
        -- Since we must not enqueue an empty queue to the priority
        -- queue, we spawn a thread to ensure that the designated
        -- queue is not empty.
        void $ forkIO $ waiter tvar sq strm outputQ
        atomically $ writeTVar tvar $ SyncNext out
        let write chunk = do
                atomically $ writeTBQueue sq $ case chunk of
                    BuilderChunk b -> SBuilder b
                    FileChunk path part -> SFile path part
                tickle
            flush  = atomically $ writeTBQueue sq SFlush
        trailers <- strmbdy write flush
        atomically $ writeTBQueue sq $ SFinish trailers

    -- | Handle abnormal termination of a stream: mark it as closed, send a reset
    -- frame, and call the user's 'settingsOnException' handler if applicable.
    cleanupStream :: Context -> S.Settings -> Stream -> Maybe Request -> Maybe SomeException -> IO ()
    cleanupStream Context{outputQ} set strm req me = do
        closed strm Killed
        let sid = streamNumber strm
            frame = resetFrame InternalError sid
        enqueueControl outputQ sid $ OFrame frame
        case me of
            Nothing -> return ()
            Just e -> S.settingsOnException set req e

    -- | Push the given 'Responder' to the client if the settings allow it
    -- (specifically 'enablePush' and 'maxConcurrentStreams').  Returns 'True' if
    -- the stream was actually pushed.
    --
    -- This is the push function given to an 'HTTP2Application'.
    pushResponder :: Context -> S.Settings -> Stream -> PushPromise -> Responder -> IO Bool
    pushResponder ctx set strm promise responder = do
        let Context{ http2settings
                   , pushConcurrency
                   } = ctx
        cnt <- readIORef pushConcurrency
        settings <- readIORef http2settings
        let enabled = enablePush settings
            fits = maybe True (cnt <) $ maxConcurrentStreams settings
            canPush = fits && enabled
        if canPush then
            actuallyPushResponder ctx set strm promise responder
          else
            return False

    -- | Set up a pushed stream and run the 'Responder' in its own thread.  Waits
    -- for the sender thread to handle the push request.  This can fail to push the
    -- stream and return 'False' if the sender dequeued the push request after the
    -- associated stream was closed.
    actuallyPushResponder :: Context -> S.Settings -> Stream -> PushPromise -> Responder -> IO Bool
    actuallyPushResponder ctx set strm promise responder = do
        let Context{ http2settings
                   , nextPushStreamId
                   , pushConcurrency
                   , streamTable
                   } = ctx
        -- Claim the next outgoing stream.
        newSid <- atomicModifyIORef nextPushStreamId $ \sid -> (sid+2, sid)
        ws <- initialWindowSize <$> readIORef http2settings

        newStrm <- newStream pushConcurrency newSid ws
        -- Section 5.3.5 of RFC 7540 defines the weight of push promise is 16.
        -- But we need not to follow the spec. So, this value would change
        -- if necessary.
        writeIORef (streamPrecedence newStrm) $
            toPrecedence $ defaultPriority { streamDependency = streamNumber strm }
        opened newStrm
        insert streamTable newSid newStrm

        -- Set up a channel for the sender to report back whether it pushed the
        -- stream.
        mvar <- newEmptyMVar

        let mkOutput = OPush strm promise mvar
            tickle = return ()
            respond = runStream ctx mkOutput

        -- TODO(awpr): synthesize a Request for 'settingsOnException'?
        _ <- forkIO $ runResponder responder (respond tickle newStrm) `E.catch`
            (cleanupStream ctx set strm Nothing . Just)

        takeMVar mvar

    data Break = Break deriving (Show, Typeable)

    instance Exception Break

    worker :: Context
           -> S.Settings
           -> T.Manager
           -> HTTP2Application
           -> (ThreadContinue -> Respond)
           -> IO ()
    worker ctx@Context{inputQ} set tm app respond = do
        tid <- myThreadId
        sinfo <- newStreamInfo
        tcont <- newThreadContinue
        let setup = T.register tm $ E.throwTo tid Break
        E.bracket setup T.cancel $ go sinfo tcont
      where
        go sinfo tcont th = do
            setThreadContinue tcont True

            ex <- E.try $ do
                T.pause th
                Input strm req <- atomically $ readTQueue inputQ
                setStreamInfo sinfo strm req
                T.resume th
                T.tickle th
                let responder = app req $ pushResponder ctx set strm
                runResponder responder $ respond tcont (T.tickle th) strm
            cont1 <- case ex of
                Right () -> return True
                Left  e@(SomeException _)
                  | Just Break        <- E.fromException e -> do
                      cleanup sinfo Nothing
                      return True
                  -- killed by the sender
                  | Just ThreadKilled <- E.fromException e -> do
                      cleanup sinfo Nothing
                      return False
                  | otherwise -> do
                      cleanup sinfo (Just e)
                      return True
            cont2 <- getThreadContinue tcont
            when (cont1 && cont2) $ go sinfo tcont th
        cleanup sinfo me = do
            m <- getStreamInfo sinfo
            case m of
                Nothing -> return ()
                Just (strm,req) -> do
                    cleanupStream ctx set strm (Just req) me
                    clearStreamInfo sinfo

    -- | A dedicated waiter thread to re-enqueue the stream in the priority tree
    -- whenever output becomes available.  When the sender drains the queue and
    -- moves on to another stream, it drops a message in the 'TVar', and this
    -- thread wakes up, waits for more output to become available, and re-enqueues
    -- the stream.
    waiter :: TVar Sync -> TBQueue Sequence -> Stream -> PriorityTree Output -> IO ()
    waiter tvar sq strm outQ = do
        -- waiting for actions other than SyncNone
        mx <- atomically $ do
            mout <- readTVar tvar
            case mout of
                SyncNone     -> retry
                SyncNext out -> do
                    writeTVar tvar SyncNone
                    return $ Just out
                SyncFinish   -> return Nothing
        case mx of
            Nothing  -> return ()
            Just out -> do
                -- ensuring that the streaming queue is not empty.
                atomically $ do
                    isEmpty <- isEmptyTBQueue sq
                    when isEmpty retry
                -- ensuring that stream window is greater than 0.
                enqueueWhenWindowIsOpen outQ out
                waiter tvar sq strm outQ

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

    -- | It would nice if responders could return values to workers.
    --   Unfortunately, 'ResponseReceived' is already defined in WAI 2.0.
    --   It is not wise to change this type.
    --   So, a reference is shared by a 'Respond' and its worker.
    --   The reference refers a value of this type as a return value.
    --   If 'True', the worker continue to serve requests.
    --   Otherwise, the worker get finished.
    newtype ThreadContinue = ThreadContinue (IORef Bool)

    newThreadContinue :: IO ThreadContinue
    newThreadContinue = ThreadContinue <$> newIORef True

    setThreadContinue :: ThreadContinue -> Bool -> IO ()
    setThreadContinue (ThreadContinue ref) x = writeIORef ref x

    getThreadContinue :: ThreadContinue -> IO Bool
    getThreadContinue (ThreadContinue ref) = readIORef ref

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

    -- | The type to store enough information for 'settingsOnException'.
    newtype StreamInfo = StreamInfo (IORef (Maybe (Stream,Request)))

    newStreamInfo :: IO StreamInfo
    newStreamInfo = StreamInfo <$> newIORef Nothing

    clearStreamInfo :: StreamInfo -> IO ()
    clearStreamInfo (StreamInfo ref) = writeIORef ref Nothing

    setStreamInfo :: StreamInfo -> Stream -> Request -> IO ()
    setStreamInfo (StreamInfo ref) strm req = writeIORef ref $ Just (strm,req)

    getStreamInfo :: StreamInfo -> IO (Maybe (Stream, Request))
    getStreamInfo (StreamInfo ref) = readIORef ref