#include "test_framework.h"
#include "../src/types.h"
#include "../src/connection.h"
#include "../src/buffer.h"
#include "../src/monitor.h"
#include "../src/config.h"
#include <string.h>
#include <unistd.h>
#include <fcntl.h>
#include <sys/socket.h>
#include <netinet/in.h>
#include <netinet/tcp.h>
#include <arpa/inet.h>
#include <sys/epoll.h>
#include <time.h>
#include <errno.h>

extern connection_t connections[MAX_FDS];
extern int epoll_fd;

void test_connection_init_all(void) {
    TEST_SUITE_BEGIN("Connection Init All");

    connection_init_all();

    TEST_ASSERT_EQ(CONN_TYPE_UNUSED, connections[0].type, "Connection 0 is unused");
    TEST_ASSERT_EQ(CONN_TYPE_UNUSED, connections[100].type, "Connection 100 is unused");

    TEST_SUITE_END();
}

void test_connection_set_non_blocking(void) {
    TEST_SUITE_BEGIN("Connection Set Non-Blocking");

    int pipefd[2];
    int result = pipe(pipefd);
    TEST_ASSERT_EQ(0, result, "Pipe created");

    result = connection_set_non_blocking(pipefd[0]);
    TEST_ASSERT_EQ(0, result, "Set non-blocking returns 0");

    int flags = fcntl(pipefd[0], F_GETFL, 0);
    TEST_ASSERT((flags & O_NONBLOCK) != 0, "O_NONBLOCK flag is set");

    close(pipefd[0]);
    close(pipefd[1]);

    result = connection_set_non_blocking(-1);
    TEST_ASSERT_EQ(-1, result, "Invalid fd returns -1");

    TEST_SUITE_END();
}

void test_connection_close_unused(void) {
    TEST_SUITE_BEGIN("Connection Close Unused");

    connection_init_all();

    connection_close(-1);
    TEST_ASSERT(1, "Close invalid fd doesn't crash");

    connection_close(MAX_FDS + 100);
    TEST_ASSERT(1, "Close out of range fd doesn't crash");

    TEST_SUITE_END();
}

void test_connection_do_read_null(void) {
    TEST_SUITE_BEGIN("Connection Do Read NULL");

    int result = connection_do_read(NULL);
    TEST_ASSERT_EQ(-1, result, "NULL connection returns -1");

    TEST_SUITE_END();
}

void test_connection_do_write_null(void) {
    TEST_SUITE_BEGIN("Connection Do Write NULL");

    int result = connection_do_write(NULL);
    TEST_ASSERT_EQ(-1, result, "NULL connection returns -1");

    TEST_SUITE_END();
}

void test_connection_do_read_basic(void) {
    TEST_SUITE_BEGIN("Connection Do Read Basic");

    int sockfd[2];
    int result = socketpair(AF_UNIX, SOCK_STREAM, 0, sockfd);
    TEST_ASSERT_EQ(0, result, "Socket pair created");

    connection_set_non_blocking(sockfd[0]);
    connection_set_non_blocking(sockfd[1]);

    connection_t conn;
    memset(&conn, 0, sizeof(conn));
    conn.fd = sockfd[0];
    conn.type = CONN_TYPE_CLIENT;
    buffer_init(&conn.read_buf, 4096);

    const char *test_data = "Hello, World!";
    ssize_t written = write(sockfd[1], test_data, strlen(test_data));
    TEST_ASSERT(written > 0, "Data written to socket");

    int bytes = connection_do_read(&conn);
    TEST_ASSERT(bytes > 0, "Read returned positive bytes");
    TEST_ASSERT(buffer_available_read(&conn.read_buf) > 0, "Data in read buffer");

    buffer_free(&conn.read_buf);
    close(sockfd[0]);
    close(sockfd[1]);

    TEST_SUITE_END();
}

void test_connection_do_write_basic(void) {
    TEST_SUITE_BEGIN("Connection Do Write Basic");

    int sockfd[2];
    int result = socketpair(AF_UNIX, SOCK_STREAM, 0, sockfd);
    TEST_ASSERT_EQ(0, result, "Socket pair created");

    connection_set_non_blocking(sockfd[0]);
    connection_set_non_blocking(sockfd[1]);

    connection_t conn;
    memset(&conn, 0, sizeof(conn));
    conn.fd = sockfd[0];
    conn.type = CONN_TYPE_CLIENT;
    buffer_init(&conn.write_buf, 4096);

    const char *test_data = "Hello, World!";
    memcpy(conn.write_buf.data, test_data, strlen(test_data));
    conn.write_buf.tail = strlen(test_data);

    int bytes = connection_do_write(&conn);
    TEST_ASSERT(bytes > 0, "Write returned positive bytes");

    char recv_buf[100] = {0};
    int recv_bytes = read(sockfd[1], recv_buf, sizeof(recv_buf));
    TEST_ASSERT(recv_bytes > 0, "Data received on other end");

    buffer_free(&conn.write_buf);
    close(sockfd[0]);
    close(sockfd[1]);

    TEST_SUITE_END();
}

void test_connection_tcp_keepalive(void) {
    TEST_SUITE_BEGIN("Connection TCP Keepalive");

    int sockfd = socket(AF_INET, SOCK_STREAM, 0);
    TEST_ASSERT(sockfd >= 0, "Socket created");

    connection_set_tcp_keepalive(sockfd);

    int optval;
    socklen_t optlen = sizeof(optval);
    getsockopt(sockfd, SOL_SOCKET, SO_KEEPALIVE, &optval, &optlen);
    TEST_ASSERT_EQ(1, optval, "SO_KEEPALIVE is enabled");

    close(sockfd);

    TEST_SUITE_END();
}

void test_connection_buffer_operations(void) {
    TEST_SUITE_BEGIN("Connection Buffer Operations");

    connection_t conn;
    memset(&conn, 0, sizeof(conn));

    int result = buffer_init(&conn.read_buf, 1024);
    TEST_ASSERT_EQ(0, result, "Read buffer initialized");

    result = buffer_init(&conn.write_buf, 1024);
    TEST_ASSERT_EQ(0, result, "Write buffer initialized");

    const char *data = "Test data for buffer";
    memcpy(conn.write_buf.data, data, strlen(data));
    conn.write_buf.tail = strlen(data);

    TEST_ASSERT_EQ(strlen(data), buffer_available_read(&conn.write_buf), "Write buffer has data");

    buffer_consume(&conn.write_buf, 5);
    TEST_ASSERT_EQ(strlen(data) - 5, buffer_available_read(&conn.write_buf), "Data consumed from buffer");

    buffer_free(&conn.read_buf);
    buffer_free(&conn.write_buf);

    TEST_SUITE_END();
}

void test_connection_send_error_response(void) {
    TEST_SUITE_BEGIN("Connection Send Error Response");

    connection_init_all();
    int old_epoll = epoll_fd;
    epoll_fd = epoll_create1(0);
    TEST_ASSERT(epoll_fd >= 0, "Epoll created");

    int sockfd[2];
    int result = socketpair(AF_UNIX, SOCK_STREAM, 0, sockfd);
    TEST_ASSERT_EQ(0, result, "Socket pair created");

    connection_t *conn = &connections[sockfd[0]];
    memset(conn, 0, sizeof(connection_t));
    conn->fd = sockfd[0];
    conn->type = CONN_TYPE_CLIENT;
    conn->state = CLIENT_STATE_READING_HEADERS;
    buffer_init(&conn->read_buf, 4096);
    buffer_init(&conn->write_buf, 4096);

    struct epoll_event ev = { .events = EPOLLIN, .data.fd = sockfd[0] };
    epoll_ctl(epoll_fd, EPOLL_CTL_ADD, sockfd[0], &ev);

    connection_send_error_response(conn, 404, "Not Found", "Resource not found");

    TEST_ASSERT(conn->write_buf.tail > 0, "Error response written");
    TEST_ASSERT(strstr((char*)conn->write_buf.data, "404") != NULL, "404 in response");

    buffer_free(&conn->read_buf);
    buffer_free(&conn->write_buf);
    close(sockfd[0]);
    close(sockfd[1]);
    close(epoll_fd);
    epoll_fd = old_epoll;

    TEST_SUITE_END();
}

void test_connection_epoll_operations(void) {
    TEST_SUITE_BEGIN("Connection Epoll Operations");

    int old_epoll = epoll_fd;
    epoll_fd = epoll_create1(0);
    TEST_ASSERT(epoll_fd >= 0, "Epoll created");

    int sockfd[2];
    int result = socketpair(AF_UNIX, SOCK_STREAM, 0, sockfd);
    TEST_ASSERT_EQ(0, result, "Socket pair created");

    connection_add_to_epoll(sockfd[0], EPOLLIN);
    TEST_ASSERT(1, "Add to epoll doesn't crash");

    connection_modify_epoll(sockfd[0], EPOLLIN | EPOLLOUT);
    TEST_ASSERT(1, "Modify epoll doesn't crash");

    close(sockfd[0]);
    close(sockfd[1]);
    close(epoll_fd);
    epoll_fd = old_epoll;

    TEST_SUITE_END();
}

void test_connection_send_auth_required(void) {
    TEST_SUITE_BEGIN("Connection Send Auth Required");

    connection_init_all();
    int old_epoll = epoll_fd;
    epoll_fd = epoll_create1(0);
    TEST_ASSERT(epoll_fd >= 0, "Epoll created");

    int sockfd[2];
    int result = socketpair(AF_UNIX, SOCK_STREAM, 0, sockfd);
    TEST_ASSERT_EQ(0, result, "Socket pair created");

    connection_t *conn = &connections[sockfd[0]];
    memset(conn, 0, sizeof(connection_t));
    conn->fd = sockfd[0];
    conn->type = CONN_TYPE_CLIENT;
    conn->state = CLIENT_STATE_READING_HEADERS;
    buffer_init(&conn->read_buf, 4096);
    buffer_init(&conn->write_buf, 4096);

    struct epoll_event ev = { .events = EPOLLIN, .data.fd = sockfd[0] };
    epoll_ctl(epoll_fd, EPOLL_CTL_ADD, sockfd[0], &ev);

    connection_send_auth_required(conn, "RProxy Dashboard");

    TEST_ASSERT(conn->write_buf.tail > 0, "Auth response written");
    TEST_ASSERT(strstr((char*)conn->write_buf.data, "401") != NULL, "401 status in response");
    TEST_ASSERT(strstr((char*)conn->write_buf.data, "WWW-Authenticate") != NULL, "WWW-Authenticate header");

    buffer_free(&conn->read_buf);
    buffer_free(&conn->write_buf);
    close(sockfd[0]);
    close(sockfd[1]);
    close(epoll_fd);
    epoll_fd = old_epoll;

    TEST_SUITE_END();
}

