connection_handler.cpp

// Author: Suchismita Bose (psychomita)
// Year: 2025
// License: MIT
// This file is part of the HTTP Load Balancer project.
// You may use, modify, and distribute this file under the terms of the MIT License.

#include "connection_handler.h"
#include <iostream>
#include <sstream>
#include <iomanip>
#include <chrono>
#include <algorithm>
#include <unistd.h>
#include <poll.h>
#include <sys/socket.h>
#include <cstring>
#include <limits>

ConnectionHandler::ConnectionHandler(int health_check_interval, bool enable_logging)
    : total_requests_(0), health_check_interval_(health_check_interval), enable_logging_(enable_logging) {}

void ConnectionHandler::add_backend_server(const std::string& ip, int port, int weight, int max_connections) {
    std::lock_guard<std::mutex> lock(servers_mutex_);
    backend_servers_.push_back(std::make_unique<BackendServer>(ip, port, weight, max_connections));

    if (enable_logging_) {
        std::cout << "[CONFIG] Added backend server: " << ip << ":" << port
                  << " (weight: " << weight << ", max_conn: " << max_connections << ")\n";
    }
}

size_t ConnectionHandler::get_backend_count() const {
    std::lock_guard<std::mutex> lock(servers_mutex_);
    return backend_servers_.size();
}

BackendServer* ConnectionHandler::select_least_loaded_server() {
    std::lock_guard<std::mutex> lock(servers_mutex_);

    if (backend_servers_.empty()) return nullptr;

    BackendServer* best_server = nullptr;
    double lowest_load = std::numeric_limits<double>::max();

    for (auto& server : backend_servers_) {
        if (server->can_accept_connection()) {
            double load_score = server->get_load_score();
            if (load_score < lowest_load) {
                lowest_load = load_score;
                best_server = server.get();
            }
        }
    }

    return best_server;
}

void ConnectionHandler::handle_client_connection(int client_socket, const std::string& client_ip) {
    long request_id = ++total_requests_;

    BackendServer* selected_server = select_least_loaded_server();
    if (!selected_server) {
        if (enable_logging_) {
            std::cout << "[ERROR] No available backend servers for request #" << request_id << "\n";
        }
        send_error_response(client_socket, "503 Service Unavailable", "No backend servers available");
        close(client_socket);
        return;
    }

    selected_server->increment_connections();

    if (enable_logging_) {
        log_request_routing(request_id, client_ip, *selected_server);
    }

    int backend_socket = selected_server->connect_to_server();
    if (backend_socket < 0) {
        if (enable_logging_) {
            std::cout << "[ERROR] Failed to connect to backend " << selected_server->to_string()
                      << " for request #" << request_id << "\n";
        }
        send_error_response(client_socket, "502 Bad Gateway", "Backend server connection failed");
        selected_server->decrement_connections();
        close(client_socket);
        return;
    }

    proxy_data_between_sockets(client_socket, backend_socket);

    selected_server->decrement_connections();
    close(client_socket);
    close(backend_socket);

    if (enable_logging_) {
        std::cout << "[COMPLETE] Request #" << request_id << " completed\n";
    }
}

void ConnectionHandler::perform_health_checks() const {
    std::lock_guard<std::mutex> lock(servers_mutex_);

    if (enable_logging_) {
        std::cout << "[HEALTH] Performing health checks...\n";
    }

    for (auto& server : backend_servers_) {
        if (server->needs_health_check(health_check_interval_)) {
            bool is_healthy = test_server_health(*server);
            server->set_health_status(is_healthy);
            server->update_last_health_check();

            if (enable_logging_) {
                std::cout << "[HEALTH] " << server->to_string()
                          << " - " << (is_healthy ? "HEALTHY" : "UNHEALTHY") << "\n";
            }
        }
    }
}

bool ConnectionHandler::has_healthy_servers() const {
    std::lock_guard<std::mutex> lock(servers_mutex_);

    return std::any_of(backend_servers_.begin(), backend_servers_.end(),
                       [](const auto& server) { return server->is_healthy(); });
}

