cpp_server/server.hpp

#pragma once

#include "epoll_socket.hpp"
#include "router.hpp"
#include "parser.hpp"
#include "response.hpp"
#include "static_file_handler.hpp"
#include "kv_store.hpp"
#include "session_store.hpp"
#include <iostream>
#include <string.h>
#include <string_view>
#include <array>
#include <string>
#include <thread>
#include <vector>
#include <memory>

class Server
{
public:
	KeyValueStore store;
	SessionStore sessions;

	explicit Server(uint16_t port, Router& router) : port_(port), router_(router) {}

	~Server()
	{
		stop();
		store.save();
		sessions.save();
		// Wait for all worker threads to finish
		for (auto& worker : workers_) {
			if (worker->thread.joinable()) {
				worker->thread.join();
			}
		}
	}

	bool start()
	{
		// Load persistent data
		store.load();
		sessions.load();

		// Create one worker per CPU core for optimal performance
		unsigned int num_cores = std::thread::hardware_concurrency();
		if (num_cores == 0) num_cores = 1;

		workers_.reserve(num_cores);

		// Initialize all worker sockets (SO_REUSEPORT allows this)
		for (unsigned int i = 0; i < num_cores; ++i) {
			auto worker = std::make_unique<Worker>(port_, router_, static_handler_, sessions);
			if (!worker->socket.start()) return false;
			workers_.push_back(std::move(worker));
		}

		// Start worker threads
		for (auto& worker : workers_) {
			worker->thread = std::thread([&worker]() { worker->socket.run(); });
		}

		return true;
	}

	void run()
	{
		// Wait for all workers to complete
		for (auto& worker : workers_) {
			if (worker->thread.joinable()) {
				worker->thread.join();
			}
		}
	}

	void stop()
	{
		// Signal all workers to stop
		for (auto& worker : workers_) {
			worker->socket.stop();
		}
	}

	// Utility function to extract path segments
	static std::string get_path_param(string_view path, size_t segment_index = 0)
	{
		size_t start = 0;
		size_t current_segment = 0;

		while (start < path.length()) {
			if (path[start] == '/') {
				start++;
				continue;
			}

			size_t end = path.find('/', start);
			if (end == std::string_view::npos) end = path.length();

			if (current_segment == segment_index) {
				return std::string(path.substr(start, end - start));
			}

			current_segment++;
			start = end;
		}
		return "";
	}

	// Enable static file serving
	void serve_static(const std::string& static_dir, const std::string& url_prefix = "")
	{
		static_handler_ = std::make_shared<StaticFileHandler>(static_dir, url_prefix);
	}

private:
	static constexpr int BUFFER_SIZE = 65536;
	std::shared_ptr<StaticFileHandler> static_handler_;

	// Worker handles requests in a dedicated thread
	struct Worker
	{
		EpollSocket socket;
		Router& router;
		std::shared_ptr<StaticFileHandler>& static_handler;
		SessionStore& sessions;
		std::array<char, BUFFER_SIZE> buffer;
		std::string request_accumulator; // For handling partial requests
		std::thread thread;

		Worker(uint16_t port, Router& r, std::shared_ptr<StaticFileHandler>& sh, SessionStore& s)
			: socket(port), router(r), static_handler(sh), sessions(s)
		{
			// Set up event handlers for the new epoll socket library
			socket.on_connection([this](int fd) { handle_connection(fd); });
			socket.on_data([this](int fd) { handle_data(fd); });
			socket.on_disconnect([this](int fd) { handle_disconnect(fd); });
		}

		void handle_connection(int client_fd)
		{
			// New client connected - no action needed with new library
			// The library handles connection tracking automatically
		}

		void handle_data(int client_fd)
		{
			// Read all available data using edge-triggered epoll
			while (true) {
				ssize_t bytes_read = read(client_fd, buffer.data(), buffer.size());

				if (bytes_read == -1) {
					if (errno == EAGAIN || errno == EWOULDBLOCK) {
						// No more data available
						break;
					}
					// Read error - connection will be closed by epoll library
					return;
				}

				if (bytes_read == 0) {
					// EOF - client closed connection
					return;
				}

				// Accumulate request data
				request_accumulator.append(buffer.data(), bytes_read);

				// Check if we have a complete HTTP request
				size_t header_end = request_accumulator.find("\r\n\r\n");
				if (header_end != std::string::npos) {
					// Process the complete request
					process_request(client_fd, request_accumulator);
					// Clear accumulator for next request (HTTP keep-alive)
					request_accumulator.clear();
					return;
				}

				// Prevent memory exhaustion from malicious large headers
				if (request_accumulator.size() > 64 * 1024) {
					send_error_response(client_fd, "Request Header Too Large", 431);
					request_accumulator.clear();
					return;
				}
			}
		}

		void handle_disconnect(int client_fd)
		{
			// Client disconnected - cleanup handled by epoll library
			// Clear any partial request data for this connection
			request_accumulator.clear();
		}

		void process_request(int client_fd, std::string_view request_data)
		{
			// Parse the HTTP request
			Request req = Parser::parse(request_data);

			if (!req.valid) {
				send_error_response(client_fd, "Bad Request", 400);
				return;
			}

			Response response;

			// Try to handle with router first
			if (router.handle(req, response)) {
				// Handle session management
				std::string_view existing_id = sessions.extract_session_id(req);
				std::string session_id = existing_id.empty() ?
					sessions.create() : std::string(existing_id);

				set_session_cookie(response, session_id);
				send_response(client_fd, response, req.version);
				return;
			}

			// Try static file handler if available
			if (static_handler && static_handler->handle(req, response)) {
				send_response(client_fd, response, req.version);
				return;
			}

			// No handler found - return 404
			response.status = 404;
			response.set_text("Not Found");
			send_response(client_fd, response, req.version);
		}

		void set_session_cookie(Response& response, const std::string& session_id)
		{
			// Set secure session cookie
			response.cookies.push_back("session_id=" + session_id +
				"; HttpOnly; Path=/; SameSite=Strict");
		}

		void send_response(int client_fd, const Response& response, string_view version)
		{
			// Build HTTP response string
			std::string http_response = ResponseBuilder::build_response(response, version);

			// Use the new epoll library's buffered send method
			if (!socket.send_data(client_fd, http_response.data(), http_response.size())) {
				// Send failed - connection probably closed
				// Library will handle cleanup automatically
			}
		}

		void send_error_response(int client_fd, const std::string& message, int status)
		{
			// Fast path for error responses
			std::string response = ResponseBuilder::build_error_response(status, message);

			// Use the new epoll library's send method
			socket.send_data(client_fd, response.data(), response.size());
		}
	};

	uint16_t port_;
	Router& router_;
	std::vector<std::unique_ptr<Worker>> workers_;
};