void test_connection_close_active(void) {
    TEST_SUITE_BEGIN("Connection Close Active");

    connection_init_all();
    int old_epoll = epoll_fd;
    epoll_fd = epoll_create1(0);
    TEST_ASSERT(epoll_fd >= 0, "Epoll created");

    int sockfd[2];
    int result = socketpair(AF_UNIX, SOCK_STREAM, 0, sockfd);
    TEST_ASSERT_EQ(0, result, "Socket pair created");

    connection_t *conn = &connections[sockfd[0]];
    memset(conn, 0, sizeof(connection_t));
    conn->fd = sockfd[0];
    conn->type = CONN_TYPE_CLIENT;
    conn->state = CLIENT_STATE_READING_HEADERS;
    buffer_init(&conn->read_buf, 4096);
    buffer_init(&conn->write_buf, 4096);

    struct epoll_event ev = { .events = EPOLLIN, .data.fd = sockfd[0] };
    epoll_ctl(epoll_fd, EPOLL_CTL_ADD, sockfd[0], &ev);

    connection_close(sockfd[0]);

    TEST_ASSERT_EQ(CONN_TYPE_UNUSED, conn->type, "Connection type reset to unused");

    close(sockfd[1]);
    close(epoll_fd);
    epoll_fd = old_epoll;

    TEST_SUITE_END();
}

void test_connection_cleanup_idle(void) {
    TEST_SUITE_BEGIN("Connection Cleanup Idle");

    connection_init_all();
    int old_epoll = epoll_fd;
    epoll_fd = epoll_create1(0);
    TEST_ASSERT(epoll_fd >= 0, "Epoll created");

    connection_cleanup_idle();
    TEST_ASSERT(1, "Cleanup idle doesn't crash on empty connections");

    close(epoll_fd);
    epoll_fd = old_epoll;

    TEST_SUITE_END();
}

void test_connection_do_read_eof(void) {
    TEST_SUITE_BEGIN("Connection Do Read EOF");

    int sockfd[2];
    int result = socketpair(AF_UNIX, SOCK_STREAM, 0, sockfd);
    TEST_ASSERT_EQ(0, result, "Socket pair created");

    connection_set_non_blocking(sockfd[0]);

    connection_t conn;
    memset(&conn, 0, sizeof(conn));
    conn.fd = sockfd[0];
    conn.type = CONN_TYPE_CLIENT;
    buffer_init(&conn.read_buf, 4096);

    close(sockfd[1]);

    int bytes = connection_do_read(&conn);
    TEST_ASSERT_EQ(0, bytes, "Read returns 0 on EOF");

    buffer_free(&conn.read_buf);
    close(sockfd[0]);

    TEST_SUITE_END();
}

void test_connection_do_write_full_buffer(void) {
    TEST_SUITE_BEGIN("Connection Do Write Full Buffer");

    int sockfd[2];
    int result = socketpair(AF_UNIX, SOCK_STREAM, 0, sockfd);
    TEST_ASSERT_EQ(0, result, "Socket pair created");

    connection_set_non_blocking(sockfd[0]);
    connection_set_non_blocking(sockfd[1]);

    int sndbuf = 4096;
    setsockopt(sockfd[0], SOL_SOCKET, SO_SNDBUF, &sndbuf, sizeof(sndbuf));

    connection_t conn;
    memset(&conn, 0, sizeof(conn));
    conn.fd = sockfd[0];
    conn.type = CONN_TYPE_CLIENT;
    buffer_init(&conn.write_buf, 4096);

    memset(conn.write_buf.data, 'A', 4096);
    conn.write_buf.tail = 4096;

    int bytes = connection_do_write(&conn);
    TEST_ASSERT(bytes > 0, "Write returns positive bytes");

    buffer_free(&conn.write_buf);
    close(sockfd[0]);
    close(sockfd[1]);

    TEST_SUITE_END();
}

void test_connection_multiple_error_responses(void) {
    TEST_SUITE_BEGIN("Connection Multiple Error Responses");

    connection_init_all();
    int old_epoll = epoll_fd;
    epoll_fd = epoll_create1(0);

    int sockfd[2];
    socketpair(AF_UNIX, SOCK_STREAM, 0, sockfd);

    connection_t *conn = &connections[sockfd[0]];
    memset(conn, 0, sizeof(connection_t));
    conn->fd = sockfd[0];
    conn->type = CONN_TYPE_CLIENT;
    buffer_init(&conn->read_buf, 4096);
    buffer_init(&conn->write_buf, 4096);

    struct epoll_event ev = { .events = EPOLLIN, .data.fd = sockfd[0] };
    epoll_ctl(epoll_fd, EPOLL_CTL_ADD, sockfd[0], &ev);

    connection_send_error_response(conn, 500, "Internal Server Error", "Server error");
    TEST_ASSERT(strstr((char*)conn->write_buf.data, "500") != NULL, "500 in response");

    buffer_free(&conn->read_buf);
    buffer_free(&conn->write_buf);
    close(sockfd[0]);
    close(sockfd[1]);
    close(epoll_fd);
    epoll_fd = old_epoll;

    TEST_SUITE_END();
}

void test_connection_do_read_eagain(void) {
    TEST_SUITE_BEGIN("Connection Do Read EAGAIN");

    int sockfd[2];
    int result = socketpair(AF_UNIX, SOCK_STREAM, 0, sockfd);
    TEST_ASSERT_EQ(0, result, "Socket pair created");

    connection_set_non_blocking(sockfd[0]);
    connection_set_non_blocking(sockfd[1]);

    connection_t conn;
    memset(&conn, 0, sizeof(conn));
    conn.fd = sockfd[0];
    conn.type = CONN_TYPE_CLIENT;
    buffer_init(&conn.read_buf, 4096);

    int bytes = connection_do_read(&conn);
    TEST_ASSERT(bytes <= 0, "Read returns <= 0 when no data");

    buffer_free(&conn.read_buf);
    close(sockfd[0]);
    close(sockfd[1]);

    TEST_SUITE_END();
}

void test_connection_optimize_socket(void) {
    TEST_SUITE_BEGIN("Connection Optimize Socket");

    int sockfd = socket(AF_INET, SOCK_STREAM, 0);
    TEST_ASSERT(sockfd >= 0, "Socket created");

    connection_optimize_socket(sockfd, 0);

    int optval;
    socklen_t optlen = sizeof(optval);
    getsockopt(sockfd, IPPROTO_TCP, TCP_NODELAY, &optval, &optlen);
    TEST_ASSERT_EQ(1, optval, "TCP_NODELAY is enabled for client");

    close(sockfd);

    sockfd = socket(AF_INET, SOCK_STREAM, 0);
    TEST_ASSERT(sockfd >= 0, "Second socket created");

    connection_optimize_socket(sockfd, 1);

    getsockopt(sockfd, IPPROTO_TCP, TCP_NODELAY, &optval, &optlen);
    TEST_ASSERT_EQ(1, optval, "TCP_NODELAY is enabled for upstream");

    close(sockfd);

    TEST_SUITE_END();
}

extern time_t cached_time;
void connection_update_cached_time(void);

void test_connection_cached_time(void) {
    TEST_SUITE_BEGIN("Connection Cached Time");

    time_t before = time(NULL);
    connection_update_cached_time();
    time_t after = time(NULL);

    TEST_ASSERT(cached_time >= before, "Cached time >= before");
    TEST_ASSERT(cached_time <= after, "Cached time <= after");

    TEST_SUITE_END();
}

void test_connection_epoll_event_caching(void) {
    TEST_SUITE_BEGIN("Connection Epoll Event Caching");

    connection_init_all();
    int old_epoll = epoll_fd;
    epoll_fd = epoll_create1(0);
    TEST_ASSERT(epoll_fd >= 0, "Epoll created");

    int sockfd[2];
    int result = socketpair(AF_UNIX, SOCK_STREAM, 0, sockfd);
    TEST_ASSERT_EQ(0, result, "Socket pair created");

    connection_add_to_epoll(sockfd[0], EPOLLIN);
    TEST_ASSERT_EQ((uint32_t)EPOLLIN, connections[sockfd[0]].epoll_events, "Events cached on add");

    connection_modify_epoll(sockfd[0], EPOLLIN | EPOLLOUT);
    TEST_ASSERT_EQ((uint32_t)(EPOLLIN | EPOLLOUT), connections[sockfd[0]].epoll_events, "Events cached on modify");

    connection_modify_epoll(sockfd[0], EPOLLIN | EPOLLOUT);
    TEST_ASSERT(1, "Redundant modify doesn't crash");

    close(sockfd[0]);
    close(sockfd[1]);
    close(epoll_fd);
    epoll_fd = old_epoll;

    TEST_SUITE_END();
}

void test_connection_splice_pipe_init(void) {
    TEST_SUITE_BEGIN("Connection Splice Pipe Init");

    connection_init_all();

    TEST_ASSERT_EQ(-1, connections[0].splice_pipe[0], "Splice pipe[0] initialized to -1");
    TEST_ASSERT_EQ(-1, connections[0].splice_pipe[1], "Splice pipe[1] initialized to -1");
    TEST_ASSERT_EQ(0, connections[0].can_splice, "can_splice initialized to 0");

    TEST_SUITE_END();
}

void test_connection_patch_buf_cleanup(void) {
    TEST_SUITE_BEGIN("Connection Patch Buffer Cleanup");

    connection_init_all();
    int old_epoll = epoll_fd;
    epoll_fd = epoll_create1(0);
    TEST_ASSERT(epoll_fd >= 0, "Epoll created");

    int sockfd[2];
    socketpair(AF_UNIX, SOCK_STREAM, 0, sockfd);

    connection_t *conn = &connections[sockfd[0]];
    memset(conn, 0, sizeof(connection_t));
    conn->fd = sockfd[0];
    conn->type = CONN_TYPE_CLIENT;
    buffer_init(&conn->read_buf, 4096);
    buffer_init(&conn->write_buf, 4096);

    conn->patch_buf = malloc(1024);
    conn->patch_buf_capacity = 1024;

    struct epoll_event ev = { .events = EPOLLIN, .data.fd = sockfd[0] };
    epoll_ctl(epoll_fd, EPOLL_CTL_ADD, sockfd[0], &ev);

    connection_close(sockfd[0]);

    TEST_ASSERT_EQ(CONN_TYPE_UNUSED, conn->type, "Connection closed");

    close(sockfd[1]);
    close(epoll_fd);
    epoll_fd = old_epoll;

    TEST_SUITE_END();
}

