#include "connection.hpp"

#include <cerrno>
#include <climits>
#include <cstdio>
#include <cstdlib>
#include <sys/epoll.h>

#include "metric.hpp"
#include "server.hpp" // Need this for server reference

namespace {
// Thread-local metric instances
thread_local auto connections_total =
    metric::create_counter("weaseldb_connections_total",
                           "Total number of connections accepted")
        .create({});
thread_local auto connections_active =
    metric::create_gauge("weaseldb_connections_active",
                         "Number of currently active connections")
        .create({});
thread_local auto bytes_read =
    metric::create_counter("weaseldb_bytes_read_total",
                           "Total number of bytes read from clients")
        .create({});
thread_local auto bytes_written =
    metric::create_counter("weaseldb_bytes_written_total",
                           "Total number of bytes written to clients")
        .create({});
thread_local auto write_eagain_failures =
    metric::create_counter(
        "weaseldb_write_eagain_failures_total",
        "Total number of write operations that failed with EAGAIN")
        .create({});
} // namespace

// Static thread-local storage for iovec buffer
static thread_local std::vector<struct iovec> g_iovec_buffer{IOV_MAX};
// Thread-local storage for arenas to be freed after unlocking
static thread_local std::vector<Arena> g_arenas_to_free;

Connection::Connection(struct sockaddr_storage addr, int fd, int64_t id,
                       size_t epoll_index, ConnectionHandler *handler,
                       WeakRef<Server> server)
    : fd_(fd), id_(id), epoll_index_(epoll_index), addr_(addr),
      handler_(handler), server_(std::move(server)) {
  auto server_ref = server_.lock();
  // This should only be called from a member of Server itself, so I should
  // hope it's alive.
  assert(server_ref);
  server_ref->active_connections_.fetch_add(1, std::memory_order_relaxed);

  // Increment connection metrics using thread-local instances
  connections_total.inc();
  connections_active.inc();

  assert(handler_);
  handler_->on_connection_established(*this);
}

Connection::~Connection() {
  handler_->on_connection_closed(*this);
  assert(fd_ < 0 && "Connection fd was not closed before ~Connection");
}

void Connection::close() {
  std::lock_guard lock{mutex_};
  auto server_ptr = server_.lock();
  // Should only be called from the io thread
  assert(server_ptr);
  server_ptr->active_connections_.fetch_sub(1, std::memory_order_relaxed);
  assert(fd_ >= 0);
  int e = ::close(fd_);
  if (e == -1 && errno != EINTR) {
    perror("close");
    std::abort();
  }
  // EINTR ignored - fd is guaranteed closed on Linux
  fd_ = -1;
  // Decrement active connections gauge
  connections_active.dec();
}

// May be called off the io thread!
void Connection::append_message(std::span<std::string_view> data_parts,
                                Arena arena, bool close_after_send) {
  // Calculate total bytes for this message. Don't need to hold the lock yet.
  size_t total_bytes = 0;
  for (const auto &part : data_parts) {
    total_bytes += part.size();
  }

  std::unique_lock lock(mutex_);

  // Check if queue was empty to determine if we need to enable EPOLLOUT
  bool was_empty = message_queue_.empty();

  // Add message to queue
  message_queue_.emplace_back(
      Message{std::move(arena), data_parts, close_after_send});
  outgoing_bytes_queued_ += total_bytes;

  // If queue was empty, we need to add EPOLLOUT interest.
  if (was_empty && fd_ >= 0) {
    auto server = server_.lock();
    if (server) {
      // Add EPOLLOUT interest - pipeline thread manages epoll
      struct epoll_event event;
      event.data.fd = fd_;
      event.events = EPOLLIN | EPOLLOUT;
      tsan_release();
      epoll_ctl(server->epoll_fds_[epoll_index_], EPOLL_CTL_MOD, fd_, &event);
    }
  }
}