void ConnectionHandler::print_statistics() const {
    std::lock_guard<std::mutex> lock(servers_mutex_);

    std::cout << "\n" << std::string(60, '=') << "\n";
    std::cout << "LOAD BALANCER STATISTICS\n";
    std::cout << std::string(60, '=') << "\n";
    std::cout << "Total Requests Processed: " << total_requests_.load() << "\n";

    int healthy_count = std::count_if(backend_servers_.begin(), backend_servers_.end(),
                                      [](const auto& server) { return server->is_healthy(); });

    std::cout << "Healthy Backend Servers: " << healthy_count << "/" << backend_servers_.size() << "\n";
    std::cout << std::string(60, '-') << "\n";

    for (size_t i = 0; i < backend_servers_.size(); ++i) {
        const auto& server = backend_servers_[i];
        std::cout << "Backend #" << (i + 1) << ": " << server->to_string()
                  << " | Connections: " << server->get_active_connections()
                  << "/" << server->get_max_connections()
                  << " | Weight: " << server->get_weight()
                  << " | Load Score: " << std::fixed << std::setprecision(2) << server->get_load_score()
                  << " | Status: " << (server->is_healthy() ? "HEALTHY" : "DOWN") << "\n";
    }

    std::cout << std::string(60, '=') << "\n\n";
}

long ConnectionHandler::get_total_requests() const {
    return total_requests_.load();
}

void ConnectionHandler::log_request_routing(long request_id, const std::string& client_ip,
                                            const BackendServer& server) const {
    auto now = std::chrono::system_clock::now();
    auto time_t = std::chrono::system_clock::to_time_t(now);

    std::cout << "[ROUTE] Request #" << request_id
              << " from " << client_ip
              << " -> " << server.to_string()
              << " (connections: " << server.get_active_connections()
              << ", load: " << std::fixed << std::setprecision(2) << server.get_load_score() << ")"
              << " at " << std::put_time(std::localtime(&time_t), "%H:%M:%S") << "\n";
}

void ConnectionHandler::proxy_data_between_sockets(int client_socket, int backend_socket) {
    struct pollfd fds[2];
    fds[0].fd = client_socket;
    fds[0].events = POLLIN;
    fds[1].fd = backend_socket;
    fds[1].events = POLLIN;

    char buffer[8192];
    bool client_closed = false;
    bool backend_closed = false;

    while (!client_closed && !backend_closed) {
        int poll_result = poll(fds, 2, 5000); // 5 second timeout

        if (poll_result <= 0) break;

        if (fds[0].revents & POLLIN) {
            ssize_t bytes_read = recv(client_socket, buffer, sizeof(buffer), 0);
            if (bytes_read <= 0) {
                client_closed = true;
            } else {
                send_all_data(backend_socket, buffer, bytes_read);
            }
        }

        if (fds[1].revents & POLLIN) {
            ssize_t bytes_read = recv(backend_socket, buffer, sizeof(buffer), 0);
            if (bytes_read <= 0) {
                backend_closed = true;
            } else {
                send_all_data(client_socket, buffer, bytes_read);
            }
        }

        if (fds[0].revents & (POLLERR | POLLHUP)) client_closed = true;
        if (fds[1].revents & (POLLERR | POLLHUP)) backend_closed = true;
    }
}

void ConnectionHandler::send_all_data(int socket, const char* data, size_t length) {
    size_t total_sent = 0;
    while (total_sent < length) {
        ssize_t sent = send(socket, data + total_sent, length - total_sent, MSG_NOSIGNAL);
        if (sent <= 0) break;
        total_sent += sent;
    }
}

void ConnectionHandler::send_error_response(int socket, const std::string& status, const std::string& message) {
    std::ostringstream response;
    response << "HTTP/1.1 " << status << "\r\n"
             << "Content-Type: text/plain\r\n"
             << "Content-Length: " << message.length() << "\r\n"
             << "Connection: close\r\n"
             << "Server: HTTPLoadBalancer/1.0\r\n"
             << "\r\n"
             << message;

    std::string response_str = response.str();
    send_all_data(socket, response_str.c_str(), response_str.length());
}

bool ConnectionHandler::test_server_health(BackendServer& server) const {
    int test_socket = server.connect_to_server();
    if (test_socket >= 0) {
        close(test_socket);
        return true;
    }
    return false;
}