void test_connection_tcp_nodelay(void) {
    TEST_SUITE_BEGIN("Connection TCP NODELAY");

    int sockfd = socket(AF_INET, SOCK_STREAM, 0);
    TEST_ASSERT(sockfd >= 0, "Socket created");

    connection_set_tcp_nodelay(sockfd);

    int optval;
    socklen_t optlen = sizeof(optval);
    getsockopt(sockfd, IPPROTO_TCP, TCP_NODELAY, &optval, &optlen);
    TEST_ASSERT_EQ(1, optval, "TCP_NODELAY is set");

    close(sockfd);

    TEST_SUITE_END();
}

void test_connection_close_with_upstream_pair(void) {
    TEST_SUITE_BEGIN("Connection Close With Upstream Pair");

    connection_init_all();
    int old_epoll = epoll_fd;
    epoll_fd = epoll_create1(0);
    TEST_ASSERT(epoll_fd >= 0, "Epoll created");

    int sockfd1[2], sockfd2[2];
    socketpair(AF_UNIX, SOCK_STREAM, 0, sockfd1);
    socketpair(AF_UNIX, SOCK_STREAM, 0, sockfd2);

    connection_t *client = &connections[sockfd1[0]];
    memset(client, 0, sizeof(connection_t));
    client->fd = sockfd1[0];
    client->type = CONN_TYPE_CLIENT;
    client->state = CLIENT_STATE_FORWARDING;
    buffer_init(&client->read_buf, 4096);
    buffer_init(&client->write_buf, 4096);

    connection_t *upstream = &connections[sockfd2[0]];
    memset(upstream, 0, sizeof(connection_t));
    upstream->fd = sockfd2[0];
    upstream->type = CONN_TYPE_UPSTREAM;
    upstream->state = CLIENT_STATE_FORWARDING;
    buffer_init(&upstream->read_buf, 4096);
    buffer_init(&upstream->write_buf, 4096);

    client->pair = upstream;
    upstream->pair = client;

    struct epoll_event ev1 = { .events = EPOLLIN, .data.fd = sockfd1[0] };
    struct epoll_event ev2 = { .events = EPOLLIN, .data.fd = sockfd2[0] };
    epoll_ctl(epoll_fd, EPOLL_CTL_ADD, sockfd1[0], &ev1);
    epoll_ctl(epoll_fd, EPOLL_CTL_ADD, sockfd2[0], &ev2);

    connection_close(sockfd2[0]);

    TEST_ASSERT_EQ(CLIENT_STATE_READING_HEADERS, client->state, "Client reset to READING_HEADERS");
    TEST_ASSERT(client->pair == NULL, "Client pair is NULL");

    buffer_free(&client->read_buf);
    buffer_free(&client->write_buf);
    close(sockfd1[0]);
    close(sockfd1[1]);
    close(sockfd2[1]);
    close(epoll_fd);
    epoll_fd = old_epoll;

    TEST_SUITE_END();
}

void test_connection_close_client_closes_upstream(void) {
    TEST_SUITE_BEGIN("Connection Close Client Closes Upstream");

    connection_init_all();
    int old_epoll = epoll_fd;
    epoll_fd = epoll_create1(0);
    TEST_ASSERT(epoll_fd >= 0, "Epoll created");

    int sockfd1[2], sockfd2[2];
    socketpair(AF_UNIX, SOCK_STREAM, 0, sockfd1);
    socketpair(AF_UNIX, SOCK_STREAM, 0, sockfd2);

    connection_t *client = &connections[sockfd1[0]];
    memset(client, 0, sizeof(connection_t));
    client->fd = sockfd1[0];
    client->type = CONN_TYPE_CLIENT;
    client->state = CLIENT_STATE_FORWARDING;
    buffer_init(&client->read_buf, 4096);
    buffer_init(&client->write_buf, 4096);

    connection_t *upstream = &connections[sockfd2[0]];
    memset(upstream, 0, sizeof(connection_t));
    upstream->fd = sockfd2[0];
    upstream->type = CONN_TYPE_UPSTREAM;
    upstream->state = CLIENT_STATE_FORWARDING;
    buffer_init(&upstream->read_buf, 4096);
    buffer_init(&upstream->write_buf, 4096);

    client->pair = upstream;
    upstream->pair = client;

    struct epoll_event ev1 = { .events = EPOLLIN, .data.fd = sockfd1[0] };
    struct epoll_event ev2 = { .events = EPOLLIN, .data.fd = sockfd2[0] };
    epoll_ctl(epoll_fd, EPOLL_CTL_ADD, sockfd1[0], &ev1);
    epoll_ctl(epoll_fd, EPOLL_CTL_ADD, sockfd2[0], &ev2);

    connection_close(sockfd1[0]);

    TEST_ASSERT_EQ(CONN_TYPE_UNUSED, upstream->type, "Upstream also closed");

    close(sockfd1[1]);
    close(sockfd2[1]);
    close(epoll_fd);
    epoll_fd = old_epoll;

    TEST_SUITE_END();
}

void test_connection_handle_event_listener(void) {
    TEST_SUITE_BEGIN("Connection Handle Event Listener");

    connection_init_all();
    int old_epoll = epoll_fd;
    epoll_fd = epoll_create1(0);
    TEST_ASSERT(epoll_fd >= 0, "Epoll created");

    int listener = socket(AF_INET, SOCK_STREAM, 0);
    TEST_ASSERT(listener >= 0, "Listener socket created");

    int reuse = 1;
    setsockopt(listener, SOL_SOCKET, SO_REUSEADDR, &reuse, sizeof(reuse));

    struct sockaddr_in addr = {
        .sin_family = AF_INET,
        .sin_port = htons(0),
        .sin_addr.s_addr = htonl(INADDR_LOOPBACK)
    };
    int result = bind(listener, (struct sockaddr*)&addr, sizeof(addr));
    TEST_ASSERT_EQ(0, result, "Bind succeeded");

    listen(listener, 1);
    connection_set_non_blocking(listener);

    connections[listener].type = CONN_TYPE_LISTENER;
    connections[listener].fd = listener;

    struct epoll_event ev = { .events = EPOLLIN, .data.fd = listener };
    epoll_ctl(epoll_fd, EPOLL_CTL_ADD, listener, &ev);

    struct epoll_event event = { .events = EPOLLIN, .data.fd = listener };
    connection_handle_event(&event);
    TEST_ASSERT(1, "Handle event on listener doesn't crash");

    connection_close(listener);
    close(epoll_fd);
    epoll_fd = old_epoll;

    TEST_SUITE_END();
}

void test_connection_handle_event_error(void) {
    TEST_SUITE_BEGIN("Connection Handle Event Error");

    connection_init_all();
    int old_epoll = epoll_fd;
    epoll_fd = epoll_create1(0);
    TEST_ASSERT(epoll_fd >= 0, "Epoll created");

    int sockfd[2];
    socketpair(AF_UNIX, SOCK_STREAM, 0, sockfd);

    connection_t *conn = &connections[sockfd[0]];
    memset(conn, 0, sizeof(connection_t));
    conn->fd = sockfd[0];
    conn->type = CONN_TYPE_CLIENT;
    conn->state = CLIENT_STATE_READING_HEADERS;
    buffer_init(&conn->read_buf, 4096);
    buffer_init(&conn->write_buf, 4096);

    struct epoll_event ev = { .events = EPOLLIN, .data.fd = sockfd[0] };
    epoll_ctl(epoll_fd, EPOLL_CTL_ADD, sockfd[0], &ev);

    struct epoll_event error_event = { .events = EPOLLERR, .data.fd = sockfd[0] };
    connection_handle_event(&error_event);

    TEST_ASSERT_EQ(CONN_TYPE_UNUSED, conn->type, "Connection closed on EPOLLERR");

    close(sockfd[1]);
    close(epoll_fd);
    epoll_fd = old_epoll;

    TEST_SUITE_END();
}

void test_connection_handle_event_hup(void) {
    TEST_SUITE_BEGIN("Connection Handle Event HUP");

    connection_init_all();
    int old_epoll = epoll_fd;
    epoll_fd = epoll_create1(0);
    TEST_ASSERT(epoll_fd >= 0, "Epoll created");

    int sockfd[2];
    socketpair(AF_UNIX, SOCK_STREAM, 0, sockfd);

    connection_t *conn = &connections[sockfd[0]];
    memset(conn, 0, sizeof(connection_t));
    conn->fd = sockfd[0];
    conn->type = CONN_TYPE_CLIENT;
    conn->state = CLIENT_STATE_READING_HEADERS;
    buffer_init(&conn->read_buf, 4096);
    buffer_init(&conn->write_buf, 4096);

    struct epoll_event ev = { .events = EPOLLIN, .data.fd = sockfd[0] };
    epoll_ctl(epoll_fd, EPOLL_CTL_ADD, sockfd[0], &ev);

    struct epoll_event hup_event = { .events = EPOLLHUP, .data.fd = sockfd[0] };
    connection_handle_event(&hup_event);

    TEST_ASSERT_EQ(CONN_TYPE_UNUSED, conn->type, "Connection closed on EPOLLHUP");

    close(sockfd[1]);
    close(epoll_fd);
    epoll_fd = old_epoll;

    TEST_SUITE_END();
}

void test_connection_handle_event_invalid_fd(void) {
    TEST_SUITE_BEGIN("Connection Handle Event Invalid FD");

    struct epoll_event event = { .events = EPOLLIN, .data.fd = -1 };
    connection_handle_event(&event);
    TEST_ASSERT(1, "Invalid fd doesn't crash");

    event.data.fd = MAX_FDS + 100;
    connection_handle_event(&event);
    TEST_ASSERT(1, "Out of range fd doesn't crash");

    TEST_SUITE_END();
}

void test_connection_handle_event_unused(void) {
    TEST_SUITE_BEGIN("Connection Handle Event Unused");

    connection_init_all();

    struct epoll_event event = { .events = EPOLLIN, .data.fd = 5 };
    connection_handle_event(&event);
    TEST_ASSERT(1, "Event on unused connection doesn't crash");

    TEST_SUITE_END();
}

