#include "common.hh"

#include "buffer.hh"
#include "fcgi_protocol.hh"

#include <algorithm>
#include <limits>
#include <optional>
#include <utility>

namespace fcgi {

namespace {

constexpr const uint8_t FCGI_VERSION_1 = 1;

struct RawRecord {
  uint8_t version;
  uint8_t type;
  uint8_t request_id_b1;
  uint8_t request_id_b0;
  uint8_t content_length_b1;
  uint8_t content_length_b0;
  uint8_t padding_length;
  uint8_t reserved;
};

class RecordImpl : public Record {
public:
  RecordImpl(uint8_t type, uint16_t request_id, uint16_t content_length,
             uint8_t padding_length)
    : good_(true), type_(type), request_id_(request_id),
      content_length_(content_length), padding_length_(padding_length) {}

  RecordImpl()
    : good_(false), type_(RecordType::UnknownType), request_id_(0),
      content_length_(0), padding_length_(0) {}

  bool good() const override {
    return good_;
  }

  uint8_t type() const override {
    return type_;
  }

  uint16_t request_id() const override {
    return request_id_;
  }

  uint16_t content_length() const override {
    return content_length_;
  }

  uint8_t padding_length() const override {
    return padding_length_;
  }

private:
  bool const good_;
  uint8_t const type_;
  uint16_t const request_id_;
  uint16_t const content_length_;
  uint8_t const padding_length_;
};

class BeginRequestBodyImpl : public BeginRequestBody {
public:
  BeginRequestBodyImpl()
    : good_(false), role_(0), flags_(0) {}

  BeginRequestBodyImpl(uint16_t role, uint8_t flags)
    : good_(true), role_(role), flags_(flags) {
  }

  bool good() const override {
    return good_;
  }

  uint16_t role() const override {
    return role_;
  }

  uint8_t flags() const override {
    return flags_;
  }

  static std::unique_ptr<BeginRequestBody> parse(uint8_t const* data,
                                                 size_t len) {
    if (len != 8)
      return std::make_unique<BeginRequestBodyImpl>();
    return std::make_unique<BeginRequestBodyImpl>(
        static_cast<uint16_t>(data[0]) << 8 | data[1], data[2]);
  }

private:
  bool const good_;
  uint16_t const role_;
  uint8_t const flags_;
};

uint8_t calc_padding(uint16_t length) {
  auto extra = length % 8;
  return extra ? 8 - extra : 0;
}

class RecordBuilderImpl : public RecordBuilder {
public:
  RecordBuilderImpl(RecordType type, uint16_t request_id,
                    uint16_t content_length, uint8_t padding_length)
    : type_(type), request_id_(request_id), content_length_(content_length),
      padding_length_(padding_length) {}

  RecordBuilderImpl(RecordType type, uint16_t request_id, std::string body)
    : type_(type), request_id_(request_id), content_length_(body.size()),
      padding_length_(calc_padding(content_length_)),
      body_(std::move(body)) {
    assert(body_->size() <= std::numeric_limits<uint16_t>::max());
  }

  bool build(Buffer* dst) const override {
    size_t avail;
    size_t need = sizeof(RawRecord) +
      (body_ ? content_length_ + padding_length_ : 0);
    auto* ptr = dst->wbuf(need, avail);
    if (!build(ptr, avail))
      return false;
    dst->wcommit(need);
    return true;
  }

  bool build(char* ptr, size_t avail) const override {
    if (avail < sizeof(RawRecord) +
        (body_ ? content_length_ + padding_length_ : 0))
      return false;
    auto* raw = reinterpret_cast<RawRecord*>(ptr);
    raw->version = FCGI_VERSION_1;
    raw->type = type_;
    raw->request_id_b1 = request_id_ >> 8;
    raw->request_id_b0 = request_id_ & 0xff;
    raw->content_length_b1 = content_length_ >> 8;
    raw->content_length_b0 = content_length_ & 0xff;
    raw->padding_length = padding_length_;
    if (body_) {
      std::copy_n(body_->data(), body_->size(), ptr + sizeof(RawRecord));
      std::fill_n(ptr + sizeof(RawRecord) + body_->size(), padding_length_,
                  '\0');
    }
    return true;
  }

  size_t size() const override {
    return sizeof(RawRecord) + content_length_ + padding_length_;
  }

  bool padding(Buffer* dst) const override {
    size_t avail;
    auto* ptr = dst->wbuf(padding_length_, avail);
    if (!padding(ptr, avail))
      return false;
    dst->wcommit(padding_length_);
    return true;
  }

