/*
 *  Copyright (c) 2015, Facebook, Inc.
 *  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. An additional grant
 *  of patent rights can be found in the PATENTS file in the same directory.
 *
 */
#pragma once

#include <folly/ScopeGuard.h>
#include <proxygen/httpserver/ResponseHandler.h>

namespace proxygen {

/**
 * Helps you make responses and send them on demand.
 *
 * NOTE: We don't do any correctness checks here, we depend on
 *       state machine in HTTPTransaction to tell us when an
 *       error occurs
 *
 * Three expected use cases are
 *
 * 1. Send all response at once. If this is an error
 *    response, most probably you also want 'closeConnection'.
 *
 * ResponseBuilder(handler)
 *    .status(200, "OK")
 *    .body(...)
 *    .sendWithEOM();
 *
 * 2. Sending back response in chunks.
 *
 * ResponseBuilder(handler)
 *    .status(200, "OK")
 *    .body(...)
 *    .send();  // Without `WithEOM` we make it chunked
 *
 * 1 or more time
 *
 * ResponseBuilder(handler)
 *    .body(...)
 *    .send();
 *
 * At last
 *
 * ResponseBuilder(handler)
 *    .body(...)
 *    .sendWithEOM();
 *
 * 3. Accept or reject Upgrade Requests
 *
 * ResponseBuilder(handler)
 *    .acceptUpgradeRequest() // send '200 OK' without EOM
 *
 * or
 *
 * ResponseBuilder(handler)
 *    .rejectUpgradeRequest() // send '400 Bad Request'
 *
 */
class ResponseBuilder {
 public:
  explicit ResponseBuilder(ResponseHandler* txn): txn_(txn) {
  }

  ResponseBuilder& status(uint16_t code, std::string message) {
    headers_ = folly::make_unique<HTTPMessage>();
    headers_->setHTTPVersion(1, 1);
    headers_->setStatusCode(code);
    headers_->setStatusMessage(message);
    return *this;
  }

  template <typename T>
  ResponseBuilder& header(const std::string& headerIn, const T& value) {
    CHECK(headers_) << "You need to call `status` before adding headers";
    headers_->getHeaders().add(headerIn, value);
    return *this;
  }

  template <typename T>
  ResponseBuilder& header(HTTPHeaderCode code, const T& value) {
    CHECK(headers_) << "You need to call `status` before adding headers";
    headers_->getHeaders().add(code, value);
    return *this;
  }

  ResponseBuilder& body(std::unique_ptr<folly::IOBuf> bodyIn) {
    if (bodyIn) {
      if (body_) {
        body_->prependChain(std::move(bodyIn));
      } else {
        body_ = std::move(bodyIn);
      }
    }

    return *this;
  }

  template <typename T>
  ResponseBuilder& body(T&& t) {
    return body(folly::IOBuf::maybeCopyBuffer(
        folly::to<std::string>(std::forward<T>(t))));
  }

  ResponseBuilder& closeConnection() {
    return header(HTTP_HEADER_CONNECTION, "close");
  }

  void sendWithEOM() {
    sendEOM_ = true;
    send();
  }

  void send() {
    // Once we send them, we don't want to send them again
    SCOPE_EXIT { headers_.reset(); };

    // By default, chunked
    bool chunked = true;

    // If we have complete response, we can use Content-Length and get done
    if (headers_ && sendEOM_) {
      chunked = false;
    }

    if (headers_) {
      // We don't need to add Content-Length or Encoding for 1xx responses
      if (headers_->getStatusCode() >= 200) {
        if (chunked) {
          headers_->setIsChunked(true);
        } else {
          const auto len = body_ ? body_->computeChainDataLength() : 0;
          headers_->getHeaders().add(
              HTTP_HEADER_CONTENT_LENGTH,
              folly::to<std::string>(len));
        }
      }

      txn_->sendHeaders(*headers_);
    }

    if (body_) {
      if (chunked) {
        txn_->sendChunkHeader(body_->computeChainDataLength());
        txn_->sendBody(std::move(body_));
        txn_->sendChunkTerminator();
      } else {
        txn_->sendBody(std::move(body_));
      }
    }

    if (sendEOM_) {
      txn_->sendEOM();
    }
  }

  enum class UpgradeType {
    CONNECT_REQUEST = 0,
    HTTP_UPGRADE,
  };

  void acceptUpgradeRequest(UpgradeType upgradeType,
                            const std::string upgradeProtocol = "") {
    headers_ = folly::make_unique<HTTPMessage>();
    if (upgradeType == UpgradeType::CONNECT_REQUEST) {
      headers_->constructDirectResponse({1, 1}, 200, "OK");
    } else {
      CHECK(!upgradeProtocol.empty());
      headers_->constructDirectResponse({1, 1}, 101, "Switching Protocols");
      headers_->getHeaders().add(HTTP_HEADER_UPGRADE, upgradeProtocol);
      headers_->getHeaders().add(HTTP_HEADER_CONNECTION, "Upgrade");
    }
    txn_->sendHeaders(*headers_);
  }

  void rejectUpgradeRequest() {
    headers_ = folly::make_unique<HTTPMessage>();
    headers_->constructDirectResponse({1, 1}, 400, "Bad Request");
    txn_->sendHeaders(*headers_);
    txn_->sendEOM();
  }

 private:
  ResponseHandler* const txn_{nullptr};

  std::unique_ptr<HTTPMessage> headers_;
  std::unique_ptr<folly::IOBuf> body_;

  // If true, sends EOM.
  bool sendEOM_{false};
};

}