void test_connection_cleanup_idle_with_idle_conn(void) {
    TEST_SUITE_BEGIN("Connection Cleanup Idle With Idle Connection");

    connection_init_all();
    int old_epoll = epoll_fd;
    epoll_fd = epoll_create1(0);
    TEST_ASSERT(epoll_fd >= 0, "Epoll created");

    int sockfd[2];
    socketpair(AF_UNIX, SOCK_STREAM, 0, sockfd);

    connection_t *conn = &connections[sockfd[0]];
    memset(conn, 0, sizeof(connection_t));
    conn->fd = sockfd[0];
    conn->type = CONN_TYPE_CLIENT;
    conn->state = CLIENT_STATE_READING_HEADERS;
    conn->last_activity = time(NULL) - 1000;
    buffer_init(&conn->read_buf, 4096);
    buffer_init(&conn->write_buf, 4096);

    struct epoll_event ev = { .events = EPOLLIN, .data.fd = sockfd[0] };
    epoll_ctl(epoll_fd, EPOLL_CTL_ADD, sockfd[0], &ev);

    connection_cleanup_idle();

    TEST_ASSERT_EQ(CONN_TYPE_UNUSED, conn->type, "Idle connection closed");

    close(sockfd[1]);
    close(epoll_fd);
    epoll_fd = old_epoll;

    TEST_SUITE_END();
}

void test_connection_send_error_null_params(void) {
    TEST_SUITE_BEGIN("Connection Send Error NULL Params");

    connection_send_error_response(NULL, 500, "Error", "Body");
    TEST_ASSERT(1, "NULL connection doesn't crash");

    connection_init_all();
    int old_epoll = epoll_fd;
    epoll_fd = epoll_create1(0);

    int sockfd[2];
    socketpair(AF_UNIX, SOCK_STREAM, 0, sockfd);

    connection_t *conn = &connections[sockfd[0]];
    memset(conn, 0, sizeof(connection_t));
    conn->fd = sockfd[0];
    conn->type = CONN_TYPE_CLIENT;
    buffer_init(&conn->read_buf, 4096);
    buffer_init(&conn->write_buf, 4096);

    struct epoll_event ev = { .events = EPOLLIN, .data.fd = sockfd[0] };
    epoll_ctl(epoll_fd, EPOLL_CTL_ADD, sockfd[0], &ev);

    connection_send_error_response(conn, 500, NULL, "Body");
    TEST_ASSERT(1, "NULL status doesn't crash");

    connection_send_error_response(conn, 500, "Error", NULL);
    TEST_ASSERT(1, "NULL body doesn't crash");

    buffer_free(&conn->read_buf);
    buffer_free(&conn->write_buf);
    close(sockfd[0]);
    close(sockfd[1]);
    close(epoll_fd);
    epoll_fd = old_epoll;

    TEST_SUITE_END();
}

void test_connection_send_auth_null_conn(void) {
    TEST_SUITE_BEGIN("Connection Send Auth NULL Connection");

    connection_send_auth_required(NULL, "Test");
    TEST_ASSERT(1, "NULL connection doesn't crash");

    TEST_SUITE_END();
}

void test_connection_buffer_compaction(void) {
    TEST_SUITE_BEGIN("Connection Buffer Compaction");

    int sockfd[2];
    socketpair(AF_UNIX, SOCK_STREAM, 0, sockfd);

    connection_set_non_blocking(sockfd[0]);
    connection_set_non_blocking(sockfd[1]);

    connection_t conn;
    memset(&conn, 0, sizeof(conn));
    conn.fd = sockfd[0];
    conn.type = CONN_TYPE_CLIENT;
    buffer_init(&conn.read_buf, 4096);

    conn.read_buf.head = 2000;
    conn.read_buf.tail = 2500;

    const char *test_data = "Test data for compaction";
    write(sockfd[1], test_data, strlen(test_data));

    int bytes = connection_do_read(&conn);
    TEST_ASSERT(bytes > 0 || (bytes < 0 && conn.read_buf.head < 2000), "Buffer compacted or data read");

    buffer_free(&conn.read_buf);
    close(sockfd[0]);
    close(sockfd[1]);

    TEST_SUITE_END();
}

void test_connection_do_write_empty_buffer(void) {
    TEST_SUITE_BEGIN("Connection Do Write Empty Buffer");

    int sockfd[2];
    socketpair(AF_UNIX, SOCK_STREAM, 0, sockfd);

    connection_t conn;
    memset(&conn, 0, sizeof(conn));
    conn.fd = sockfd[0];
    conn.type = CONN_TYPE_CLIENT;
    buffer_init(&conn.write_buf, 4096);

    int bytes = connection_do_write(&conn);
    TEST_ASSERT_EQ(0, bytes, "Empty buffer write returns 0");

    buffer_free(&conn.write_buf);
    close(sockfd[0]);
    close(sockfd[1]);

    TEST_SUITE_END();
}

void test_connection_modify_epoll_same_events(void) {
    TEST_SUITE_BEGIN("Connection Modify Epoll Same Events");

    connection_init_all();
    int old_epoll = epoll_fd;
    epoll_fd = epoll_create1(0);
    TEST_ASSERT(epoll_fd >= 0, "Epoll created");

    int sockfd[2];
    socketpair(AF_UNIX, SOCK_STREAM, 0, sockfd);

    connection_add_to_epoll(sockfd[0], EPOLLIN);

    connection_modify_epoll(sockfd[0], EPOLLIN);
    TEST_ASSERT_EQ((uint32_t)EPOLLIN, connections[sockfd[0]].epoll_events, "Same events not modified");

    close(sockfd[0]);
    close(sockfd[1]);
    close(epoll_fd);
    epoll_fd = old_epoll;

    TEST_SUITE_END();
}

void test_connection_accept_new(void) {
    TEST_SUITE_BEGIN("Connection Accept New");

    connection_init_all();
    connection_update_cached_time();
    int old_epoll = epoll_fd;
    epoll_fd = epoll_create1(0);
    TEST_ASSERT(epoll_fd >= 0, "Epoll created");

    int listener = socket(AF_INET, SOCK_STREAM, 0);
    TEST_ASSERT(listener >= 0, "Listener socket created");

    int reuse = 1;
    setsockopt(listener, SOL_SOCKET, SO_REUSEADDR, &reuse, sizeof(reuse));

    struct sockaddr_in addr = {
        .sin_family = AF_INET,
        .sin_port = htons(0),
        .sin_addr.s_addr = htonl(INADDR_LOOPBACK)
    };
    int result = bind(listener, (struct sockaddr*)&addr, sizeof(addr));
    TEST_ASSERT_EQ(0, result, "Bind succeeded");

    socklen_t addr_len = sizeof(addr);
    getsockname(listener, (struct sockaddr*)&addr, &addr_len);
    int port = ntohs(addr.sin_port);

    listen(listener, 5);
    connection_set_non_blocking(listener);

    connections[listener].type = CONN_TYPE_LISTENER;
    connections[listener].fd = listener;

    struct epoll_event ev = { .events = EPOLLIN, .data.fd = listener };
    epoll_ctl(epoll_fd, EPOLL_CTL_ADD, listener, &ev);

    int client = socket(AF_INET, SOCK_STREAM, 0);
    TEST_ASSERT(client >= 0, "Client socket created");

    struct sockaddr_in client_addr = {
        .sin_family = AF_INET,
        .sin_port = htons(port),
        .sin_addr.s_addr = htonl(INADDR_LOOPBACK)
    };
    result = connect(client, (struct sockaddr*)&client_addr, sizeof(client_addr));
    TEST_ASSERT_EQ(0, result, "Connect succeeded");

    connection_accept(listener);
    TEST_ASSERT(1, "Accept completed");

    close(client);
    connection_close(listener);
    for (int i = 0; i < MAX_FDS; i++) {
        if (connections[i].type != CONN_TYPE_UNUSED && connections[i].type != CONN_TYPE_LISTENER) {
            connection_close(i);
        }
    }
    close(epoll_fd);
    epoll_fd = old_epoll;

    TEST_SUITE_END();
}

void test_connection_handle_event_write(void) {
    TEST_SUITE_BEGIN("Connection Handle Event Write");

    connection_init_all();
    connection_update_cached_time();
    int old_epoll = epoll_fd;
    epoll_fd = epoll_create1(0);
    TEST_ASSERT(epoll_fd >= 0, "Epoll created");

    int sockfd[2];
    socketpair(AF_UNIX, SOCK_STREAM, 0, sockfd);

    connection_set_non_blocking(sockfd[0]);
    connection_set_non_blocking(sockfd[1]);

    connection_t *conn = &connections[sockfd[0]];
    memset(conn, 0, sizeof(connection_t));
    conn->fd = sockfd[0];
    conn->type = CONN_TYPE_CLIENT;
    conn->state = CLIENT_STATE_SERVING_INTERNAL;
    buffer_init(&conn->read_buf, 4096);
    buffer_init(&conn->write_buf, 4096);

    const char *data = "HTTP/1.1 200 OK\r\n\r\n";
    memcpy(conn->write_buf.data, data, strlen(data));
    conn->write_buf.tail = strlen(data);

    struct epoll_event ev = { .events = EPOLLIN | EPOLLOUT, .data.fd = sockfd[0] };
    epoll_ctl(epoll_fd, EPOLL_CTL_ADD, sockfd[0], &ev);
    conn->epoll_events = EPOLLIN | EPOLLOUT;

    struct epoll_event event = { .events = EPOLLOUT, .data.fd = sockfd[0] };
    connection_handle_event(&event);

    TEST_ASSERT(conn->write_buf.head > 0 || conn->type == CONN_TYPE_UNUSED, "Data written or connection closed");

    if (conn->type != CONN_TYPE_UNUSED) {
        buffer_free(&conn->read_buf);
        buffer_free(&conn->write_buf);
    }
    close(sockfd[0]);
    close(sockfd[1]);
    close(epoll_fd);
    epoll_fd = old_epoll;

    TEST_SUITE_END();
}

