#include "config.hpp"
#include "connection.hpp"
#include "connection_handler.hpp"
#include "server.hpp"

#include <doctest/doctest.h>
#include <latch>
#include <string_view>
#include <thread>

struct EchoHandler : ConnectionHandler {
  Arena arena;
  std::span<std::string_view> reply;
  WeakRef<MessageSender> wconn;
  std::latch done{1};
  void on_data_arrived(std::string_view data, Connection &conn) override {
    reply = arena.allocate_span<std::string_view>(1);
    reply[0] = arena.copy_string(data);
    wconn = conn.get_weak_ref();
    CHECK(wconn.lock());
    done.count_down();
  }
};

TEST_CASE("Echo test") {
  EchoHandler handler;
  weaseldb::Config config;
  auto server = Server::create(config, handler, {});
  int fd = server->create_local_connection();

  auto runThread = std::thread{[&]() { server->run(); }};

  int w = write(fd, "hello", 5);
  REQUIRE(w == 5);

  handler.done.wait();
  if (auto conn = handler.wconn.lock()) {
    // Cast to Connection* to access append_bytes (not available on
    // MessageSender)
    auto *conn_ptr = static_cast<Connection *>(conn.get());
    conn_ptr->append_bytes(std::exchange(handler.reply, {}),
                           std::move(handler.arena), ConnectionShutdown::None);
  } else {
    REQUIRE(false);
  }

  char buf[6];
  buf[5] = 0;
  int r = read(fd, buf, 5);
  REQUIRE(r == 5);
  CHECK(std::string(buf) == "hello");

  close(fd);

  server->shutdown();
  runThread.join();
}

struct ShutdownTestHandler : ConnectionHandler {
  Arena arena;
  std::span<std::string_view> reply;
  WeakRef<MessageSender> wconn;
  std::latch received_data{1};
  std::latch connection_closed_latch{1};
  ConnectionShutdown shutdown_mode = ConnectionShutdown::None;
  std::atomic<bool> connection_closed{false};

  void on_data_arrived(std::string_view data, Connection &conn) override {
    reply = arena.allocate_span<std::string_view>(1);
    reply[0] = arena.copy_string(data);
    wconn = conn.get_weak_ref();
    received_data.count_down();
  }

  void on_connection_closed(Connection &) override {
    connection_closed = true;
    connection_closed_latch.count_down();
  }
};

TEST_CASE("Connection shutdown write-only mode") {
  ShutdownTestHandler handler;
  handler.shutdown_mode = ConnectionShutdown::WriteOnly;
  weaseldb::Config config;
  auto server = Server::create(config, handler, {});
  int fd = server->create_local_connection();

  auto runThread = std::thread{[&]() { server->run(); }};

  // Send data to trigger handler
  int w = write(fd, "test", 4);
  REQUIRE(w == 4);

  handler.received_data.wait();

  // Send response with write shutdown
  if (auto conn = handler.wconn.lock()) {
    auto *conn_ptr = static_cast<Connection *>(conn.get());
    conn_ptr->append_bytes(std::exchange(handler.reply, {}),
                           std::move(handler.arena),
                           ConnectionShutdown::WriteOnly);
  } else {
    REQUIRE(false);
  }

  // Read the response
  char buf[5];
  buf[4] = 0;
  int r = read(fd, buf, 4);
  REQUIRE(r == 4);
  CHECK(std::string(buf) == "test");

  // After write shutdown, we should get EOF when trying to read more
  char extra_buf[1];
  int eof_result = read(fd, extra_buf, 1);
  CHECK(eof_result == 0); // EOF indicates successful write shutdown

  // Connection should still be alive (not closed) after write shutdown
  // We can verify this by checking that we can still write to the socket
  int write_result = write(fd, "x", 1);
  CHECK(write_result == 1); // Should succeed - connection still alive
  CHECK(handler.connection_closed.load() ==
        false); // Connection should still be alive

  close(fd);
  server->shutdown();
  runThread.join();
}

TEST_CASE("Connection shutdown full mode") {
  ShutdownTestHandler handler;
  handler.shutdown_mode = ConnectionShutdown::Full;
  weaseldb::Config config;
  auto server = Server::create(config, handler, {});
  int fd = server->create_local_connection();

  auto runThread = std::thread{[&]() { server->run(); }};

  // Send data to trigger handler
  int w = write(fd, "test", 4);
  REQUIRE(w == 4);

  handler.received_data.wait();

  // Send response with full shutdown
  if (auto conn = handler.wconn.lock()) {
    auto *conn_ptr = static_cast<Connection *>(conn.get());
    conn_ptr->append_bytes(std::exchange(handler.reply, {}),
                           std::move(handler.arena), ConnectionShutdown::Full);
  } else {
    REQUIRE(false);
  }

  // Read the response - connection should close after this
  char buf[5];
  buf[4] = 0;
  int r = read(fd, buf, 4);
  REQUIRE(r == 4);
  CHECK(std::string(buf) == "test");

  // Connection should be closed by server (full shutdown)
  char extra_buf[1];
  int close_result = read(fd, extra_buf, 1);
  CHECK(close_result == 0); // EOF indicates connection was closed

  // Wait for connection closed callback to be called
  handler.connection_closed_latch.wait();
  CHECK(handler.connection_closed.load() == true);

  close(fd);
  server->shutdown();
  runThread.join();
}