#include "connection_registry.hpp"
#include "connection.hpp"
#include <atomic>
#include <cstdlib>
#include <cstring>
#include <unistd.h>

ConnectionRegistry::ConnectionRegistry() : connections_(nullptr), max_fds_(0) {
  // Get the process file descriptor limit
  struct rlimit rlim;
  if (getrlimit(RLIMIT_NOFILE, &rlim) == -1) {
    perror("getrlimit");
    std::abort();
  }
  max_fds_ = rlim.rlim_cur;

  // Calculate size rounded up to page boundary
  size_t array_size = max_fds_ * sizeof(Connection *);
  size_t page_size = getpagesize();
  size_t aligned_size = (array_size + page_size - 1) & ~(page_size - 1);

  // Allocate virtual address space using mmap
  // MAP_ANONYMOUS provides zero-initialized pages on-demand (lazy allocation)
  connections_ = static_cast<std::atomic<Connection *> *>(
      mmap(nullptr, aligned_size, PROT_READ | PROT_WRITE,
           MAP_PRIVATE | MAP_ANONYMOUS, -1, 0));

  if (connections_ == MAP_FAILED) {
    perror("mmap");
    std::abort();
  }

  // Store aligned size for munmap
  aligned_size_ = aligned_size;
}

ConnectionRegistry::~ConnectionRegistry() {
  if (connections_ != nullptr) {
    for (int fd = 0; fd < static_cast<int>(max_fds_); ++fd) {
      delete connections_[fd].load(std::memory_order_relaxed);
    }
    if (munmap(connections_, aligned_size_) == -1) {
      perror("munmap");
    }
  }
}

void ConnectionRegistry::store(int fd, std::unique_ptr<Connection> connection) {
  if (fd < 0 || static_cast<size_t>(fd) >= max_fds_) {
    std::abort();
  }
  // Release ownership from unique_ptr and store raw pointer
  connections_[fd].store(connection.release(), std::memory_order_release);
}

std::unique_ptr<Connection> ConnectionRegistry::remove(int fd) {
  if (fd < 0 || static_cast<size_t>(fd) >= max_fds_) {
    std::abort();
  }
  return std::unique_ptr<Connection>(
      connections_[fd].exchange(nullptr, std::memory_order_acquire));
}