void test_connection_close_with_vhost_stats(void) {
    TEST_SUITE_BEGIN("Connection Close With VHost Stats");

    connection_init_all();
    int old_epoll = epoll_fd;
    epoll_fd = epoll_create1(0);
    TEST_ASSERT(epoll_fd >= 0, "Epoll created");

    int sockfd[2];
    socketpair(AF_UNIX, SOCK_STREAM, 0, sockfd);

    connection_t *conn = &connections[sockfd[0]];
    memset(conn, 0, sizeof(connection_t));
    conn->fd = sockfd[0];
    conn->type = CONN_TYPE_CLIENT;
    conn->state = CLIENT_STATE_FORWARDING;
    buffer_init(&conn->read_buf, 4096);
    buffer_init(&conn->write_buf, 4096);

    conn->vhost_stats = monitor_get_or_create_vhost_stats("test.example.com");
    struct timespec ts;
    clock_gettime(CLOCK_MONOTONIC, &ts);
    conn->request_start_time = ts.tv_sec + ts.tv_nsec / 1e9;

    struct epoll_event ev = { .events = EPOLLIN, .data.fd = sockfd[0] };
    epoll_ctl(epoll_fd, EPOLL_CTL_ADD, sockfd[0], &ev);

    connection_close(sockfd[0]);
    TEST_ASSERT_EQ(CONN_TYPE_UNUSED, conn->type, "Connection closed with stats");

    close(sockfd[1]);
    close(epoll_fd);
    epoll_fd = old_epoll;

    TEST_SUITE_END();
}

void test_connection_close_with_config(void) {
    TEST_SUITE_BEGIN("Connection Close With Config Ref");

    connection_init_all();
    int old_epoll = epoll_fd;
    epoll_fd = epoll_create1(0);
    TEST_ASSERT(epoll_fd >= 0, "Epoll created");

    int sockfd[2];
    socketpair(AF_UNIX, SOCK_STREAM, 0, sockfd);

    connection_t *conn = &connections[sockfd[0]];
    memset(conn, 0, sizeof(connection_t));
    conn->fd = sockfd[0];
    conn->type = CONN_TYPE_CLIENT;
    buffer_init(&conn->read_buf, 4096);
    buffer_init(&conn->write_buf, 4096);

    struct epoll_event ev = { .events = EPOLLIN, .data.fd = sockfd[0] };
    epoll_ctl(epoll_fd, EPOLL_CTL_ADD, sockfd[0], &ev);

    connection_close(sockfd[0]);
    TEST_ASSERT_EQ(CONN_TYPE_UNUSED, conn->type, "Connection closed");

    close(sockfd[1]);
    close(epoll_fd);
    epoll_fd = old_epoll;

    TEST_SUITE_END();
}

void test_connection_send_error_various_codes(void) {
    TEST_SUITE_BEGIN("Connection Send Error Various Codes");

    connection_init_all();
    int old_epoll = epoll_fd;
    epoll_fd = epoll_create1(0);

    int sockfd[2];
    socketpair(AF_UNIX, SOCK_STREAM, 0, sockfd);

    connection_t *conn = &connections[sockfd[0]];
    memset(conn, 0, sizeof(connection_t));
    conn->fd = sockfd[0];
    conn->type = CONN_TYPE_CLIENT;
    buffer_init(&conn->read_buf, 4096);
    buffer_init(&conn->write_buf, 4096);

    struct epoll_event ev = { .events = EPOLLIN, .data.fd = sockfd[0] };
    epoll_ctl(epoll_fd, EPOLL_CTL_ADD, sockfd[0], &ev);

    connection_send_error_response(conn, 400, "Bad Request", "Invalid request");
    TEST_ASSERT(strstr((char*)conn->write_buf.data, "400") != NULL, "400 error generated");

    conn->write_buf.head = 0;
    conn->write_buf.tail = 0;

    connection_send_error_response(conn, 502, "Bad Gateway", "Upstream error");
    TEST_ASSERT(strstr((char*)conn->write_buf.data, "502") != NULL, "502 error generated");

    buffer_free(&conn->read_buf);
    buffer_free(&conn->write_buf);
    close(sockfd[0]);
    close(sockfd[1]);
    close(epoll_fd);
    epoll_fd = old_epoll;

    TEST_SUITE_END();
}

void test_connection_send_auth_null_realm(void) {
    TEST_SUITE_BEGIN("Connection Send Auth NULL Realm");

    connection_init_all();
    int old_epoll = epoll_fd;
    epoll_fd = epoll_create1(0);

    int sockfd[2];
    socketpair(AF_UNIX, SOCK_STREAM, 0, sockfd);

    connection_t *conn = &connections[sockfd[0]];
    memset(conn, 0, sizeof(connection_t));
    conn->fd = sockfd[0];
    conn->type = CONN_TYPE_CLIENT;
    buffer_init(&conn->read_buf, 4096);
    buffer_init(&conn->write_buf, 4096);

    struct epoll_event ev = { .events = EPOLLIN, .data.fd = sockfd[0] };
    epoll_ctl(epoll_fd, EPOLL_CTL_ADD, sockfd[0], &ev);

    connection_send_auth_required(conn, NULL);
    TEST_ASSERT(strstr((char*)conn->write_buf.data, "Protected Area") != NULL, "Default realm used");

    buffer_free(&conn->read_buf);
    buffer_free(&conn->write_buf);
    close(sockfd[0]);
    close(sockfd[1]);
    close(epoll_fd);
    epoll_fd = old_epoll;

    TEST_SUITE_END();
}

void test_connection_do_read_grow_buffer(void) {
    TEST_SUITE_BEGIN("Connection Do Read Grow Buffer");

    int sockfd[2];
    socketpair(AF_UNIX, SOCK_STREAM, 0, sockfd);

    connection_set_non_blocking(sockfd[0]);
    connection_set_non_blocking(sockfd[1]);

    connection_t conn;
    memset(&conn, 0, sizeof(conn));
    conn.fd = sockfd[0];
    conn.type = CONN_TYPE_CLIENT;
    buffer_init(&conn.read_buf, 64);

    conn.read_buf.tail = conn.read_buf.capacity;

    char large_data[128];
    memset(large_data, 'A', sizeof(large_data));
    write(sockfd[1], large_data, sizeof(large_data));

    int bytes = connection_do_read(&conn);
    TEST_ASSERT(bytes > 0 || conn.read_buf.capacity > 64, "Buffer grew or data read");

    buffer_free(&conn.read_buf);
    close(sockfd[0]);
    close(sockfd[1]);

    TEST_SUITE_END();
}

void test_connection_close_with_splice_pipes(void) {
    TEST_SUITE_BEGIN("Connection Close With Splice Pipes");

    connection_init_all();
    int old_epoll = epoll_fd;
    epoll_fd = epoll_create1(0);
    TEST_ASSERT(epoll_fd >= 0, "Epoll created");

    int sockfd[2];
    socketpair(AF_UNIX, SOCK_STREAM, 0, sockfd);

    connection_t *conn = &connections[sockfd[0]];
    memset(conn, 0, sizeof(connection_t));
    conn->fd = sockfd[0];
    conn->type = CONN_TYPE_UPSTREAM;
    buffer_init(&conn->read_buf, 4096);
    buffer_init(&conn->write_buf, 4096);

    int pipefd[2];
    pipe(pipefd);
    conn->splice_pipe[0] = pipefd[0];
    conn->splice_pipe[1] = pipefd[1];
    conn->can_splice = 1;

    struct epoll_event ev = { .events = EPOLLIN, .data.fd = sockfd[0] };
    epoll_ctl(epoll_fd, EPOLL_CTL_ADD, sockfd[0], &ev);

    connection_close(sockfd[0]);
    TEST_ASSERT_EQ(CONN_TYPE_UNUSED, conn->type, "Connection with splice pipes closed");
    TEST_ASSERT_EQ(-1, conn->splice_pipe[0], "Splice pipe[0] reset");
    TEST_ASSERT_EQ(-1, conn->splice_pipe[1], "Splice pipe[1] reset");

    close(sockfd[1]);
    close(epoll_fd);
    epoll_fd = old_epoll;

    TEST_SUITE_END();
}

void test_connection_modify_epoll_invalid(void) {
    TEST_SUITE_BEGIN("Connection Modify Epoll Invalid");

    connection_init_all();
    int old_epoll = epoll_fd;
    epoll_fd = epoll_create1(0);

    connection_modify_epoll(-1, EPOLLIN);
    TEST_ASSERT(1, "Invalid fd doesn't crash");

    connection_modify_epoll(MAX_FDS + 100, EPOLLIN);
    TEST_ASSERT(1, "Out of range fd doesn't crash");

    connection_modify_epoll(5, EPOLLIN);
    TEST_ASSERT(1, "Non-existent fd in epoll doesn't crash");

    close(epoll_fd);
    epoll_fd = old_epoll;

    TEST_SUITE_END();
}

void test_connection_add_epoll_failure(void) {
    TEST_SUITE_BEGIN("Connection Add Epoll Handles Failure");

    connection_init_all();
    int old_epoll = epoll_fd;
    epoll_fd = -1;

    int sockfd[2];
    socketpair(AF_UNIX, SOCK_STREAM, 0, sockfd);

    connection_add_to_epoll(sockfd[0], EPOLLIN);
    TEST_ASSERT(1, "Add to invalid epoll doesn't crash");

    close(sockfd[1]);
    epoll_fd = old_epoll;

    TEST_SUITE_END();
}

void test_connection_connect_upstream_no_config(void) {
    TEST_SUITE_BEGIN("Connection Connect Upstream No Config");

    connection_init_all();
    int old_epoll = epoll_fd;
    epoll_fd = epoll_create1(0);
    TEST_ASSERT(epoll_fd >= 0, "Epoll created");

    int sockfd[2];
    socketpair(AF_UNIX, SOCK_STREAM, 0, sockfd);

    connection_t *conn = &connections[sockfd[0]];
    memset(conn, 0, sizeof(connection_t));
    conn->fd = sockfd[0];
    conn->type = CONN_TYPE_CLIENT;
    conn->state = CLIENT_STATE_READING_HEADERS;
    buffer_init(&conn->read_buf, 4096);
    buffer_init(&conn->write_buf, 4096);

    struct epoll_event ev = { .events = EPOLLIN, .data.fd = sockfd[0] };
    epoll_ctl(epoll_fd, EPOLL_CTL_ADD, sockfd[0], &ev);

    app_config_t *old_config = config;
    config = NULL;

    connection_connect_to_upstream(conn, "GET / HTTP/1.1\r\n\r\n", 18);
    TEST_ASSERT(conn->write_buf.tail > 0 || conn->state == CLIENT_STATE_ERROR, "Error response or state changed");

    config = old_config;
    buffer_free(&conn->read_buf);
    buffer_free(&conn->write_buf);
    close(sockfd[0]);
    close(sockfd[1]);
    close(epoll_fd);
    epoll_fd = old_epoll;

    TEST_SUITE_END();
}

void test_connection_connect_upstream_null_params(void) {
    TEST_SUITE_BEGIN("Connection Connect Upstream NULL Params");

    connection_connect_to_upstream(NULL, "data", 4);
    TEST_ASSERT(1, "NULL connection doesn't crash");

    connection_t conn;
    memset(&conn, 0, sizeof(conn));
    connection_connect_to_upstream(&conn, NULL, 0);
    TEST_ASSERT(1, "NULL data doesn't crash");

    TEST_SUITE_END();
}