  bool padding(char* ptr, size_t avail) const override {
    if (avail < padding_length_)
      return false;
    std::fill_n(ptr, padding_length_, '\0');
    return true;
  }

private:
  RecordType const type_;
  uint16_t const request_id_;
  uint16_t const content_length_;
  uint8_t const padding_length_;
  std::optional<std::string> const body_;
};

std::unique_ptr<Record> parse_raw(RawRecord const& raw) {
  if (raw.version != FCGI_VERSION_1)
    return std::make_unique<RecordImpl>();
  return std::make_unique<RecordImpl>(
      raw.type,
      static_cast<uint16_t>(raw.request_id_b1) << 8 | raw.request_id_b0,
      static_cast<uint16_t>(raw.content_length_b1) << 8 | raw.content_length_b0,
      raw.padding_length);
}

bool parse_pair(RecordStream* stream, RoBuffer* buf,
                std::pair<std::string, std::string>& pair) {
  size_t avail;
  auto* ptr = stream->rbuf(buf, 2, avail);
  if (avail < 2)
    return false;
  auto* u8 = reinterpret_cast<uint8_t const*>(ptr);
  if (u8[0] & 0x80) {
    size_t need = 5;
    ptr = stream->rbuf(buf, need, avail);
    if (avail < need)
      return false;
    u8 = reinterpret_cast<uint8_t const*>(ptr);
    auto name_len = static_cast<uint32_t>(u8[0] & 0x7f) << 24 |
      static_cast<uint32_t>(u8[1]) << 16 |
      static_cast<uint32_t>(u8[2]) << 8 |
      u8[3];
    if (u8[4] & 0x80) {
      need = 8 + name_len;
      ptr = stream->rbuf(buf, need, avail);
      if (avail < need)
        return false;
      u8 = reinterpret_cast<uint8_t const*>(ptr);
      auto value_len = static_cast<uint32_t>(u8[4] & 0x7f) << 24 |
        static_cast<uint32_t>(u8[5]) << 16 |
        static_cast<uint32_t>(u8[6]) << 8 |
        u8[7];
      need = 8 + name_len + value_len;
      ptr = stream->rbuf(buf, need, avail);
      if (avail < need)
        return false;
      pair.first.assign(ptr + 8, name_len);
      pair.second.assign(ptr + 8 + name_len, value_len);
      stream->rcommit(buf, need);
      return true;
    } else {
      auto value_len = u8[4];
      need = 5 + name_len + value_len;
      ptr = stream->rbuf(buf, need, avail);
      if (avail < need)
        return false;
      pair.first.assign(ptr + 5, name_len);
      pair.second.assign(ptr + 5 + name_len, value_len);
      stream->rcommit(buf, need);
      return true;
    }
  } else if (u8[1] & 0x80) {
    auto name_len = u8[0];
    size_t need = 5 + name_len;
    ptr = stream->rbuf(buf, need, avail);
    if (avail < need)
      return false;
    u8 = reinterpret_cast<uint8_t const*>(ptr);
    auto value_len = static_cast<uint32_t>(u8[1] & 0x7f) << 24 |
      static_cast<uint32_t>(u8[2]) << 16 |
      static_cast<uint32_t>(u8[3]) << 8 |
      u8[4];
    need = 5 + name_len + value_len;
    ptr = stream->rbuf(buf, need, avail);
    if (avail < need)
      return false;
    pair.first.assign(ptr + 5, name_len);
    pair.second.assign(ptr + 5 + name_len, value_len);
    stream->rcommit(buf, need);
    return true;
  } else {
    auto name_len = u8[0];
    auto value_len = u8[1];
    size_t need = 2 + name_len + value_len;
    ptr = stream->rbuf(buf, need, avail);
    if (avail < need)
      return false;
    pair.first.assign(ptr + 2, name_len);
    pair.second.assign(ptr + 2 + name_len, value_len);
    stream->rcommit(buf, need);
    return true;
  }
}

class RecordStreamImpl : public RecordStream {
public:
  RecordStreamImpl(Record const* record, bool ended)
    : type_(record->type()),
      padding_(record->padding_length()),
      left_(static_cast<size_t>(record->content_length()) + padding_),
      leftover_(Buffer::growing(64, 8192)),
      ended_(ended || record->content_length() == 0) {
    check_end_of_stream();
  }

