#pragma once

#include <sys/epoll.h>
#include <sys/socket.h>
#include <netinet/in.h>
#include <netinet/tcp.h>
#include <fcntl.h>
#include <unistd.h>
#include <errno.h>
#include <array>
#include <functional>
#include <chrono>
#include <unordered_map>
#include <vector>
#include <string> // Added for std::string

namespace sockeye {

class Socket
{
public:
	using ConnectionHandler = std::function<void(int client_fd)>;
	using DataHandler = std::function<void(int client_fd, const char* data, size_t len)>;
	using DisconnectHandler = std::function<void(int client_fd)>;

	explicit Socket(uint16_t port = 8080, int timeout_ms = 5000)
		: port_(port), timeout_ms_(timeout_ms) {}

	~Socket()
	{
		if (server_fd_ != -1) close(server_fd_);
		if (epoll_fd_ != -1) close(epoll_fd_);
	}

	bool start()
	{
		return create_server_socket() &&
		       create_epoll() &&
		       add_server_to_epoll();
	}

	void run()
	{
		std::array<epoll_event, MAX_EVENTS> events;

		while (running_) {
			int num_events = epoll_wait(epoll_fd_, events.data(), MAX_EVENTS, 1000);
			if (num_events == -1) {
				if (errno == EINTR) continue;
				break;
			}

			for (int i = 0; i < num_events; ++i) {
				if (events[i].data.fd == server_fd_) {
					accept_connections();
				} else {
					handle_client_data(events[i].data.fd);
				}
			}

			check_timeouts();
		}
	}

	void stop() { running_ = false; }

	// New: Send data to a client
	bool send(int client_fd, const std::string& data) {
		ssize_t total_sent = 0;
		const char* p_data = data.c_str();
		size_t len = data.length();

		while (total_sent < static_cast<ssize_t>(len)) {
			ssize_t sent = ::send(client_fd, p_data + total_sent, len - total_sent, MSG_NOSIGNAL);
			if (sent == -1) {
				if (errno == EAGAIN || errno == EWOULDBLOCK) {
					// Can't send right now, handle as an error for simplicity in this context
					return false;
				}
				// Other error
				return false;
			}
			total_sent += sent;
		}
		return true;
	}

	void on_connection(ConnectionHandler handler) { on_connection_ = std::move(handler); }
	void on_data(DataHandler handler) { on_data_ = std::move(handler); }
	void on_disconnect(DisconnectHandler handler) { on_disconnect_ = std::move(handler); }

private:
	struct Client {
		int fd;
		std::chrono::steady_clock::time_point last_activity;
	};

	static constexpr int MAX_EVENTS = 1024;
	static constexpr int BUFFER_SIZE = 8192;

	uint16_t port_;
	int timeout_ms_;
	int server_fd_ = -1;
	int epoll_fd_ = -1;
	bool running_ = true;

	std::unordered_map<int, Client> clients_;

	ConnectionHandler on_connection_;
	DataHandler on_data_;
	DisconnectHandler on_disconnect_;

	bool create_server_socket()
	{
		server_fd_ = socket(AF_INET, SOCK_STREAM | SOCK_NONBLOCK | SOCK_CLOEXEC, 0);
		if (server_fd_ == -1) return false;

		int opt = 1;
		if (setsockopt(server_fd_, SOL_SOCKET, SO_REUSEADDR, &opt, sizeof(opt)) == -1 ||
		    setsockopt(server_fd_, SOL_SOCKET, SO_REUSEPORT, &opt, sizeof(opt)) == -1 ||
		    setsockopt(server_fd_, IPPROTO_TCP, TCP_NODELAY, &opt, sizeof(opt)) == -1) {
			return false;
		}

		sockaddr_in addr{};
		addr.sin_family = AF_INET;
		addr.sin_addr.s_addr = INADDR_ANY;
		addr.sin_port = htons(port_);

		return bind(server_fd_, reinterpret_cast<sockaddr*>(&addr), sizeof(addr)) != -1 &&
		       listen(server_fd_, SOMAXCONN) != -1;
	}

	bool create_epoll()
	{
		epoll_fd_ = epoll_create1(EPOLL_CLOEXEC);
		return epoll_fd_ != -1;
	}

	bool add_server_to_epoll()
	{
		epoll_event event{};
		event.events = EPOLLIN | EPOLLET;
		event.data.fd = server_fd_;
		return epoll_ctl(epoll_fd_, EPOLL_CTL_ADD, server_fd_, &event) != -1;
	}

	bool add_client(int client_fd)
	{
		epoll_event event{};
		event.events = EPOLLIN | EPOLLET;
		event.data.fd = client_fd;
		if (epoll_ctl(epoll_fd_, EPOLL_CTL_ADD, client_fd, &event) == -1) {
			return false;
		}
		clients_[client_fd] = {client_fd, std::chrono::steady_clock::now()};
		return true;
	}

	void remove_client(int client_fd)
	{
		epoll_ctl(epoll_fd_, EPOLL_CTL_DEL, client_fd, nullptr);
		close(client_fd);
		clients_.erase(client_fd);
		if (on_disconnect_) on_disconnect_(client_fd);
	}

	inline void accept_connections()
	{
		while (true) {
			sockaddr_in client_addr{};
			socklen_t client_len = sizeof(client_addr);

			int client_fd = accept4(server_fd_,
				reinterpret_cast<sockaddr*>(&client_addr),
				&client_len,
				SOCK_NONBLOCK | SOCK_CLOEXEC);

			if (client_fd == -1) {
				if (errno == EAGAIN || errno == EWOULDBLOCK) break;
				continue;
			}

			int opt = 1;
			setsockopt(client_fd, IPPROTO_TCP, TCP_NODELAY, &opt, sizeof(opt));

			if (add_client(client_fd)) {
				if (on_connection_) on_connection_(client_fd);
			} else {
				close(client_fd);
			}
		}
	}

	inline void handle_client_data(int client_fd)
	{
		auto it = clients_.find(client_fd);
		if (it != clients_.end()) {
			it->second.last_activity = std::chrono::steady_clock::now();
		}

		if (!on_data_) return;

		char buffer[BUFFER_SIZE];
		while (true) {
			ssize_t bytes = recv(client_fd, buffer, BUFFER_SIZE, 0);
			if (bytes > 0) {
				on_data_(client_fd, buffer, bytes);
			} else if (bytes == 0) {
				remove_client(client_fd);
				break;
			} else {
				if (errno == EAGAIN || errno == EWOULDBLOCK) break;
				remove_client(client_fd);
				break;
			}
		}
	}

	void check_timeouts() {
		if (timeout_ms_ <= 0) return;

		auto now = std::chrono::steady_clock::now();
		std::vector<int> timed_out_clients;

		for (const auto& pair : clients_) {
			auto duration = std::chrono::duration_cast<std::chrono::milliseconds>(now - pair.second.last_activity);
			if (duration.count() > timeout_ms_) {
				timed_out_clients.push_back(pair.first);
			}
		}

		for (int client_fd : timed_out_clients) {
			remove_client(client_fd);
		}
	}
};

} // namespace sockeye