void test_connection_handle_event_client_read(void) {
    TEST_SUITE_BEGIN("Connection Handle Event Client Read");

    connection_init_all();
    connection_update_cached_time();
    int old_epoll = epoll_fd;
    epoll_fd = epoll_create1(0);
    TEST_ASSERT(epoll_fd >= 0, "Epoll created");

    int sockfd[2];
    socketpair(AF_UNIX, SOCK_STREAM, 0, sockfd);

    connection_set_non_blocking(sockfd[0]);
    connection_set_non_blocking(sockfd[1]);

    connection_t *conn = &connections[sockfd[0]];
    memset(conn, 0, sizeof(connection_t));
    conn->fd = sockfd[0];
    conn->type = CONN_TYPE_CLIENT;
    conn->state = CLIENT_STATE_READING_HEADERS;
    buffer_init(&conn->read_buf, 4096);
    buffer_init(&conn->write_buf, 4096);

    struct epoll_event ev = { .events = EPOLLIN, .data.fd = sockfd[0] };
    epoll_ctl(epoll_fd, EPOLL_CTL_ADD, sockfd[0], &ev);
    conn->epoll_events = EPOLLIN;

    const char *partial_request = "GET / HTTP/1.1\r\n";
    write(sockfd[1], partial_request, strlen(partial_request));

    struct epoll_event event = { .events = EPOLLIN, .data.fd = sockfd[0] };
    connection_handle_event(&event);
    TEST_ASSERT(conn->state == CLIENT_STATE_READING_HEADERS, "Still reading headers with partial");

    if (conn->type != CONN_TYPE_UNUSED) {
        buffer_free(&conn->read_buf);
        buffer_free(&conn->write_buf);
    }
    close(sockfd[0]);
    close(sockfd[1]);
    close(epoll_fd);
    epoll_fd = old_epoll;

    TEST_SUITE_END();
}

void test_connection_handle_complete_request(void) {
    TEST_SUITE_BEGIN("Connection Handle Complete Request");

    connection_init_all();
    connection_update_cached_time();
    int old_epoll = epoll_fd;
    epoll_fd = epoll_create1(0);
    TEST_ASSERT(epoll_fd >= 0, "Epoll created");

    int sockfd[2];
    socketpair(AF_UNIX, SOCK_STREAM, 0, sockfd);

    connection_set_non_blocking(sockfd[0]);
    connection_set_non_blocking(sockfd[1]);

    connection_t *conn = &connections[sockfd[0]];
    memset(conn, 0, sizeof(connection_t));
    conn->fd = sockfd[0];
    conn->type = CONN_TYPE_CLIENT;
    conn->state = CLIENT_STATE_READING_HEADERS;
    buffer_init(&conn->read_buf, 4096);
    buffer_init(&conn->write_buf, 4096);

    struct epoll_event ev = { .events = EPOLLIN, .data.fd = sockfd[0] };
    epoll_ctl(epoll_fd, EPOLL_CTL_ADD, sockfd[0], &ev);
    conn->epoll_events = EPOLLIN;

    const char *request = "GET / HTTP/1.1\r\nHost: test.example.com\r\n\r\n";
    write(sockfd[1], request, strlen(request));

    struct epoll_event event = { .events = EPOLLIN, .data.fd = sockfd[0] };
    connection_handle_event(&event);
    TEST_ASSERT(1, "Complete request handled");

    if (conn->type != CONN_TYPE_UNUSED) {
        buffer_free(&conn->read_buf);
        buffer_free(&conn->write_buf);
    }
    close(sockfd[0]);
    close(sockfd[1]);
    close(epoll_fd);
    epoll_fd = old_epoll;

    TEST_SUITE_END();
}

void test_connection_handle_dashboard_request(void) {
    TEST_SUITE_BEGIN("Connection Handle Dashboard Request");

    connection_init_all();
    connection_update_cached_time();
    int old_epoll = epoll_fd;
    epoll_fd = epoll_create1(0);
    TEST_ASSERT(epoll_fd >= 0, "Epoll created");

    int sockfd[2];
    socketpair(AF_UNIX, SOCK_STREAM, 0, sockfd);

    connection_set_non_blocking(sockfd[0]);
    connection_set_non_blocking(sockfd[1]);

    connection_t *conn = &connections[sockfd[0]];
    memset(conn, 0, sizeof(connection_t));
    conn->fd = sockfd[0];
    conn->type = CONN_TYPE_CLIENT;
    conn->state = CLIENT_STATE_READING_HEADERS;
    buffer_init(&conn->read_buf, 4096);
    buffer_init(&conn->write_buf, 4096);

    struct epoll_event ev = { .events = EPOLLIN, .data.fd = sockfd[0] };
    epoll_ctl(epoll_fd, EPOLL_CTL_ADD, sockfd[0], &ev);
    conn->epoll_events = EPOLLIN;

    const char *request = "GET /rproxy/dashboard HTTP/1.1\r\nHost: localhost\r\n\r\n";
    write(sockfd[1], request, strlen(request));

    struct epoll_event event = { .events = EPOLLIN, .data.fd = sockfd[0] };
    connection_handle_event(&event);
    TEST_ASSERT(conn->state == CLIENT_STATE_SERVING_INTERNAL || conn->type == CONN_TYPE_UNUSED || conn->write_buf.tail > 0, "Dashboard request processed");

    if (conn->type != CONN_TYPE_UNUSED) {
        buffer_free(&conn->read_buf);
        buffer_free(&conn->write_buf);
    }
    close(sockfd[0]);
    close(sockfd[1]);
    close(epoll_fd);
    epoll_fd = old_epoll;

    TEST_SUITE_END();
}

void test_connection_handle_stats_api(void) {
    TEST_SUITE_BEGIN("Connection Handle Stats API");

    connection_init_all();
    connection_update_cached_time();
    int old_epoll = epoll_fd;
    epoll_fd = epoll_create1(0);
    TEST_ASSERT(epoll_fd >= 0, "Epoll created");

    int sockfd[2];
    socketpair(AF_UNIX, SOCK_STREAM, 0, sockfd);

    connection_set_non_blocking(sockfd[0]);
    connection_set_non_blocking(sockfd[1]);

    connection_t *conn = &connections[sockfd[0]];
    memset(conn, 0, sizeof(connection_t));
    conn->fd = sockfd[0];
    conn->type = CONN_TYPE_CLIENT;
    conn->state = CLIENT_STATE_READING_HEADERS;
    buffer_init(&conn->read_buf, 4096);
    buffer_init(&conn->write_buf, 4096);

    struct epoll_event ev = { .events = EPOLLIN, .data.fd = sockfd[0] };
    epoll_ctl(epoll_fd, EPOLL_CTL_ADD, sockfd[0], &ev);
    conn->epoll_events = EPOLLIN;

    const char *request = "GET /rproxy/api/stats HTTP/1.1\r\nHost: localhost\r\n\r\n";
    write(sockfd[1], request, strlen(request));

    struct epoll_event event = { .events = EPOLLIN, .data.fd = sockfd[0] };
    connection_handle_event(&event);
    TEST_ASSERT(conn->state == CLIENT_STATE_SERVING_INTERNAL || conn->type == CONN_TYPE_UNUSED || conn->write_buf.tail > 0, "Stats API request processed");

    if (conn->type != CONN_TYPE_UNUSED) {
        buffer_free(&conn->read_buf);
        buffer_free(&conn->write_buf);
    }
    close(sockfd[0]);
    close(sockfd[1]);
    close(epoll_fd);
    epoll_fd = old_epoll;

    TEST_SUITE_END();
}

void test_connection_handle_oversized_header(void) {
    TEST_SUITE_BEGIN("Connection Handle Oversized Header");

    connection_init_all();
    connection_update_cached_time();
    int old_epoll = epoll_fd;
    epoll_fd = epoll_create1(0);
    TEST_ASSERT(epoll_fd >= 0, "Epoll created");

    int sockfd[2];
    socketpair(AF_UNIX, SOCK_STREAM, 0, sockfd);

    connection_set_non_blocking(sockfd[0]);
    connection_set_non_blocking(sockfd[1]);

    connection_t *conn = &connections[sockfd[0]];
    memset(conn, 0, sizeof(connection_t));
    conn->fd = sockfd[0];
    conn->type = CONN_TYPE_CLIENT;
    conn->state = CLIENT_STATE_READING_HEADERS;
    buffer_init(&conn->read_buf, 65536);
    buffer_init(&conn->write_buf, 4096);

    struct epoll_event ev = { .events = EPOLLIN, .data.fd = sockfd[0] };
    epoll_ctl(epoll_fd, EPOLL_CTL_ADD, sockfd[0], &ev);
    conn->epoll_events = EPOLLIN;

    char huge_header[16384];
    memset(huge_header, 'X', sizeof(huge_header));
    huge_header[0] = 'G';
    huge_header[1] = 'E';
    huge_header[2] = 'T';
    huge_header[3] = ' ';
    huge_header[sizeof(huge_header)-1] = '\0';
    write(sockfd[1], huge_header, sizeof(huge_header) - 1);

    struct epoll_event event = { .events = EPOLLIN, .data.fd = sockfd[0] };
    connection_handle_event(&event);
    TEST_ASSERT(1, "Oversized header handled");

    if (conn->type != CONN_TYPE_UNUSED) {
        buffer_free(&conn->read_buf);
        buffer_free(&conn->write_buf);
    }
    close(sockfd[0]);
    close(sockfd[1]);
    close(epoll_fd);
    epoll_fd = old_epoll;

    TEST_SUITE_END();
}

