#include "test_framework.h"
#include "../src/types.h"
#include "../src/connection.h"
#include "../src/buffer.h"
#include "../src/monitor.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>

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 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();
}