#define DOCTEST_CONFIG_IMPLEMENT_WITH_MAIN
#include <doctest/doctest.h>

#include <cstring>
#include <stdexcept>
#include <sys/mman.h>
#include <sys/resource.h>
#include <unistd.h>

// Forward declare Connection for registry
class Connection;

// Simplified connection registry for testing (avoid linking issues)
class TestConnectionRegistry {
public:
  TestConnectionRegistry() : connections_(nullptr), max_fds_(0) {
    struct rlimit rlim;
    if (getrlimit(RLIMIT_NOFILE, &rlim) == -1) {
      throw std::runtime_error("Failed to get RLIMIT_NOFILE");
    }
    max_fds_ = rlim.rlim_cur;

    connections_ = static_cast<Connection **>(
        mmap(nullptr, max_fds_ * sizeof(Connection *), PROT_READ | PROT_WRITE,
             MAP_PRIVATE | MAP_ANONYMOUS, -1, 0));

    if (connections_ == MAP_FAILED) {
      throw std::runtime_error("Failed to mmap for connection registry");
    }

    memset(connections_, 0, max_fds_ * sizeof(Connection *));
  }

  ~TestConnectionRegistry() {
    if (connections_ != MAP_FAILED && connections_ != nullptr) {
      munmap(connections_, max_fds_ * sizeof(Connection *));
    }
  }

  void store(int fd, Connection *connection) {
    if (fd < 0 || static_cast<size_t>(fd) >= max_fds_) {
      return;
    }
    connections_[fd] = connection;
  }

  Connection *get(int fd) const {
    if (fd < 0 || static_cast<size_t>(fd) >= max_fds_) {
      return nullptr;
    }
    return connections_[fd];
  }

  Connection *remove(int fd) {
    if (fd < 0 || static_cast<size_t>(fd) >= max_fds_) {
      return nullptr;
    }

    Connection *conn = connections_[fd];
    connections_[fd] = nullptr;
    return conn;
  }

  size_t max_fds() const { return max_fds_; }

private:
  Connection **connections_;
  size_t max_fds_;
};

// Mock Connection class for testing
class MockConnection {
public:
  MockConnection(int id) : id_(id) {}
  int getId() const { return id_; }

private:
  int id_;
};

TEST_CASE("ConnectionRegistry basic functionality") {
  TestConnectionRegistry registry;

  SUBCASE("max_fds returns valid limit") {
    struct rlimit rlim;
    getrlimit(RLIMIT_NOFILE, &rlim);
    CHECK(registry.max_fds() == rlim.rlim_cur);
    CHECK(registry.max_fds() > 0);
  }

  SUBCASE("get returns nullptr for empty registry") {
    CHECK(registry.get(0) == nullptr);
    CHECK(registry.get(100) == nullptr);
    CHECK(registry.get(1000) == nullptr);
  }

  SUBCASE("get handles invalid file descriptors") {
    CHECK(registry.get(-1) == nullptr);
    CHECK(registry.get(static_cast<int>(registry.max_fds())) == nullptr);
  }
}

TEST_CASE("ConnectionRegistry store and retrieve") {
  TestConnectionRegistry registry;

  // Create some mock connections (using reinterpret_cast for testing)
  MockConnection mock1(1);
  MockConnection mock2(2);
  Connection *conn1 = reinterpret_cast<Connection *>(&mock1);
  Connection *conn2 = reinterpret_cast<Connection *>(&mock2);

  SUBCASE("store and get single connection") {
    registry.store(5, conn1);
    CHECK(registry.get(5) == conn1);

    // Other fds should still return nullptr
    CHECK(registry.get(4) == nullptr);
    CHECK(registry.get(6) == nullptr);
  }

  SUBCASE("store multiple connections") {
    registry.store(5, conn1);
    registry.store(10, conn2);

    CHECK(registry.get(5) == conn1);
    CHECK(registry.get(10) == conn2);
    CHECK(registry.get(7) == nullptr);
  }

  SUBCASE("overwrite existing connection") {
    registry.store(5, conn1);
    CHECK(registry.get(5) == conn1);

    registry.store(5, conn2);
    CHECK(registry.get(5) == conn2);
  }

  SUBCASE("store handles invalid file descriptors safely") {
    registry.store(-1, conn1); // Should not crash
    registry.store(static_cast<int>(registry.max_fds()),
                   conn1); // Should not crash

    CHECK(registry.get(-1) == nullptr);
    CHECK(registry.get(static_cast<int>(registry.max_fds())) == nullptr);
  }
}