void test_connection_handle_empty_host(void) {
    TEST_SUITE_BEGIN("Connection Handle Empty Host");

    connection_init_all();
    connection_update_cached_time();
    int old_epoll = epoll_fd;
    epoll_fd = epoll_create1(0);
    TEST_ASSERT(epoll_fd >= 0, "Epoll created");

    int sockfd[2];
    socketpair(AF_UNIX, SOCK_STREAM, 0, sockfd);

    connection_set_non_blocking(sockfd[0]);
    connection_set_non_blocking(sockfd[1]);

    connection_t *conn = &connections[sockfd[0]];
    memset(conn, 0, sizeof(connection_t));
    conn->fd = sockfd[0];
    conn->type = CONN_TYPE_CLIENT;
    conn->state = CLIENT_STATE_READING_HEADERS;
    buffer_init(&conn->read_buf, 4096);
    buffer_init(&conn->write_buf, 4096);

    struct epoll_event ev = { .events = EPOLLIN, .data.fd = sockfd[0] };
    epoll_ctl(epoll_fd, EPOLL_CTL_ADD, sockfd[0], &ev);
    conn->epoll_events = EPOLLIN;

    const char *request = "GET / HTTP/1.1\r\n\r\n";
    write(sockfd[1], request, strlen(request));

    struct epoll_event event = { .events = EPOLLIN, .data.fd = sockfd[0] };
    connection_handle_event(&event);
    TEST_ASSERT(conn->state == CLIENT_STATE_ERROR || conn->write_buf.tail > 0 || conn->type == CONN_TYPE_UNUSED, "Empty host handled with error");

    if (conn->type != CONN_TYPE_UNUSED) {
        buffer_free(&conn->read_buf);
        buffer_free(&conn->write_buf);
    }
    close(sockfd[0]);
    close(sockfd[1]);
    close(epoll_fd);
    epoll_fd = old_epoll;

    TEST_SUITE_END();
}

void test_connection_handle_malformed_request(void) {
    TEST_SUITE_BEGIN("Connection Handle Malformed Request");

    connection_init_all();
    connection_update_cached_time();
    int old_epoll = epoll_fd;
    epoll_fd = epoll_create1(0);
    TEST_ASSERT(epoll_fd >= 0, "Epoll created");

    int sockfd[2];
    socketpair(AF_UNIX, SOCK_STREAM, 0, sockfd);

    connection_set_non_blocking(sockfd[0]);
    connection_set_non_blocking(sockfd[1]);

    connection_t *conn = &connections[sockfd[0]];
    memset(conn, 0, sizeof(connection_t));
    conn->fd = sockfd[0];
    conn->type = CONN_TYPE_CLIENT;
    conn->state = CLIENT_STATE_READING_HEADERS;
    buffer_init(&conn->read_buf, 4096);
    buffer_init(&conn->write_buf, 4096);

    struct epoll_event ev = { .events = EPOLLIN, .data.fd = sockfd[0] };
    epoll_ctl(epoll_fd, EPOLL_CTL_ADD, sockfd[0], &ev);
    conn->epoll_events = EPOLLIN;

    const char *request = "INVALID REQUEST\r\n\r\n";
    write(sockfd[1], request, strlen(request));

    struct epoll_event event = { .events = EPOLLIN, .data.fd = sockfd[0] };
    connection_handle_event(&event);
    TEST_ASSERT(conn->state == CLIENT_STATE_ERROR || conn->write_buf.tail > 0 || conn->type == CONN_TYPE_UNUSED, "Malformed request handled");

    if (conn->type != CONN_TYPE_UNUSED) {
        buffer_free(&conn->read_buf);
        buffer_free(&conn->write_buf);
    }
    close(sockfd[0]);
    close(sockfd[1]);
    close(epoll_fd);
    epoll_fd = old_epoll;

    TEST_SUITE_END();
}

void test_connection_cleanup_preserves_listener(void) {
    TEST_SUITE_BEGIN("Connection Cleanup Preserves Listener");

    connection_init_all();
    int old_epoll = epoll_fd;
    epoll_fd = epoll_create1(0);
    TEST_ASSERT(epoll_fd >= 0, "Epoll created");

    connections[10].type = CONN_TYPE_LISTENER;
    connections[10].fd = 10;
    connections[10].last_activity = time(NULL) - 1000;

    connection_cleanup_idle();

    TEST_ASSERT_EQ(CONN_TYPE_LISTENER, connections[10].type, "Listener preserved");

    connections[10].type = CONN_TYPE_UNUSED;
    connections[10].fd = -1;

    close(epoll_fd);
    epoll_fd = old_epoll;

    TEST_SUITE_END();
}

void test_connection_setup_listener(void) {
    TEST_SUITE_BEGIN("Connection Setup Listener");

    connection_init_all();
    int old_epoll = epoll_fd;
    epoll_fd = epoll_create1(0);
    TEST_ASSERT(epoll_fd >= 0, "Epoll created");

    connection_setup_listener(0);

    int found_listener = 0;
    for (int i = 0; i < MAX_FDS; i++) {
        if (connections[i].type == CONN_TYPE_LISTENER) {
            found_listener = 1;
            connection_close(i);
            break;
        }
    }
    TEST_ASSERT(found_listener, "Listener created on port 0");

    close(epoll_fd);
    epoll_fd = old_epoll;

    TEST_SUITE_END();
}

void test_connection_handle_forwarding_state(void) {
    TEST_SUITE_BEGIN("Connection Handle Forwarding State");

    connection_init_all();
    connection_update_cached_time();
    int old_epoll = epoll_fd;
    epoll_fd = epoll_create1(0);
    TEST_ASSERT(epoll_fd >= 0, "Epoll created");

    int sockfd1[2], sockfd2[2];
    socketpair(AF_UNIX, SOCK_STREAM, 0, sockfd1);
    socketpair(AF_UNIX, SOCK_STREAM, 0, sockfd2);

    connection_set_non_blocking(sockfd1[0]);
    connection_set_non_blocking(sockfd1[1]);
    connection_set_non_blocking(sockfd2[0]);
    connection_set_non_blocking(sockfd2[1]);

    connection_t *client = &connections[sockfd1[0]];
    memset(client, 0, sizeof(connection_t));
    client->fd = sockfd1[0];
    client->type = CONN_TYPE_CLIENT;
    client->state = CLIENT_STATE_FORWARDING;
    buffer_init(&client->read_buf, 4096);
    buffer_init(&client->write_buf, 4096);
    client->splice_pipe[0] = -1;
    client->splice_pipe[1] = -1;

    connection_t *upstream = &connections[sockfd2[0]];
    memset(upstream, 0, sizeof(connection_t));
    upstream->fd = sockfd2[0];
    upstream->type = CONN_TYPE_UPSTREAM;
    upstream->state = CLIENT_STATE_FORWARDING;
    buffer_init(&upstream->read_buf, 4096);
    buffer_init(&upstream->write_buf, 4096);
    upstream->splice_pipe[0] = -1;
    upstream->splice_pipe[1] = -1;

    client->pair = upstream;
    upstream->pair = client;

    struct epoll_event ev1 = { .events = EPOLLIN, .data.fd = sockfd1[0] };
    struct epoll_event ev2 = { .events = EPOLLIN, .data.fd = sockfd2[0] };
    epoll_ctl(epoll_fd, EPOLL_CTL_ADD, sockfd1[0], &ev1);
    epoll_ctl(epoll_fd, EPOLL_CTL_ADD, sockfd2[0], &ev2);
    client->epoll_events = EPOLLIN;
    upstream->epoll_events = EPOLLIN;

    const char *response = "HTTP/1.1 200 OK\r\nContent-Length: 2\r\n\r\nOK";
    write(sockfd2[1], response, strlen(response));

    struct epoll_event event = { .events = EPOLLIN, .data.fd = sockfd2[0] };
    connection_handle_event(&event);
    TEST_ASSERT(1, "Forwarding from upstream handled");

    if (client->type != CONN_TYPE_UNUSED) {
        buffer_free(&client->read_buf);
        buffer_free(&client->write_buf);
    }
    if (upstream->type != CONN_TYPE_UNUSED) {
        buffer_free(&upstream->read_buf);
        buffer_free(&upstream->write_buf);
    }
    close(sockfd1[0]);
    close(sockfd1[1]);
    close(sockfd2[0]);
    close(sockfd2[1]);
    close(epoll_fd);
    epoll_fd = old_epoll;

    TEST_SUITE_END();
}

void test_connection_handle_upstream_eof(void) {
    TEST_SUITE_BEGIN("Connection Handle Upstream EOF");

    connection_init_all();
    connection_update_cached_time();
    int old_epoll = epoll_fd;
    epoll_fd = epoll_create1(0);
    TEST_ASSERT(epoll_fd >= 0, "Epoll created");

    int sockfd1[2], sockfd2[2];
    socketpair(AF_UNIX, SOCK_STREAM, 0, sockfd1);
    socketpair(AF_UNIX, SOCK_STREAM, 0, sockfd2);

    connection_set_non_blocking(sockfd1[0]);
    connection_set_non_blocking(sockfd2[0]);

    connection_t *client = &connections[sockfd1[0]];
    memset(client, 0, sizeof(connection_t));
    client->fd = sockfd1[0];
    client->type = CONN_TYPE_CLIENT;
    client->state = CLIENT_STATE_FORWARDING;
    buffer_init(&client->read_buf, 4096);
    buffer_init(&client->write_buf, 4096);
    client->splice_pipe[0] = -1;
    client->splice_pipe[1] = -1;

    connection_t *upstream = &connections[sockfd2[0]];
    memset(upstream, 0, sizeof(connection_t));
    upstream->fd = sockfd2[0];
    upstream->type = CONN_TYPE_UPSTREAM;
    upstream->state = CLIENT_STATE_FORWARDING;
    buffer_init(&upstream->read_buf, 4096);
    buffer_init(&upstream->write_buf, 4096);
    upstream->splice_pipe[0] = -1;
    upstream->splice_pipe[1] = -1;

    client->pair = upstream;
    upstream->pair = client;

    struct epoll_event ev1 = { .events = EPOLLIN, .data.fd = sockfd1[0] };
    struct epoll_event ev2 = { .events = EPOLLIN, .data.fd = sockfd2[0] };
    epoll_ctl(epoll_fd, EPOLL_CTL_ADD, sockfd1[0], &ev1);
    epoll_ctl(epoll_fd, EPOLL_CTL_ADD, sockfd2[0], &ev2);
    client->epoll_events = EPOLLIN;
    upstream->epoll_events = EPOLLIN;

    close(sockfd2[1]);

    struct epoll_event event = { .events = EPOLLIN, .data.fd = sockfd2[0] };
    connection_handle_event(&event);
    TEST_ASSERT(1, "Upstream EOF handled");

    for (int i = 0; i < MAX_FDS; i++) {
        if (connections[i].type != CONN_TYPE_UNUSED) {
            connection_close(i);
        }
    }
    close(sockfd1[1]);
    close(epoll_fd);
    epoll_fd = old_epoll;

    TEST_SUITE_END();
}

