{-
  Copyright (c) Meta Platforms, Inc. and affiliates.
  All rights reserved.

  This source code is licensed under the BSD-style license found in the
  LICENSE file in the root directory of this source tree.
-}

-- | Support for creating a Thrift-over-HTTP server

{-# LANGUAGE TypeApplications #-}
module Thrift.Server.HTTP (
    ServerOptions(..),
    Server(..),
    defaultOptions,
    withBackgroundServer,
    withBackgroundServer',
    thriftApplication,
  ) where

import Data.ByteString (ByteString)
import Control.Concurrent
import Control.Concurrent.Async
import Control.Exception
import qualified Data.ByteString.Lazy as LBS
import Data.Proxy
import Data.String
import Data.Streaming.Network (bindPortTCP, bindRandomPortTCP)
import Network.HTTP.Types
import Network.Socket (close)
import Network.Wai.Handler.Warp
import Network.Wai

import Thrift.Processor hiding (Header)
import qualified Thrift.Processor as Thrift
import Thrift.Protocol
import Thrift.Protocol.Binary
import Thrift.Protocol.Compact
import Thrift.Protocol.JSON

-- TODO:
--  - one-way requests are currently treated as 2-way. Can we do any
--    better?

-- | Options for creating a Thrift-over-HTTP service.
data ServerOptions = ServerOptions
  { desiredPort :: Maybe Int
     -- ^ If 'Nothing', creates the server on a random free port,
     -- passing the actual port number as 'serverPort'.
  , numWorkerThreads :: Maybe Int
     -- ^ Currently ignored, provided for compatibility with CppServer
  , warpSettings :: Settings
  }

-- | Default options for creating a Thrift-over-HTTP service.
defaultOptions :: ServerOptions
defaultOptions = ServerOptions
  { desiredPort = Nothing
  , numWorkerThreads = Nothing
  , warpSettings = setHost (fromString "!6") defaultSettings
      -- IPv6 only by default
  }

-- | A running HTTP server.
data Server = Server
  { serverPort :: Int
      -- The actual port number, which might be useful if
      -- 'desiredPort' was 'Nothing'.
  , serverAsync :: Async ()
      -- The 'Async' running the HTTP server
  }

-- | Create an HTTP server for a Thrift service from the given 'ServerOptions'.
-- This is a simple wrapper around Warp's 'runSettings' that optionally creates
-- the server on a random port, and also wait for the server to start before
-- invoking the given action. Shuts down the server when the action returns.
withBackgroundServer
  :: forall s a . (Processor s)
  => (forall r . s r -> IO r) -- ^ handler to use
  -> ServerOptions
  -> (Server -> IO a)  -- ^ action to run while the server is up
  -> IO a
withBackgroundServer handler = withBackgroundServer' handler (\_ _ -> [])

-- | Create an HTTP server for a Thrift service from the given 'ServerOptions'.
-- This is a simple wrapper around Warp's 'runSettings' that optionally creates
-- the server on a random port, and also wait for the server to start before
-- invoking the given action. Shuts down the server when the action returns.
withBackgroundServer'
  :: forall s a . (Processor s)
  => (forall r . s r -> IO r) -- ^ handler to use
  -> (forall r . s r -> Either SomeException r -> Thrift.Header)
  -> ServerOptions
  -> (Server -> IO a)  -- ^ action to run while the server is up
  -> IO a
withBackgroundServer' handler postProcess ServerOptions{..} action = do
  ready <- newEmptyMVar
  let
    host = getHost warpSettings
    application = thriftApplication handler postProcess

    settings =
      maybe id setPort desiredPort $
      setBeforeMainLoop (putMVar ready ())
      warpSettings

    go port sock =
      withAsync (runSettingsSocket settings sock application) $ \a -> do
        takeMVar ready
        action (Server port a)

  case desiredPort of
    Nothing ->
      bracket (bindRandomPortTCP host) (close . snd) $ \(port,sock) ->
        go port sock
    Just port ->
      bracket (bindPortTCP port host) close (go port)

-- | Make a WAI 'Application' for a Thrift service. Use this with a
-- transport layer such as Warp to make a complete server, or call
-- 'withBackgroundServer' to do it all.
thriftApplication
  :: forall s . (Processor s)
  => (forall r . s r -> IO r) -- ^ handler to use
  -> (forall r . s r -> Either SomeException r -> Thrift.Header)
  -> Application
thriftApplication handler postProcess req respond = do
  body <- strictRequestBody req
  withProto (requestHeaders req) $ \proto contentType -> do
    (res, _maybeEx, _headers) <-
      process proto 0 handler postProcess (LBS.toStrict body)
    respond $ responseLBS
      status200
      [(hContentType, contentType)]
      (LBS.fromStrict res)
  where
  withProto
   :: [Header]
   -> (forall p . Protocol p => Proxy p -> ByteString -> IO b)
   -> IO b
  withProto hdrs f =
    case [ t | (header,t) <- hdrs, header == hContentType ] of
      t@"application/x-thrift-binary" : _ -> f (Proxy @Binary) t
      t@"application/x-thrift-compact" : _ -> f (Proxy @Compact) t
      t@"application/x-thrift-json" : _ -> f (Proxy @JSON) t
      _ -> f (Proxy @Binary) "application/x-thrift-binary"