{-# OPTIONS_GHC -fno-warn-orphans #-}
{-# OPTIONS_GHC -fno-warn-type-defaults #-}
module Servant.Streaming.ServerSpec (spec) where

import           Control.Concurrent
import           Control.Exception            (bracket)
import           Control.Monad.Trans.Resource (ResourceT, getInternalState,
                                               runInternalState, runResourceT)
import           Data.ByteString              (ByteString)
import qualified Data.ByteString.Streaming    as BSS
import qualified Data.ByteString              as BS
import qualified Data.ByteString.Lazy         as BSL

import           Data.IORef
import           Data.String                  (fromString, IsString)
import           GHC.Stats
import qualified Network.HTTP.Media           as M
import           Network.HTTP.Types           (status200, status405, status406,
                                               status415)
import           Network.Wai.Handler.Warp
import           Pipes.HTTP                   (Request, RequestBody (..),
                                               Response, brRead,
                                               defaultManagerSettings,
                                               defaultRequest, httpLbs, method,
                                               newManager, path, port,
                                               requestBody, requestHeaders,
                                               responseBody, responseStatus,
                                               responseTimeout,
                                               responseTimeoutNone,
                                               withResponse)
import           Servant                      ((:<|>) ((:<|>)), (:>), JSON,
                                               MimeRender (..), PlainText, Post,
                                               Proxy (..), Server, serve)
import           Servant.Streaming.Server
import           Streaming
import qualified Streaming.Prelude            as S
import           System.Directory             (removeFile)
import           Test.Hspec


spec :: Spec
spec = do
  streamBodySpec
  streamResponseSpec

streamBodySpec :: Spec
streamBodySpec = describe "StreamBody instance" $ around withServer $ do

  it "streams the request body" $ \port' -> do
    let req = streamReq port' "length" (S.each ["h","i"])
    responseBody <$> makeRequest req `shouldReturn` "2"

  it "does not keep the request in memory" $ \port' -> do
    let req = streamReq port' "length"
            $ S.replicate megabyte
            $ BS.replicate 1000 97 -- 1000 MB total
    responseBody <$> makeRequest req
      `shouldReturn` fromString (show (1000 * megabyte :: Int))
    bytes <- max_live_bytes <$> getRTSStats
    bytes < 200 * megabyte `shouldBe` True

  it "passes as argument the content-type" $ \port' -> do
    let req = streamReq port' "contentType" (S.each ["h","i"])
    responseBody <$> makeRequest req `shouldReturn` "application/json"
    let req' = req { requestHeaders = [("Content-Type", "text/plain;charset=utf-8")] }
    responseBody <$> makeRequest req' `shouldReturn` "text/plain;charset=utf-8"

  it "responds with '415 - Unsupported Media Type' on wrong content type" $ \port' -> do
    let req' = streamReq port' "length" (S.each ["h","i"])
        req = req' { requestHeaders = [("Content-Type", "bla/bla")] }
    responseStatus <$> makeRequest req `shouldReturn` status415

streamResponseSpec :: Spec
streamResponseSpec = describe "StreamResponse instance" $ around withServer $ do

  it "streams the response body" $ \port' -> do
    let req = streamReq port' "echo" (S.each ["h","i"])
    responseBody <$> makeRequest req `shouldReturn` "hi"

  it "does not keep the response in memory" $ \port' -> do
    let req = streamReq port' "echo"
            $ S.replicate 100
            $ BS.replicate megabyte 97 -- 100 MB total
    _ <- makeRequestStreamResponse req (runResourceT . S.length)
    bytes <- max_live_bytes <$> getRTSStats
    bytes < 200 * megabyte `shouldBe` True

  it "sets the specified status code" $ \port' -> do
    let req = streamReq port' "length" (S.each ["h","i"])
    responseStatus <$> makeRequest req `shouldReturn` status200

  it "responds with '405 - Method Not Allowed' on wrong method" $ \port' -> do
    let req' = streamReq port' "echo" (S.each ["h","i"])
        req = req' { method = "GET" }
    responseStatus <$> makeRequest req `shouldReturn` status405

  it "responds with '406 - Not Acceptable' on wrong content type" $ \port' -> do
    let req' = streamReq port' "echo" (S.each ["h","i"])
        req = req' { requestHeaders = ("Accept", "bla/bla"):requestHeaders req' }
    responseStatus <$> makeRequest req `shouldReturn` status406

  it "handles resource deallocation correctly" $ \port' -> do
    let req' = streamReq port' "getfile" (S.each []) -- TODO: simplify
        contents :: IsString a => a
        contents = "foobar"
    x <- bracket
      (BS.writeFile "hello.txt" contents)
      (\_ -> removeFile "hello.txt")
      (\_ -> responseBody <$> makeRequest req')
    x `shouldBe` fromString contents


------------------------------------------------------------------------------
-- API

type API
  =    "length" :> StreamBody '[JSON] :> Post '[PlainText] Int
  :<|> "contentType" :> StreamBody '[JSON, PlainText] :> Post '[PlainText] M.MediaType
  :<|> "echo" :> StreamBody '[JSON] :> StreamResponsePost '[JSON]
  :<|> "getfile" :> StreamResponsePost '[PlainText]

api :: Proxy API
api = Proxy

server :: Server API
server = lengthH :<|> contentTypeH :<|> echoH :<|> getfileH
  where
    lengthH      (_contentType, stream')
      = liftIO . runResourceT $ S.sum_ $ S.subst (\x -> BS.length x :> ()) stream'
    contentTypeH (contentType, _stream')
      = return contentType
    echoH        (_contentType, stream')
      = return stream'
    getfileH = return $ BSS.toChunks (BSS.readFile "hello.txt")

withServer :: (Port -> IO ()) -> IO ()
withServer = withApplicationSettings settings (return $ serve api server)
  where
    settings = setTimeout 60 defaultSettings

------------------------------------------------------------------------------
-- Utils

makeRequest
  :: Request -> IO (Response BSL.ByteString)
makeRequest req = do
  manager <- newManager defaultManagerSettings
  httpLbs req manager

makeRequestStreamResponse
  :: Request -> (Stream (Of ByteString) (ResourceT IO) () -> IO r) -> IO r
makeRequestStreamResponse req responseAction = do
  manager <- newManager defaultManagerSettings
  withResponse req manager (responseAction . go . responseBody)
  where
    go :: MonadIO m => IO BS.ByteString -> Stream (Of ByteString) m ()
    go respReader = do
      n <- liftIO $ brRead respReader
      if BS.null n
      then return ()
      else do
        S.yield n
        go respReader

streamReq
  :: Port -> ByteString -> Stream (Of ByteString) (ResourceT IO) () -> Request
streamReq appPort urlPath requestStream = defaultRequest
  { port = appPort
  , requestHeaders = [("Content-Type", "application/json")]
  , method = "POST"
  , path = urlPath
  , requestBody = streamReqBody requestStream
  , responseTimeout = responseTimeoutNone
  }
  where
    streamReqBody body
      = RequestBodyStreamChunked body'
      where
        body' :: (IO BS.ByteString -> IO ()) -> IO ()
        body' write = void . forkIO . runResourceT $ do
          ref <- liftIO $ newIORef body
          is <- getInternalState
          let popper :: IO BS.ByteString
              popper = do
                rsrc <- readIORef ref
                mres <- runInternalState (S.uncons rsrc) is
                case mres of
                  Nothing -> return BS.empty
                  Just (bs, str)
                    | BS.null bs -> writeIORef ref str >> popper
                    | otherwise -> writeIORef ref str >> return bs
          liftIO $ write popper

instance Show a => MimeRender PlainText a where
  mimeRender _ = fromString . show

megabyte :: Num a => a
megabyte = 1000 ^ 2