#include "common.hh"

#include "config.hh"
#include "fcgi_protocol.hh"
#include "file_test.hh"
#include "io.hh"
#include "logger.hh"
#include "looper.hh"
#include "socket_test.hh"
#include "str_buffer.hh"
#include "task_runner.hh"
#include "transport_fastcgi.hh"

#include <gtest/gtest.h>
#include <map>

namespace {

class TransportFcgiTest : public SocketTest, public Transport::Handler {
public:
  ~TransportFcgiTest() override {
    if (!path_.empty()) {
      std::error_code err;
      std::filesystem::remove(path_, err);
    }
  }

  void SetUp() override {
    fd_ = FileTest::create_temp_file(std::string(), &path_);

    auto config = Config::create_empty();
    auto factory = create_transport_factory_fastcgi();
    handler_ = Transport::create_default_handler(logger_, this);
    transport_ = factory->create(logger_, looper(), runner_, logger_.get(),
                                 config.get(), handler_.get());
  }

  void TearDown() override {
    transport_.reset();
    handler_.reset();
    runner_.reset();
  }

  void write_file(std::string_view content) {
    ASSERT_TRUE(fd_);
    auto buffer = make_strbuffer(content);
    while (!buffer->empty()) {
      ASSERT_TRUE(io::drain(buffer.get(), fd_.get()));
    }
    ASSERT_TRUE(io::close(fd_.release()));
  }

  std::unique_ptr<Transport::Response> request(
      Transport* transport, Transport::Request const* request) override {
    if (request->method() == "GET") {
      if (request->path() == "/hello_world")
        return transport->create_ok_data("Hello World!");
      if (request->path() == "/file")
        return transport->create_ok_file(path_);
    }
    return transport->create_not_found();;
  }

  Transport* transport() { return transport_.get(); }

  Logger* logger() { return logger_.get(); }

private:
  std::shared_ptr<Logger> logger_ = Logger::create_null();
  std::shared_ptr<TaskRunner> runner_ = TaskRunner::create(looper());
  std::unique_ptr<Transport::Handler> handler_;
  std::unique_ptr<Transport> transport_;
  std::filesystem::path path_;
  unique_fd fd_;
};

void make_request(SocketTest::Client* cli, uint16_t request_id,
                  std::string const& path, uint8_t flags,
                  size_t step) {
  switch (step) {
  case 0: {
    auto builder = fcgi::RecordBuilder::create_begin_request(
        request_id, fcgi::Role::Responder, flags);
    cli->write([&] (Buffer* buf) {
      ASSERT_TRUE(builder->build(buf));
    });
    break;
  }
  case 1: {
    auto pair_builder = fcgi::PairBuilder::create();
    pair_builder->add("REQUEST_METHOD", "GET");
    pair_builder->add("REQUEST_URI", path);
    auto builder = fcgi::RecordBuilder::create(fcgi::RecordType::Params,
                                               request_id,
                                               pair_builder->size());
    cli->write([&] (Buffer* buf) {
      ASSERT_TRUE(builder->build(buf));
      ASSERT_TRUE(pair_builder->build(buf));
      ASSERT_TRUE(builder->padding(buf));
    });
    break;
  }
  case 2: {
    auto builder = fcgi::RecordBuilder::create(fcgi::RecordType::Params,
                                               request_id,
                                               std::string());
    cli->write([&] (Buffer* buf) {
      ASSERT_TRUE(builder->build(buf));
    });
    break;
  }
  case 3: {
    auto builder = fcgi::RecordBuilder::create(fcgi::RecordType::Stdin,
                                               request_id,
                                               std::string());
    cli->write([&] (Buffer* buf) {
      ASSERT_TRUE(builder->build(buf));
    });
    break;
  }
  default:
    FAIL();
  }
}

void make_request(SocketTest::Client* cli, uint16_t request_id,
                  std::string const& path, uint8_t flags) {
  for (size_t step = 0; step < 4; ++step) {
    make_request(cli, request_id, path, flags, step);
  }
}

struct Response {
  std::unique_ptr<fcgi::Record> record;
  std::unique_ptr<fcgi::RecordStream> stdout_stream;
  std::string stdout;
  std::optional<std::string> ended;

