{-# LANGUAGE EmptyCase #-}

module Servant.Server.Internal.ResponseRender where

import Data.ByteString (ByteString)
import Data.Kind (Type)
import Data.Typeable
import GHC.TypeLits
import qualified Data.ByteString.Lazy as BSL
import qualified Network.Wai as Wai
import Network.HTTP.Types (Status, hContentType)
import Data.SOP
import qualified Servant.Types.SourceT                      as S
import qualified Data.ByteString.Builder                    as BB
import qualified Data.Sequence as Seq 

import Servant.API.ContentTypes (AcceptHeader (..), AllMimeRender, MimeRender, Accept, allMimeRender, mimeRender, contentType)
import Servant.API.MultiVerb
import Servant.API.Status
import Servant.API.Stream (SourceIO)
import Servant.API.UVerb.Union
import Servant.Types.Internal.Response
import qualified Network.HTTP.Media as M
import Data.Foldable (toList)
import Data.Sequence ((<|))

class (Typeable a) => IsWaiBody a where
  responseToWai :: InternalResponse a -> Wai.Response

instance IsWaiBody BSL.ByteString where
  responseToWai :: InternalResponse ByteString -> Response
responseToWai InternalResponse ByteString
r =
    Status -> ResponseHeaders -> ByteString -> Response
Wai.responseLBS
      (InternalResponse ByteString -> Status
forall a. InternalResponse a -> Status
statusCode InternalResponse ByteString
r)
      (Seq Header -> ResponseHeaders
forall a. Seq a -> [a]
forall (t :: Type -> Type) a. Foldable t => t a -> [a]
toList (InternalResponse ByteString -> Seq Header
forall a. InternalResponse a -> Seq Header
headers InternalResponse ByteString
r))
      (InternalResponse ByteString -> ByteString
forall a. InternalResponse a -> a
responseBody InternalResponse ByteString
r)

instance IsWaiBody () where
  responseToWai :: InternalResponse () -> Response
responseToWai InternalResponse ()
r =
    Status -> ResponseHeaders -> ByteString -> Response
Wai.responseLBS
      (InternalResponse () -> Status
forall a. InternalResponse a -> Status
statusCode InternalResponse ()
r)
      (Seq Header -> ResponseHeaders
forall a. Seq a -> [a]
forall (t :: Type -> Type) a. Foldable t => t a -> [a]
toList (InternalResponse () -> Seq Header
forall a. InternalResponse a -> Seq Header
headers InternalResponse ()
r))
      ByteString
forall a. Monoid a => a
mempty

instance IsWaiBody (SourceIO ByteString) where
  responseToWai :: InternalResponse (SourceIO ByteString) -> Response
responseToWai InternalResponse (SourceIO ByteString)
r =
    Status -> ResponseHeaders -> StreamingBody -> Response
Wai.responseStream
      (InternalResponse (SourceIO ByteString) -> Status
forall a. InternalResponse a -> Status
statusCode InternalResponse (SourceIO ByteString)
r)
      (Seq Header -> ResponseHeaders
forall a. Seq a -> [a]
forall (t :: Type -> Type) a. Foldable t => t a -> [a]
toList (InternalResponse (SourceIO ByteString) -> Seq Header
forall a. InternalResponse a -> Seq Header
headers InternalResponse (SourceIO ByteString)
r))
      (StreamingBody -> Response) -> StreamingBody -> Response
forall a b. (a -> b) -> a -> b
$ \Builder -> IO ()
output IO ()
flush -> do
        (String -> IO ())
-> (ByteString -> IO ()) -> SourceIO ByteString -> IO ()
forall (m :: Type -> Type) a.
Monad m =>
(String -> m ()) -> (a -> m ()) -> SourceT m a -> m ()
S.foreach
          (IO () -> String -> IO ()
forall a b. a -> b -> a
const (() -> IO ()
forall a. a -> IO a
forall (f :: Type -> Type) a. Applicative f => a -> f a
pure ()))
          (\ByteString
chunk -> Builder -> IO ()
output (ByteString -> Builder
BB.byteString ByteString
chunk) IO () -> IO () -> IO ()
forall a b. IO a -> IO b -> IO b
forall (f :: Type -> Type) a b. Applicative f => f a -> f b -> f b
*> IO ()
flush)
          (InternalResponse (SourceIO ByteString) -> SourceIO ByteString
forall a. InternalResponse a -> a
responseBody InternalResponse (SourceIO ByteString)
r)

data SomeResponse = forall a. (IsWaiBody a) => SomeResponse (InternalResponse a)

class ResponseListRender cs as where
  responseListRender
    :: AcceptHeader
    -> Union (ResponseTypes as)
    -> Maybe SomeResponse
  responseListStatuses :: [Status]

instance ResponseListRender cs '[] where
  responseListRender :: AcceptHeader -> Union (ResponseTypes '[]) -> Maybe SomeResponse
responseListRender AcceptHeader
_ Union (ResponseTypes '[])
x = case Union (ResponseTypes '[])
x of {}
  responseListStatuses :: [Status]
responseListStatuses = []

class (IsWaiBody (ResponseBody a)) => ResponseRender cs a where
  type ResponseStatus a :: Nat
  type ResponseBody a :: Type
  responseRender
    :: AcceptHeader
    -> ResponseType a
    -> Maybe (InternalResponse (ResponseBody a))

instance
  ( ResponseRender cs a,
    ResponseListRender cs as,
    KnownStatus (ResponseStatus a)
  ) =>
  ResponseListRender cs (a ': as)
  where
  responseListRender :: AcceptHeader
-> Union (ResponseTypes (a : as)) -> Maybe SomeResponse
responseListRender AcceptHeader
acc (Z (I x
x)) = (InternalResponse (ResponseBody a) -> SomeResponse)
-> Maybe (InternalResponse (ResponseBody a)) -> Maybe SomeResponse
forall a b. (a -> b) -> Maybe a -> Maybe b
forall (f :: Type -> Type) a b. Functor f => (a -> b) -> f a -> f b
fmap InternalResponse (ResponseBody a) -> SomeResponse
forall a. IsWaiBody a => InternalResponse a -> SomeResponse
SomeResponse (forall (cs :: k) a.
ResponseRender cs a =>
AcceptHeader
-> ResponseType a -> Maybe (InternalResponse (ResponseBody a))
forall {k} (cs :: k) a.
ResponseRender cs a =>
AcceptHeader
-> ResponseType a -> Maybe (InternalResponse (ResponseBody a))
responseRender @cs @a AcceptHeader
acc x
ResponseType a
x)
  responseListRender AcceptHeader
acc (S NS I xs
x) = forall (cs :: k) (as :: [Type]).
ResponseListRender cs as =>
AcceptHeader -> Union (ResponseTypes as) -> Maybe SomeResponse
forall {k} (cs :: k) (as :: [Type]).
ResponseListRender cs as =>
AcceptHeader -> Union (ResponseTypes as) -> Maybe SomeResponse
responseListRender @cs @as AcceptHeader
acc NS I xs
Union (ResponseTypes as)
x

  responseListStatuses :: [Status]
responseListStatuses = Proxy (ResponseStatus a) -> Status
forall (n :: Nat) (proxy :: Nat -> Type).
KnownStatus n =>
proxy n -> Status
forall (proxy :: Nat -> Type). proxy (ResponseStatus a) -> Status
statusVal (forall (t :: Nat). Proxy t
forall {k} (t :: k). Proxy t
Proxy @(ResponseStatus a)) Status -> [Status] -> [Status]
forall a. a -> [a] -> [a]
: forall (cs :: k) (as :: [Type]).
ResponseListRender cs as =>
[Status]
forall {k} (cs :: k) (as :: [Type]).
ResponseListRender cs as =>
[Status]
responseListStatuses @cs @as

instance
  ( AsHeaders xs (ResponseType r) a,
    ServantHeaders hs xs,
    ResponseRender cs r
  ) =>
  ResponseRender cs (WithHeaders hs a r)
  where
  type ResponseStatus (WithHeaders hs a r) = ResponseStatus r
  type ResponseBody (WithHeaders hs a r) = ResponseBody r

  responseRender :: AcceptHeader
-> ResponseType (WithHeaders hs a r)
-> Maybe (InternalResponse (ResponseBody (WithHeaders hs a r)))
responseRender AcceptHeader
acc ResponseType (WithHeaders hs a r)
x = InternalResponse (ResponseBody r)
-> InternalResponse (ResponseBody r)
addHeaders (InternalResponse (ResponseBody r)
 -> InternalResponse (ResponseBody r))
-> Maybe (InternalResponse (ResponseBody r))
-> Maybe (InternalResponse (ResponseBody r))
forall (f :: Type -> Type) a b. Functor f => (a -> b) -> f a -> f b
<$> forall (cs :: k) a.
ResponseRender cs a =>
AcceptHeader
-> ResponseType a -> Maybe (InternalResponse (ResponseBody a))
forall {k} (cs :: k) a.
ResponseRender cs a =>
AcceptHeader
-> ResponseType a -> Maybe (InternalResponse (ResponseBody a))
responseRender @cs @r AcceptHeader
acc ResponseType r
y
    where
      (NP I xs
hs, ResponseType r
y) = forall (headers :: [Type]) response returnType.
AsHeaders headers response returnType =>
returnType -> (NP I headers, response)
toHeaders @xs a
ResponseType (WithHeaders hs a r)
x
      addHeaders :: InternalResponse (ResponseBody r)
-> InternalResponse (ResponseBody r)
addHeaders InternalResponse (ResponseBody r)
r =
        InternalResponse (ResponseBody r)
r
          { headers = headers r <> Seq.fromList (constructHeaders @hs hs)
          }

instance
  ( KnownStatus s,
    MimeRender ct a
  ) =>
  ResponseRender cs (RespondAs (ct :: Type) s desc a)
  where
  type ResponseStatus (RespondAs ct s desc a) = s
  type ResponseBody (RespondAs ct s desc a) = BSL.ByteString

  responseRender :: AcceptHeader
-> ResponseType (RespondAs ct s desc a)
-> Maybe (InternalResponse (ResponseBody (RespondAs ct s desc a)))
responseRender AcceptHeader
_ ResponseType (RespondAs ct s desc a)
x =
    InternalResponse ByteString -> Maybe (InternalResponse ByteString)
InternalResponse ByteString
-> Maybe (InternalResponse (ResponseBody (RespondAs ct s desc a)))
forall a. a -> Maybe a
forall (f :: Type -> Type) a. Applicative f => a -> f a
pure (InternalResponse ByteString
 -> Maybe (InternalResponse (ResponseBody (RespondAs ct s desc a))))
-> (InternalResponse ByteString -> InternalResponse ByteString)
-> InternalResponse ByteString
-> Maybe (InternalResponse (ResponseBody (RespondAs ct s desc a)))
forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall {k} (ct :: k) a.
Accept ct =>
InternalResponse a -> InternalResponse a
forall ct a. Accept ct => InternalResponse a -> InternalResponse a
addContentType @ct (InternalResponse ByteString
 -> Maybe (InternalResponse (ResponseBody (RespondAs ct s desc a))))
-> InternalResponse ByteString
-> Maybe (InternalResponse (ResponseBody (RespondAs ct s desc a)))
forall a b. (a -> b) -> a -> b
$
      InternalResponse
        { statusCode :: Status
statusCode = Proxy s -> Status
forall (n :: Nat) (proxy :: Nat -> Type).
KnownStatus n =>
proxy n -> Status
forall (proxy :: Nat -> Type). proxy s -> Status
statusVal (forall (t :: Nat). Proxy t
forall {k} (t :: k). Proxy t
Proxy @s),
          responseBody :: ByteString
responseBody = Proxy ct -> a -> ByteString
forall {k} (ctype :: k) a.
MimeRender ctype a =>
Proxy ctype -> a -> ByteString
mimeRender (forall t. Proxy t
forall {k} (t :: k). Proxy t
Proxy @ct) a
ResponseType (RespondAs ct s desc a)
x,
          headers :: Seq Header
headers = Seq Header
forall a. Monoid a => a
mempty
        }

instance (KnownStatus s) => ResponseRender cs (RespondAs '() s desc ()) where
  type ResponseStatus (RespondAs '() s desc ()) = s
  type ResponseBody (RespondAs '() s desc ()) = ()

  responseRender :: AcceptHeader
-> ResponseType (RespondAs '() s desc ())
-> Maybe
     (InternalResponse (ResponseBody (RespondAs '() s desc ())))
responseRender AcceptHeader
_ ResponseType (RespondAs '() s desc ())
_ =
    InternalResponse (ResponseBody (RespondAs '() s desc ()))
-> Maybe
     (InternalResponse (ResponseBody (RespondAs '() s desc ())))
forall a. a -> Maybe a
forall (f :: Type -> Type) a. Applicative f => a -> f a
pure (InternalResponse (ResponseBody (RespondAs '() s desc ()))
 -> Maybe
      (InternalResponse (ResponseBody (RespondAs '() s desc ()))))
-> InternalResponse (ResponseBody (RespondAs '() s desc ()))
-> Maybe
     (InternalResponse (ResponseBody (RespondAs '() s desc ())))
forall a b. (a -> b) -> a -> b
$
      InternalResponse
        { statusCode :: Status
statusCode = Proxy s -> Status
forall (n :: Nat) (proxy :: Nat -> Type).
KnownStatus n =>
proxy n -> Status
forall (proxy :: Nat -> Type). proxy s -> Status
statusVal (forall (t :: Nat). Proxy t
forall {k} (t :: k). Proxy t
Proxy @s),
          responseBody :: ()
responseBody = (),
          headers :: Seq Header
headers = Seq Header
forall a. Monoid a => a
mempty
        }

instance
  (Accept ct, KnownStatus s)
  => ResponseRender cs (RespondStreaming s desc framing ct)
  where
  type ResponseStatus (RespondStreaming s desc framing ct) = s
  type ResponseBody (RespondStreaming s desc framing ct) = SourceIO ByteString 
  responseRender :: AcceptHeader
-> ResponseType (RespondStreaming s desc framing ct)
-> Maybe
     (InternalResponse
        (ResponseBody (RespondStreaming s desc framing ct)))
responseRender AcceptHeader
_ ResponseType (RespondStreaming s desc framing ct)
x =
    InternalResponse (SourceIO ByteString)
-> Maybe (InternalResponse (SourceIO ByteString))
InternalResponse (SourceIO ByteString)
-> Maybe
     (InternalResponse
        (ResponseBody (RespondStreaming s desc framing ct)))
forall a. a -> Maybe a
forall (f :: Type -> Type) a. Applicative f => a -> f a
pure (InternalResponse (SourceIO ByteString)
 -> Maybe
      (InternalResponse
         (ResponseBody (RespondStreaming s desc framing ct))))
-> (InternalResponse (SourceIO ByteString)
    -> InternalResponse (SourceIO ByteString))
-> InternalResponse (SourceIO ByteString)
-> Maybe
     (InternalResponse
        (ResponseBody (RespondStreaming s desc framing ct)))
forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall {k} (ct :: k) a.
Accept ct =>
InternalResponse a -> InternalResponse a
forall ct a. Accept ct => InternalResponse a -> InternalResponse a
addContentType @ct (InternalResponse (SourceIO ByteString)
 -> Maybe
      (InternalResponse
         (ResponseBody (RespondStreaming s desc framing ct))))
-> InternalResponse (SourceIO ByteString)
-> Maybe
     (InternalResponse
        (ResponseBody (RespondStreaming s desc framing ct)))
forall a b. (a -> b) -> a -> b
$
      InternalResponse
        { statusCode :: Status
statusCode = Proxy s -> Status
forall (n :: Nat) (proxy :: Nat -> Type).
KnownStatus n =>
proxy n -> Status
forall (proxy :: Nat -> Type). proxy s -> Status
statusVal (forall (t :: Nat). Proxy t
forall {k} (t :: k). Proxy t
Proxy @s),
          responseBody :: SourceIO ByteString
responseBody = SourceIO ByteString
ResponseType (RespondStreaming s desc framing ct)
x,
          headers :: Seq Header
headers = Seq Header
forall a. Monoid a => a
mempty
        }

instance
  (AllMimeRender cs a, KnownStatus s)
  => ResponseRender cs (Respond s desc a) where
  type ResponseStatus (Respond s desc a) = s
  type ResponseBody (Respond s desc a) = BSL.ByteString

  -- Note: here it seems like we are rendering for all possible content types,
  -- only to choose the correct one afterwards. However, render results besides the
  -- one picked by 'M.mapAcceptMedia' are not evaluated, and therefore nor are the
  -- corresponding rendering functions.
  responseRender :: AcceptHeader
-> ResponseType (Respond s desc a)
-> Maybe (InternalResponse (ResponseBody (Respond s desc a)))
responseRender (AcceptHeader ByteString
acc) ResponseType (Respond s desc a)
x =
    [(MediaType, InternalResponse ByteString)]
-> ByteString -> Maybe (InternalResponse ByteString)
forall b. [(MediaType, b)] -> ByteString -> Maybe b
M.mapAcceptMedia (((MediaType, ByteString)
 -> (MediaType, InternalResponse ByteString))
-> [(MediaType, ByteString)]
-> [(MediaType, InternalResponse ByteString)]
forall a b. (a -> b) -> [a] -> [b]
map ((MediaType
 -> ByteString -> (MediaType, InternalResponse ByteString))
-> (MediaType, ByteString)
-> (MediaType, InternalResponse ByteString)
forall a b c. (a -> b -> c) -> (a, b) -> c
uncurry MediaType -> ByteString -> (MediaType, InternalResponse ByteString)
mkRenderOutput) (Proxy cs -> a -> [(MediaType, ByteString)]
forall (list :: [Type]) a.
AllMimeRender list a =>
Proxy list -> a -> [(MediaType, ByteString)]
allMimeRender (forall (t :: [Type]). Proxy t
forall {k} (t :: k). Proxy t
Proxy @cs) a
ResponseType (Respond s desc a)
x)) ByteString
acc
    where
      mkRenderOutput :: M.MediaType -> BSL.ByteString -> (M.MediaType, InternalResponse BSL.ByteString)
      mkRenderOutput :: MediaType -> ByteString -> (MediaType, InternalResponse ByteString)
mkRenderOutput MediaType
c ByteString
body =
        (MediaType
c,) (InternalResponse ByteString
 -> (MediaType, InternalResponse ByteString))
-> (InternalResponse ByteString -> InternalResponse ByteString)
-> InternalResponse ByteString
-> (MediaType, InternalResponse ByteString)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. MediaType
-> InternalResponse ByteString -> InternalResponse ByteString
forall a. MediaType -> InternalResponse a -> InternalResponse a
addContentType' MediaType
c (InternalResponse ByteString
 -> (MediaType, InternalResponse ByteString))
-> InternalResponse ByteString
-> (MediaType, InternalResponse ByteString)
forall a b. (a -> b) -> a -> b
$
          InternalResponse
            { statusCode :: Status
statusCode = Proxy s -> Status
forall (n :: Nat) (proxy :: Nat -> Type).
KnownStatus n =>
proxy n -> Status
forall (proxy :: Nat -> Type). proxy s -> Status
statusVal (forall (t :: Nat). Proxy t
forall {k} (t :: k). Proxy t
Proxy @s),
              responseBody :: ByteString
responseBody = ByteString
body,
              headers :: Seq Header
headers = Seq Header
forall a. Monoid a => a
mempty
            }

addContentType :: forall ct a. (Accept ct) => InternalResponse a -> InternalResponse a
addContentType :: forall {k} (ct :: k) a.
Accept ct =>
InternalResponse a -> InternalResponse a
addContentType = MediaType -> InternalResponse a -> InternalResponse a
forall a. MediaType -> InternalResponse a -> InternalResponse a
addContentType' (Proxy ct -> MediaType
forall {k} (ctype :: k). Accept ctype => Proxy ctype -> MediaType
contentType (forall (t :: k). Proxy t
forall {k} (t :: k). Proxy t
Proxy @ct))

addContentType' :: M.MediaType -> InternalResponse a -> InternalResponse a
addContentType' :: forall a. MediaType -> InternalResponse a -> InternalResponse a
addContentType' MediaType
c InternalResponse a
r = InternalResponse a
r {headers = (hContentType, M.renderHeader c) <| headers r}

setEmptyBody :: SomeResponse -> SomeResponse
setEmptyBody :: SomeResponse -> SomeResponse
setEmptyBody (SomeResponse InternalResponse a
r) = InternalResponse ByteString -> SomeResponse
forall a. IsWaiBody a => InternalResponse a -> SomeResponse
SomeResponse (InternalResponse a -> InternalResponse ByteString
forall a. InternalResponse a -> InternalResponse ByteString
go InternalResponse a
r)
  where
    go :: InternalResponse a -> InternalResponse BSL.ByteString
    go :: forall a. InternalResponse a -> InternalResponse ByteString
go InternalResponse {a
Seq Header
Status
statusCode :: forall a. InternalResponse a -> Status
headers :: forall a. InternalResponse a -> Seq Header
responseBody :: forall a. InternalResponse a -> a
statusCode :: Status
headers :: Seq Header
responseBody :: a
..} = InternalResponse {responseBody :: ByteString
responseBody = ByteString
forall a. Monoid a => a
mempty, Seq Header
Status
statusCode :: Status
headers :: Seq Header
statusCode :: Status
headers :: Seq Header
..}

someResponseToWai :: SomeResponse -> Wai.Response
someResponseToWai :: SomeResponse -> Response
someResponseToWai (SomeResponse InternalResponse a
r) = InternalResponse a -> Response
forall a. IsWaiBody a => InternalResponse a -> Response
responseToWai InternalResponse a
r