#include "../src/thread_pipeline.hpp"
#include "config.hpp"
#include "connection.hpp"
#include "perfetto_categories.hpp"
#include "server.hpp"
#include <cstring>
#include <doctest/doctest.h>
#include <thread>

// Perfetto static storage for tests
PERFETTO_TRACK_EVENT_STATIC_STORAGE();

struct Message {
  std::unique_ptr<Connection> conn;
  std::string data;
  bool done;
};

struct EchoHandler : public ConnectionHandler {
private:
  StaticThreadPipeline<Message, WaitStrategy::WaitIfStageEmpty, 1> &pipeline;

public:
  explicit EchoHandler(
      StaticThreadPipeline<Message, WaitStrategy::WaitIfStageEmpty, 1>
          &pipeline)
      : pipeline(pipeline) {}

  void on_data_arrived(std::string_view data,
                       std::unique_ptr<Connection> &conn_ptr) override {
    assert(conn_ptr);
    auto guard = pipeline.push(1, true);
    for (auto &message : guard.batch) {
      message.conn = std::move(conn_ptr);
      message.data = data;
      message.done = false;
    }
  }
};

TEST_CASE(
    "Server correctly handles connection ownership transfer via pipeline") {
  weaseldb::Config config;
  config.server.io_threads = 1;
  config.server.epoll_instances = 1;

  StaticThreadPipeline<Message, WaitStrategy::WaitIfStageEmpty, 1> pipeline{10};
  EchoHandler handler{pipeline};
  auto echoThread = std::thread{[&]() {
    for (;;) {
      auto guard = pipeline.acquire<0, 0>();
      for (auto &message : guard.batch) {
        bool done = message.done;
        if (done) {
          return;
        }
        assert(message.conn);
        message.conn->append_message(message.data);
        Server::release_back_to_server(std::move(message.conn));
      }
    }
  }};

  // Create server with NO listen sockets (empty vector)
  auto server = Server::create(config, handler, {});

  std::thread server_thread([&server]() { server->run(); });

  // Create local connection
  int client_fd = server->create_local_connection();
  REQUIRE(client_fd > 0);

  // Write some test data
  const char *test_message = "Hello, World!";
  ssize_t bytes_written;
  do {
    bytes_written = write(client_fd, test_message, std::strlen(test_message));
  } while (bytes_written == -1 && errno == EINTR);
  REQUIRE(bytes_written == std::strlen(test_message));

  // Read the echoed response
  char buffer[1024] = {0};
  ssize_t bytes_read;
  do {
    bytes_read = read(client_fd, buffer, sizeof(buffer) - 1);
  } while (bytes_read == -1 && errno == EINTR);
  if (bytes_read == -1) {
    perror("read failed");
  }
  REQUIRE(bytes_read == std::strlen(test_message));

  // Verify we got back exactly what we sent
  CHECK(std::string(buffer, bytes_read) == std::string(test_message));

  // Cleanup
  int e = close(client_fd);
  if (e == -1 && errno != EINTR) {
    perror("close client_fd");
    std::abort();
  }
  server->shutdown();
  server_thread.join();
  {
    auto guard = pipeline.push(1, true);
    for (auto &message : guard.batch) {
      message.done = true;
    }
  }
  echoThread.join();
}