TEST_CASE("ConnectionRegistry remove functionality") {
  TestConnectionRegistry registry;

  MockConnection mock1(1);
  MockConnection mock2(2);
  Connection *conn1 = reinterpret_cast<Connection *>(&mock1);
  Connection *conn2 = reinterpret_cast<Connection *>(&mock2);

  SUBCASE("remove existing connection") {
    registry.store(5, conn1);
    CHECK(registry.get(5) == conn1);

    Connection *removed = registry.remove(5);
    CHECK(removed == conn1);
    CHECK(registry.get(5) == nullptr);
  }

  SUBCASE("remove non-existing connection") {
    Connection *removed = registry.remove(5);
    CHECK(removed == nullptr);
  }

  SUBCASE("remove after remove returns nullptr") {
    registry.store(5, conn1);
    Connection *removed1 = registry.remove(5);
    Connection *removed2 = registry.remove(5);

    CHECK(removed1 == conn1);
    CHECK(removed2 == nullptr);
  }

  SUBCASE("remove handles invalid file descriptors") {
    CHECK(registry.remove(-1) == nullptr);
    CHECK(registry.remove(static_cast<int>(registry.max_fds())) == nullptr);
  }

  SUBCASE("remove doesn't affect other connections") {
    registry.store(5, conn1);
    registry.store(10, conn2);

    Connection *removed = registry.remove(5);
    CHECK(removed == conn1);
    CHECK(registry.get(5) == nullptr);
    CHECK(registry.get(10) == conn2); // Should remain unchanged
  }
}

TEST_CASE("ConnectionRegistry large file descriptor handling") {
  TestConnectionRegistry registry;

  MockConnection mock1(1);
  Connection *conn1 = reinterpret_cast<Connection *>(&mock1);

  // Test with a large but valid file descriptor
  int large_fd = static_cast<int>(registry.max_fds()) - 1;

  SUBCASE("large valid fd works") {
    registry.store(large_fd, conn1);
    CHECK(registry.get(large_fd) == conn1);

    Connection *removed = registry.remove(large_fd);
    CHECK(removed == conn1);
    CHECK(registry.get(large_fd) == nullptr);
  }
}

TEST_CASE("ConnectionRegistry critical ordering simulation") {
  TestConnectionRegistry registry;

  MockConnection mock1(1);
  Connection *conn1 = reinterpret_cast<Connection *>(&mock1);
  int fd = 5;

  SUBCASE("simulate proper cleanup ordering") {
    // Step 1: Store connection
    registry.store(fd, conn1);
    CHECK(registry.get(fd) == conn1);

    // Step 2: Remove from registry (critical ordering step 1)
    Connection *removed = registry.remove(fd);
    CHECK(removed == conn1);
    CHECK(registry.get(fd) == nullptr);

    // Steps 2 & 3 would be close(fd) and delete conn
    // but we can't test those with mock objects
  }

  SUBCASE("simulate fd reuse safety") {
    // Store connection
    registry.store(fd, conn1);

    // Remove from registry first (step 1)
    Connection *removed = registry.remove(fd);
    CHECK(removed == conn1);

    // Registry is now clear - safe for fd reuse
    CHECK(registry.get(fd) == nullptr);

    // New connection could safely use same fd
    MockConnection mock2(2);
    Connection *conn2 = reinterpret_cast<Connection *>(&mock2);
    registry.store(fd, conn2);
    CHECK(registry.get(fd) == conn2);
  }
}