void test_connection_handle_client_forwarding_read(void) {
    TEST_SUITE_BEGIN("Connection Handle Client Forwarding Read");

    connection_init_all();
    connection_update_cached_time();
    int old_epoll = epoll_fd;
    epoll_fd = epoll_create1(0);
    TEST_ASSERT(epoll_fd >= 0, "Epoll created");

    int sockfd1[2], sockfd2[2];
    socketpair(AF_UNIX, SOCK_STREAM, 0, sockfd1);
    socketpair(AF_UNIX, SOCK_STREAM, 0, sockfd2);

    connection_set_non_blocking(sockfd1[0]);
    connection_set_non_blocking(sockfd1[1]);
    connection_set_non_blocking(sockfd2[0]);
    connection_set_non_blocking(sockfd2[1]);

    connection_t *client = &connections[sockfd1[0]];
    memset(client, 0, sizeof(connection_t));
    client->fd = sockfd1[0];
    client->type = CONN_TYPE_CLIENT;
    client->state = CLIENT_STATE_FORWARDING;
    buffer_init(&client->read_buf, 4096);
    buffer_init(&client->write_buf, 4096);
    client->splice_pipe[0] = -1;
    client->splice_pipe[1] = -1;

    connection_t *upstream = &connections[sockfd2[0]];
    memset(upstream, 0, sizeof(connection_t));
    upstream->fd = sockfd2[0];
    upstream->type = CONN_TYPE_UPSTREAM;
    upstream->state = CLIENT_STATE_FORWARDING;
    buffer_init(&upstream->read_buf, 4096);
    buffer_init(&upstream->write_buf, 4096);
    upstream->splice_pipe[0] = -1;
    upstream->splice_pipe[1] = -1;

    client->pair = upstream;
    upstream->pair = client;

    struct epoll_event ev1 = { .events = EPOLLIN, .data.fd = sockfd1[0] };
    struct epoll_event ev2 = { .events = EPOLLIN, .data.fd = sockfd2[0] };
    epoll_ctl(epoll_fd, EPOLL_CTL_ADD, sockfd1[0], &ev1);
    epoll_ctl(epoll_fd, EPOLL_CTL_ADD, sockfd2[0], &ev2);
    client->epoll_events = EPOLLIN;
    upstream->epoll_events = EPOLLIN;

    const char *data = "More client data\r\n";
    write(sockfd1[1], data, strlen(data));

    struct epoll_event event = { .events = EPOLLIN, .data.fd = sockfd1[0] };
    connection_handle_event(&event);
    TEST_ASSERT(1, "Client data forwarded to upstream");

    if (client->type != CONN_TYPE_UNUSED) {
        buffer_free(&client->read_buf);
        buffer_free(&client->write_buf);
    }
    if (upstream->type != CONN_TYPE_UNUSED) {
        buffer_free(&upstream->read_buf);
        buffer_free(&upstream->write_buf);
    }
    close(sockfd1[0]);
    close(sockfd1[1]);
    close(sockfd2[0]);
    close(sockfd2[1]);
    close(epoll_fd);
    epoll_fd = old_epoll;

    TEST_SUITE_END();
}

void test_connection_handle_write_complete_close(void) {
    TEST_SUITE_BEGIN("Connection Handle Write Complete Close");

    connection_init_all();
    connection_update_cached_time();
    int old_epoll = epoll_fd;
    epoll_fd = epoll_create1(0);
    TEST_ASSERT(epoll_fd >= 0, "Epoll created");

    int sockfd[2];
    socketpair(AF_UNIX, SOCK_STREAM, 0, sockfd);

    connection_set_non_blocking(sockfd[0]);
    connection_set_non_blocking(sockfd[1]);

    connection_t *conn = &connections[sockfd[0]];
    memset(conn, 0, sizeof(connection_t));
    conn->fd = sockfd[0];
    conn->type = CONN_TYPE_CLIENT;
    conn->state = CLIENT_STATE_CLOSING;
    buffer_init(&conn->read_buf, 4096);
    buffer_init(&conn->write_buf, 4096);
    conn->splice_pipe[0] = -1;
    conn->splice_pipe[1] = -1;

    struct epoll_event ev = { .events = EPOLLIN | EPOLLOUT, .data.fd = sockfd[0] };
    epoll_ctl(epoll_fd, EPOLL_CTL_ADD, sockfd[0], &ev);
    conn->epoll_events = EPOLLIN | EPOLLOUT;

    struct epoll_event event = { .events = EPOLLOUT, .data.fd = sockfd[0] };
    connection_handle_event(&event);
    TEST_ASSERT(1, "Write complete handled without crash");

    if (conn->type != CONN_TYPE_UNUSED) {
        buffer_free(&conn->read_buf);
        buffer_free(&conn->write_buf);
        close(sockfd[0]);
    }
    close(sockfd[1]);
    close(epoll_fd);
    epoll_fd = old_epoll;

    TEST_SUITE_END();
}

void test_connection_handle_error_state(void) {
    TEST_SUITE_BEGIN("Connection Handle Error State");

    connection_init_all();
    connection_update_cached_time();
    int old_epoll = epoll_fd;
    epoll_fd = epoll_create1(0);
    TEST_ASSERT(epoll_fd >= 0, "Epoll created");

    int sockfd[2];
    socketpair(AF_UNIX, SOCK_STREAM, 0, sockfd);

    connection_set_non_blocking(sockfd[0]);
    connection_set_non_blocking(sockfd[1]);

    connection_t *conn = &connections[sockfd[0]];
    memset(conn, 0, sizeof(connection_t));
    conn->fd = sockfd[0];
    conn->type = CONN_TYPE_CLIENT;
    conn->state = CLIENT_STATE_ERROR;
    buffer_init(&conn->read_buf, 4096);
    buffer_init(&conn->write_buf, 4096);
    conn->splice_pipe[0] = -1;
    conn->splice_pipe[1] = -1;

    const char *response = "HTTP/1.1 500 Internal Server Error\r\n\r\n";
    memcpy(conn->write_buf.data, response, strlen(response));
    conn->write_buf.tail = strlen(response);

    struct epoll_event ev = { .events = EPOLLIN | EPOLLOUT, .data.fd = sockfd[0] };
    epoll_ctl(epoll_fd, EPOLL_CTL_ADD, sockfd[0], &ev);
    conn->epoll_events = EPOLLIN | EPOLLOUT;

    struct epoll_event event = { .events = EPOLLOUT, .data.fd = sockfd[0] };
    connection_handle_event(&event);
    TEST_ASSERT(1, "Error state write handled");

    if (conn->type != CONN_TYPE_UNUSED) {
        buffer_free(&conn->read_buf);
        buffer_free(&conn->write_buf);
        close(sockfd[0]);
    }
    close(sockfd[1]);
    close(epoll_fd);
    epoll_fd = old_epoll;

    TEST_SUITE_END();
}

void test_connection_cleanup_active_conn(void) {
    TEST_SUITE_BEGIN("Connection Cleanup Active Connection");

    connection_init_all();
    connection_update_cached_time();
    int old_epoll = epoll_fd;
    epoll_fd = epoll_create1(0);
    TEST_ASSERT(epoll_fd >= 0, "Epoll created");

    int sockfd[2];
    socketpair(AF_UNIX, SOCK_STREAM, 0, sockfd);

    connection_t *conn = &connections[sockfd[0]];
    memset(conn, 0, sizeof(connection_t));
    conn->fd = sockfd[0];
    conn->type = CONN_TYPE_CLIENT;
    conn->state = CLIENT_STATE_READING_HEADERS;
    conn->last_activity = time(NULL);
    buffer_init(&conn->read_buf, 4096);
    buffer_init(&conn->write_buf, 4096);
    conn->splice_pipe[0] = -1;
    conn->splice_pipe[1] = -1;

    struct epoll_event ev = { .events = EPOLLIN, .data.fd = sockfd[0] };
    epoll_ctl(epoll_fd, EPOLL_CTL_ADD, sockfd[0], &ev);

    connection_cleanup_idle();

    TEST_ASSERT_EQ(CONN_TYPE_CLIENT, conn->type, "Active connection preserved");

    buffer_free(&conn->read_buf);
    buffer_free(&conn->write_buf);
    close(sockfd[0]);
    close(sockfd[1]);
    close(epoll_fd);
    epoll_fd = old_epoll;

    TEST_SUITE_END();
}

void run_connection_tests(void) {
    test_connection_init_all();
    test_connection_set_non_blocking();
    test_connection_close_unused();
    test_connection_do_read_null();
    test_connection_do_write_null();
    test_connection_do_read_basic();
    test_connection_do_write_basic();
    test_connection_tcp_keepalive();
    test_connection_buffer_operations();
    test_connection_send_error_response();
    test_connection_epoll_operations();
    test_connection_send_auth_required();
    test_connection_close_active();
    test_connection_cleanup_idle();
    test_connection_do_read_eof();
    test_connection_do_write_full_buffer();
    test_connection_multiple_error_responses();
    test_connection_do_read_eagain();
    test_connection_optimize_socket();
    test_connection_cached_time();
    test_connection_epoll_event_caching();
    test_connection_splice_pipe_init();
    test_connection_patch_buf_cleanup();
    test_connection_tcp_nodelay();
    test_connection_close_with_upstream_pair();
    test_connection_close_client_closes_upstream();
    test_connection_handle_event_listener();
    test_connection_handle_event_error();
    test_connection_handle_event_hup();
    test_connection_handle_event_invalid_fd();
    test_connection_handle_event_unused();
    test_connection_cleanup_idle_with_idle_conn();
    test_connection_send_error_null_params();
    test_connection_send_auth_null_conn();
    test_connection_buffer_compaction();
    test_connection_do_write_empty_buffer();
    test_connection_modify_epoll_same_events();
    test_connection_accept_new();
    test_connection_handle_event_write();
    test_connection_close_with_vhost_stats();
    test_connection_close_with_config();
    test_connection_send_error_various_codes();
    test_connection_send_auth_null_realm();
    test_connection_do_read_grow_buffer();
    test_connection_close_with_splice_pipes();
    test_connection_modify_epoll_invalid();
    test_connection_add_epoll_failure();
    test_connection_connect_upstream_no_config();
    test_connection_connect_upstream_null_params();
    test_connection_handle_event_client_read();
    test_connection_handle_complete_request();
    test_connection_handle_empty_host();
    test_connection_handle_malformed_request();
    test_connection_cleanup_preserves_listener();
    test_connection_setup_listener();
    test_connection_handle_forwarding_state();
    test_connection_handle_upstream_eof();
    test_connection_handle_client_forwarding_read();
    test_connection_handle_write_complete_close();
    test_connection_handle_error_state();
    test_connection_cleanup_active_conn();
}