  void reset() {
    record.reset();
    stdout_stream.reset();
    stdout.clear();
    ended.reset();
  }
};

void read_response_content(SocketTest::Client* cli, uint16_t request_id,
                           Response* response, bool* need_more) {
  assert(!response->ended.has_value());
  while (true) {
    assert(response->record);
    ASSERT_EQ(request_id, response->record->request_id());
    if (response->record->type() == fcgi::RecordType::EndRequest) {
      if (cli->received().size() <
          response->record->content_length() +
          response->record->padding_length()) {
        ASSERT_FALSE(cli->closed());
        *need_more = true;
        return;
      }
      EXPECT_EQ(8, response->record->content_length());
      response->ended = cli->received().substr(
          0, response->record->content_length());
      cli->forget(response->record->content_length() +
                  response->record->padding_length());
      response->record.reset();
      *need_more = false;
      return;
    }
    ASSERT_EQ(fcgi::RecordType::Stdout, response->record->type());
    cli->read([&] (RoBuffer* buf) {
      size_t avail;
      auto* ptr = response->stdout_stream->rbuf(buf, 1, avail);
      response->stdout.append(ptr, avail);
      response->stdout_stream->rcommit(buf, avail);
    });
    if (response->stdout_stream->end_of_record()) {
      response->record.reset();
      *need_more = false;
      return;
    }
  }
}

void read_responses(SocketTest::Client* cli,
                    std::map<uint16_t, Response*>const& responses,
                    bool* need_more) {
  *need_more = false;
  auto it = responses.begin();
  while (it != responses.end()) {
    if (it->second->record) {
      read_response_content(cli, it->first, it->second, need_more);
      if (*need_more)
        return;
    } else {
      ++it;
    }
  }

  while (true) {
    bool all_ended = true;
    for (auto& pair : responses) {
      if (!pair.second->ended.has_value()) {
        all_ended = false;
        break;
      }
    }
    if (all_ended) {
      *need_more = false;
      return;
    }

    std::unique_ptr<fcgi::Record> record;
    cli->read([&] (RoBuffer* buf) {
      record = fcgi::Record::parse(buf);
    });
    if (!record) {
      ASSERT_FALSE(cli->closed());
      *need_more = true;
      return;
    }
    ASSERT_TRUE(record->good());
    it = responses.find(record->request_id());
    if (it == responses.end() || it->second->ended.has_value()) {
      FAIL();
      return;
    }
    auto* response = it->second;
    response->record = std::move(record);
    if (response->record->type() == fcgi::RecordType::Stdout) {
      if (!response->stdout_stream) {
        response->stdout_stream =
          fcgi::RecordStream::create_stream(response->record.get());
      } else {
        ASSERT_FALSE(response->stdout_stream->end_of_stream());
        response->stdout_stream->add(response->record.get());
      }
    }

    do {
      *need_more = false;
      read_response_content(cli, it->first, response, need_more);
      if (*need_more)
        return;
    } while (response->record);
  }
}

void read_response(SocketTest::Client* cli, uint16_t request_id,
                   Response* response, bool* need_more) {
  std::map<uint16_t, Response*> responses;
  responses[request_id] = response;
  read_responses(cli, responses, need_more);
}

struct GetValuesResponse {
  std::unique_ptr<fcgi::Record> record;
  std::unique_ptr<fcgi::RecordStream> stream;
  std::unique_ptr<fcgi::Pair> pair;
  std::vector<std::pair<std::string, std::string>> values;
};

void read_response(SocketTest::Client* cli, GetValuesResponse* response,
                   bool* need_more) {
  while (true) {
    if (!response->record) {
      cli->read([&] (RoBuffer* buf) {
        response->record = fcgi::Record::parse(buf);
      });
      if (!response->record) {
        ASSERT_FALSE(cli->closed());
        *need_more = true;
        return;
      }
      ASSERT_TRUE(response->record->good());
      ASSERT_EQ(fcgi::RecordType::GetValuesResult, response->record->type());
      if (!response->stream) {
        response->stream =
          fcgi::RecordStream::create_single(response->record.get());
      }
      if (!response->pair) {
        cli->read([&] (RoBuffer* buf) {
          response->pair = fcgi::Pair::start(response->stream.get(), buf);
        });
        if (!response->pair) {
          ASSERT_TRUE(response->stream->end_of_record());
          response->record.reset();
          *need_more = true;
          return;
        }
        ASSERT_TRUE(response->pair->good());
        response->values.emplace_back(response->pair->name(),
                                      response->pair->value());
      }
    }

    ASSERT_TRUE(response->pair);

    while (true) {
      bool next;
      cli->read([&] (RoBuffer* buf) {
        next = response->pair->next(response->stream.get(), buf);
      });
      if (next) {
        response->values.emplace_back(response->pair->name(),
                                      response->pair->value());
      } else {
        ASSERT_TRUE(response->pair->good());
        response->record.reset();
        if (response->stream->end_of_stream()) {
          *need_more = false;
          return;
        }
        *need_more = true;
        return;
      }
    }
  }
}

}  // namespace

TEST_F(TransportFcgiTest, sanity) {
  auto pair = create_pair();
  ASSERT_TRUE(pair.first && pair.second);
  transport()->add_client(std::move(pair.first));
  auto cli = create_client(std::move(pair.second));

  make_request(cli.get(), 1, "/hello_world", 0);

  Response response;

  while (true) {
    bool need_more;
    read_response(cli.get(), 1, &response, &need_more);
    if (need_more) {
      cli->wait(logger());
      continue;
    } else {
      break;
    }
  }

  ASSERT_TRUE(response.stdout_stream);
  EXPECT_TRUE(response.stdout_stream->end_of_stream());
  EXPECT_EQ("Status: 200\r\nContent-Length: 12\r\n\r\nHello World!",
            response.stdout);
  EXPECT_EQ(std::string_view("\0\0\0\0\0\0\0\0", 8), response.ended);
  if (!cli->closed())
    cli->wait(logger());
  EXPECT_TRUE(cli->closed());
}