int Connection::readBytes(char *buf, size_t buffer_size) {
  int r;
  for (;;) {
    r = read(fd_, buf, buffer_size);
    if (r == -1) {
      if (errno == EINTR) {
        continue;
      }
      if (errno == EAGAIN) {
        return 0;
      }
      perror("read");
      return -1;
    }
    if (r == 0) {
      return -1;
    }

    // Increment bytes read metric
    if (r > 0) {
      bytes_read.inc(r);
    }

    return r;
  }
}

uint32_t Connection::write_bytes() {
  ssize_t total_bytes_written = 0;

  uint32_t result = 0;

  while (true) {
    // Build iovec array while holding mutex using thread-local buffer
    int iov_count = 0;
    {
      std::lock_guard lock(mutex_);

      if (message_queue_.empty()) {
        break;
      }

      // Build iovec array up to IOV_MAX limit using thread-local vector
      assert(g_iovec_buffer.size() == IOV_MAX);
      struct iovec *iov = g_iovec_buffer.data();

      for (auto &message : message_queue_) {
        if (iov_count >= IOV_MAX)
          break;

        for (const auto &part : message.data_parts) {
          if (iov_count >= IOV_MAX)
            break;
          if (part.empty())
            continue;

          iov[iov_count] = {
              const_cast<void *>(static_cast<const void *>(part.data())),
              part.size()};
          iov_count++;
        }
      }

      if (iov_count == 0)
        break;
    } // Release mutex during I/O

    // Perform I/O without holding mutex
    ssize_t w;
    for (;;) {
      struct msghdr msg = {};
      msg.msg_iov = g_iovec_buffer.data();
      msg.msg_iovlen = iov_count;

      w = sendmsg(fd_, &msg, MSG_NOSIGNAL);
      if (w == -1) {
        if (errno == EINTR) {
          continue; // Standard practice: retry on signal interruption
        }
        if (errno == EAGAIN) {
          // Increment EAGAIN failure metric
          write_eagain_failures.inc();
          if (total_bytes_written > 0) {
            bytes_written.inc(total_bytes_written);
          }
          return result;
        }
        perror("sendmsg");
        result |= Error;
        return result;
      }
      break;
    }

    result |= Progress;

    assert(w > 0);
    total_bytes_written += w;

    // Handle partial writes by updating message data_parts
    {
      std::lock_guard lock(mutex_);
      outgoing_bytes_queued_ -= w;
      size_t bytes_remaining = static_cast<size_t>(w);

      while (bytes_remaining > 0 && !message_queue_.empty()) {
        auto &front_message = message_queue_.front();
        bool message_complete = true;

        for (auto &part : front_message.data_parts) {
          if (part.empty())
            continue;

          if (bytes_remaining >= part.size()) {
            // This part is completely written
            bytes_remaining -= part.size();
            part = std::string_view(); // Mark as consumed
          } else {
            // Partial write of this part
            part = std::string_view(part.data() + bytes_remaining,
                                    part.size() - bytes_remaining);
            bytes_remaining = 0;
            message_complete = false;
            break;
          }
        }

        if (message_complete) {
          if (front_message.close_after_send) {
            result |= Close;
          }
          // Move arena to thread-local vector for deferred cleanup
          g_arenas_to_free.emplace_back(std::move(front_message.arena));
          message_queue_.pop_front();
          if (result & Close) {
            break;
          }
        } else {
          break;
        }
      }
    }
  }

  // Check if queue is empty and remove EPOLLOUT interest
  {
    std::lock_guard lock(mutex_);
    if (message_queue_.empty()) {
      result |= Drained;
      auto server = server_.lock();
      if (server) {
        struct epoll_event event;
        event.data.fd = fd_;
        event.events = EPOLLIN; // Remove EPOLLOUT
        tsan_release();
        epoll_ctl(server->epoll_fds_[epoll_index_], EPOLL_CTL_MOD, fd_, &event);
      }
    }
  }

  // Increment bytes written metric
  if (total_bytes_written > 0) {
    bytes_written.inc(total_bytes_written);
  }

  // Clean up arenas after all mutex operations are complete
  // This avoids holding the connection mutex while free() potentially contends
  g_arenas_to_free.clear();

  return result;
}