  char const* rbuf(RoBuffer* buf, size_t want, size_t& avail) override {
    if (leftover_->empty()) {
      size_t content_avail = left_ - padding_;
      // Avoid want == content_avail as that might miss the padding.
      if (want < content_avail) {
        auto* ptr = buf->rbuf(want, avail);
        if (avail > content_avail)
          avail = content_avail;
        return ptr;
      }
      auto* ptr = buf->rbuf(left_, avail);
      if (avail < left_) {
        avail = 0;
        return nullptr;
      }
      Buffer::write(leftover_.get(), ptr, left_ - padding_);
      buf->rcommit(left_);
      left_ = 0;
      padding_ = 0;
    }
    auto* ptr = leftover_->rbuf(want, avail);
    if (left_ > 0 && want > avail) {
      size_t tmp_avail;
      size_t tmp_need = want - avail;
      size_t tmp_want = tmp_need;
      if (tmp_want >= left_ - padding_)
        tmp_want = left_;
      auto* tmp = buf->rbuf(tmp_want, tmp_avail);
      if (tmp_avail < tmp_want)
        return ptr;
      Buffer::write(leftover_.get(), tmp, tmp_need);
      buf->rcommit(tmp_want);
      left_ -= tmp_want;
      if (left_ == 0)
        padding_ = 0;
      ptr = leftover_->rbuf(want, avail);
    }
    return ptr;
  }

  void rcommit(RoBuffer* buf, size_t bytes) override {
    if (leftover_->empty()) {
      assert(bytes <= left_ - padding_);
      left_ -= bytes;
      assert(left_ >= padding_);
      if (left_ == padding_) {
        buf->rcommit(bytes + padding_);
        left_ = 0;
        padding_ = 0;
        check_end_of_stream();
      } else {
        buf->rcommit(bytes);
      }
    } else {
      leftover_->rcommit(bytes);
      check_end_of_stream();
    }
  }

  bool end_of_record() const override {
    return left_ == 0;
  }

  bool end_of_stream() const override {
    return end_of_stream_;
  }

  bool all_available() override {
    if (!ended_)
      return false;
    return left_ == 0;
  }

  void add(Record const* record) override {
    assert(left_ == 0);
    assert(!ended_);
    assert(record->type() == type_);

    if (record->content_length() == 0)
      ended_ = true;
    padding_ = record->padding_length();
    left_ = static_cast<size_t>(record->content_length()) + padding_;

    check_end_of_stream();
  }

private:
  void check_end_of_stream() {
    if (end_of_stream_)
      return;
    if (ended_ && left_ == 0 && leftover_->empty())
      end_of_stream_ = true;
  }

  uint8_t const type_;
  uint8_t padding_;
  size_t left_;

  std::unique_ptr<Buffer> leftover_;

  // True if zero length record was added().
  bool ended_;
  // True if ended_ is true and all content has been read.
  bool end_of_stream_{false};
};

class PairImpl : public Pair {
public:
  PairImpl(std::string name, std::string value)
    : good_(true), name_(std::move(name)), value_(std::move(value)) {}
  PairImpl()
    : good_(false) {}

  bool good() const override {
    return good_;
  }

  std::string const& name() const override {
    return name_;
  }

  std::string const& value() const override {
    return value_;
  }

  bool next(RecordStream* stream, RoBuffer* buf) override {
    if (stream->end_of_stream())
      return false;
    std::pair<std::string, std::string> tmp;
    if (parse_pair(stream, buf, tmp)) {
      if (good_) {
        name_ = std::move(tmp.first);
        value_ = std::move(tmp.second);
      }
      return true;
    }
    if (stream->all_available()) {
      good_ = false;
      name_.clear();
      value_.clear();
      return true;
    }
    return false;
  }

private:
  bool good_;
  std::string name_;
  std::string value_;
};

class PairBuilderImpl : public PairBuilder {
public:
  PairBuilderImpl() = default;

  void add(std::string name, std::string value) override {
    data_.emplace_back(std::move(name), std::move(value));
  }

  size_t size() const override {
    size_t count = 0;
    for (auto const& pair : data_) {
      count += str_need(pair.first.size());
      count += str_need(pair.second.size());
    }
    return count;
  }

  bool build(Buffer* buf) const override {
    auto need = size();
    size_t avail;
    auto* ptr = reinterpret_cast<uint8_t*>(buf->wbuf(need, avail));
    if (avail < need)
      return false;
    size_t offset = 0;
    for (auto const& pair : data_) {
      if (pair.first.size() < 128) {
        ptr[offset++] = pair.first.size();
      } else {
        writeu32(ptr + offset, pair.first.size());
        offset += 4;
      }
      if (pair.second.size() < 128) {
        ptr[offset++] = pair.second.size();
      } else {
        writeu32(ptr + offset, pair.second.size());
        offset += 4;
      }
      std::copy_n(pair.first.data(), pair.first.size(), ptr + offset);
      offset += pair.first.size();
      std::copy_n(pair.second.data(), pair.second.size(), ptr + offset);
      offset += pair.second.size();
    }
    assert(offset == need);
    buf->wcommit(offset);
    return true;
  }

private:
  static size_t str_need(size_t len) {
    return len < 128 ? 1 + len : 4 + len;
  }