TEST_F(TransportFcgiTest, reuse_conn) {
  write_file("foobar");

  auto pair = create_pair();
  ASSERT_TRUE(pair.first && pair.second);
  transport()->add_client(std::move(pair.first));
  auto cli = create_client(std::move(pair.second));

  make_request(cli.get(), 1, "/hello_world", fcgi::KeepConn);

  Response response;

  while (true) {
    bool need_more;
    read_response(cli.get(), 1, &response, &need_more);
    if (need_more) {
      cli->wait(logger());
      continue;
    } else {
      break;
    }
  }

  ASSERT_TRUE(response.stdout_stream);
  EXPECT_TRUE(response.stdout_stream->end_of_stream());
  EXPECT_EQ("Status: 200\r\nContent-Length: 12\r\n\r\nHello World!",
            response.stdout);
  EXPECT_EQ(std::string_view("\0\0\0\0\0\0\0\0", 8), response.ended);

  make_request(cli.get(), 2, "/file", fcgi::KeepConn);

  response.reset();

  while (true) {
    bool need_more;
    read_response(cli.get(), 2, &response, &need_more);
    if (need_more) {
      cli->wait(logger());
      continue;
    } else {
      break;
    }
  }

  ASSERT_TRUE(response.stdout_stream);
  EXPECT_TRUE(response.stdout_stream->end_of_stream());
  EXPECT_EQ("Status: 200\r\n\r\nfoobar", response.stdout);
  EXPECT_EQ(std::string_view("\0\0\0\0\0\0\0\0", 8), response.ended);
}

TEST_F(TransportFcgiTest, multiplexed) {
  write_file("foobar");

  auto pair = create_pair();
  ASSERT_TRUE(pair.first && pair.second);
  transport()->add_client(std::move(pair.first));
  auto cli = create_client(std::move(pair.second));

  for (size_t step = 0; step < 4; ++step) {
    make_request(cli.get(), 1, "/hello_world", fcgi::KeepConn, step);
    make_request(cli.get(), 2, "/file", fcgi::KeepConn, step);
  }

  Response response1;
  Response response2;
  std::map<uint16_t, Response*> responses;
  responses[1] = &response1;
  responses[2] = &response2;

  while (true) {
    bool need_more;
    read_responses(cli.get(), responses, &need_more);
    if (need_more) {
      cli->wait(logger());
      continue;
    } else {
      break;
    }
  }

  ASSERT_TRUE(response1.stdout_stream);
  EXPECT_TRUE(response1.stdout_stream->end_of_stream());
  EXPECT_EQ("Status: 200\r\nContent-Length: 12\r\n\r\nHello World!",
            response1.stdout);
  EXPECT_EQ(std::string_view("\0\0\0\0\0\0\0\0", 8), response1.ended);

  ASSERT_TRUE(response2.stdout_stream);
  EXPECT_TRUE(response2.stdout_stream->end_of_stream());
  EXPECT_EQ("Status: 200\r\n\r\nfoobar", response2.stdout);
  EXPECT_EQ(std::string_view("\0\0\0\0\0\0\0\0", 8), response2.ended);
}

TEST_F(TransportFcgiTest, get_values) {
  auto pair = create_pair();
  ASSERT_TRUE(pair.first && pair.second);
  transport()->add_client(std::move(pair.first));
  auto cli = create_client(std::move(pair.second));

  auto pair_builder = fcgi::PairBuilder::create();
  pair_builder->add("FCGI_MAX_CONNS", "");
  pair_builder->add("FCGI_MAX_REQS", "");
  pair_builder->add("foobar", "");
  auto builder = fcgi::RecordBuilder::create(fcgi::RecordType::GetValues,
                                             0,
                                             pair_builder->size());
  cli->write([&] (Buffer* buf) {
    ASSERT_TRUE(builder->build(buf));
    ASSERT_TRUE(pair_builder->build(buf));
    ASSERT_TRUE(builder->padding(buf));
  });

  GetValuesResponse response;

  while (true) {
    bool need_more;
    read_response(cli.get(), &response, &need_more);
    if (need_more) {
      cli->wait(logger());
      continue;
    } else {
      break;
    }
  }

  ASSERT_TRUE(response.stream);
  EXPECT_TRUE(response.stream->end_of_stream());
  EXPECT_EQ(2, response.values.size());
  for (auto const& value_pair : response.values) {
    if (value_pair.first == "FCGI_MAX_REQS") {
      EXPECT_EQ("20", value_pair.second);
    } else if (value_pair.first == "FCGI_MAX_CONNS") {
      EXPECT_EQ("10", value_pair.second);
    } else {
      EXPECT_EQ("FCGI_MAX_REQS", value_pair.first);
    }
  }
}