  static void writeu32(uint8_t* dst, uint32_t value) {
    assert(value <= 0x7ffffffful);
    dst[0] = (value >> 24) | 0x80;
    dst[1] = (value >> 16) & 0xff;
    dst[2] = (value >> 8) & 0xff;
    dst[3] = value & 0xff;
  }

  std::vector<std::pair<std::string, std::string>> data_;
};

}  // namespace

std::unique_ptr<Record> Record::parse(RoBuffer* buffer) {
  static_assert(sizeof(RawRecord) == 8);
  size_t avail;
  auto* ptr = buffer->rbuf(sizeof(RawRecord), avail);
  if (avail < sizeof(RawRecord))
    return nullptr;
  auto ret = parse_raw(*reinterpret_cast<RawRecord const*>(ptr));
  buffer->rcommit(sizeof(RawRecord));
  return ret;
}

std::unique_ptr<BeginRequestBody> BeginRequestBody::parse(Record const* record,
                                                          RoBuffer* buffer) {
  if (record->type() != RecordType::BeginRequest)
    return std::make_unique<BeginRequestBodyImpl>();
  if (record->content_length() != 8)
    return std::make_unique<BeginRequestBodyImpl>();
  auto need = static_cast<size_t>(record->content_length()) +
    record->padding_length();
  size_t avail;
  auto* ptr = reinterpret_cast<uint8_t const*>(buffer->rbuf(need, avail));
  if (avail < need)
    return nullptr;
  auto ret = BeginRequestBodyImpl::parse(ptr, record->content_length());
  buffer->rcommit(need);
  return ret;
}

std::unique_ptr<RecordStream> RecordStream::create_stream(
    Record const* record) {
  return std::make_unique<RecordStreamImpl>(record, false);
}

std::unique_ptr<RecordStream> RecordStream::create_single(
    Record const* record) {
  return std::make_unique<RecordStreamImpl>(record, true);
}

std::unique_ptr<Pair> Pair::start(RecordStream* stream, RoBuffer* buf) {
  std::pair<std::string, std::string> tmp;
  if (stream->end_of_stream())
    return nullptr;
  if (parse_pair(stream, buf, tmp))
    return std::make_unique<PairImpl>(std::move(tmp.first),
                                      std::move(tmp.second));
  if (stream->all_available())
    return std::make_unique<PairImpl>();
  return nullptr;
}

std::unique_ptr<RecordBuilder> RecordBuilder::create(RecordType type,
                                                     uint16_t request_id,
                                                     uint16_t content_length,
                                                     int16_t padding_length) {
  if (padding_length < 0)
    padding_length = calc_padding(content_length);
  return std::make_unique<RecordBuilderImpl>(type, request_id, content_length,
                                             padding_length & 0xff);
}

std::unique_ptr<RecordBuilder> RecordBuilder::create(RecordType type,
                                                     uint16_t request_id,
                                                     std::string body) {
  return std::make_unique<RecordBuilderImpl>(type, request_id, std::move(body));
}

std::unique_ptr<RecordBuilder> RecordBuilder::create_unknown_type(
    uint8_t unknown_type) {
  std::string body(8, '\0');
  body[0] = unknown_type;
  return create(RecordType::UnknownType, 0, std::move(body));
}

std::unique_ptr<RecordBuilder> RecordBuilder::create_begin_request(
    uint16_t request_id, Role role, uint8_t flags) {
  std::string body(8, '\0');
  uint16_t tmp_role = role;
  body[0] = tmp_role >> 8;
  body[1] = tmp_role & 0xff;
  body[2] = flags;
  return create(RecordType::BeginRequest, request_id, std::move(body));
}

std::unique_ptr<RecordBuilder> RecordBuilder::create_end_request(
    uint16_t request_id, uint32_t app_status, ProtocolStatus protocol_status) {
  std::string body(8, '\0');
  body[0] = app_status >> 24;
  body[1] = (app_status >> 16) & 0xff;
  body[2] = (app_status >> 8) & 0xff;
  body[3] = app_status & 0xff;
  body[4] = protocol_status;
  return create(RecordType::EndRequest, request_id, std::move(body));
}

std::unique_ptr<PairBuilder> PairBuilder::create() {
  return std::make_unique<PairBuilderImpl>( );
}

}  // namespace fcgi