#include "../lib/Socket.hpp"
#include <cassert>
#include <iostream>
#include <thread>
#include <chrono>
#include <cstring>
#include <poll.h>
#include <atomic> // Required for std::atomic

// Forward declarations of test functions
void test_socket_creation();
void test_socket_move();
void test_socket_bind_listen();
void test_socket_options();
void test_socket_connection();
void test_udp_socket();
void test_socket_shutdown();
void test_socket_error_handling();
void test_address_retrieval();
void test_self_connection_detection();


bool waitForSocketReady(int fd, short events, int timeout_ms)
{
	pollfd pfd;
	pfd.fd = fd;
	pfd.events = events;
	pfd.revents = 0;

	int result = poll(&pfd, 1, timeout_ms);
	return result > 0 && (pfd.revents & events);
}

void test_socket_creation()
{
	std::cout << "Testing socket creation...\n";

	auto tcp_socket = reactor::Socket::createTcp();
	assert(tcp_socket.fd() >= 0);

	auto tcp6_socket = reactor::Socket::createTcp(true);
	assert(tcp6_socket.fd() >= 0);

	auto udp_socket = reactor::Socket::createUdp();
	assert(udp_socket.fd() >= 0);

	auto udp6_socket = reactor::Socket::createUdp(true);
	assert(udp6_socket.fd() >= 0);

	std::cout << "✓ Socket creation passed\n";
}

void test_socket_move()
{
	std::cout << "Testing socket move semantics...\n";

	auto socket1 = reactor::Socket::createTcp();
	int fd = socket1.fd();
	assert(fd >= 0);

	auto socket2 = std::move(socket1);
	assert(socket2.fd() == fd);
	assert(socket1.fd() == -1);

	reactor::Socket socket3 = reactor::Socket::createTcp();
	socket3 = std::move(socket2);
	assert(socket3.fd() == fd);

	std::cout << "✓ Socket move semantics passed\n";
}

void test_socket_bind_listen()
{
	std::cout << "Testing socket bind and listen...\n";

	auto socket = reactor::Socket::createTcp();
	reactor::InetAddress addr(0);

	socket.setReuseAddr(true);
	socket.bind(addr);
	socket.listen();

	// FIX: Call getLocalAddr as a member function on the object.
	auto local_addr = socket.getLocalAddr();
	assert(local_addr.port() > 0);

	std::cout << "✓ Socket bind and listen passed\n";
}

void test_socket_options()
{
	std::cout << "Testing socket options...\n";

	auto socket = reactor::Socket::createTcp();

	socket.setReuseAddr(true);
	socket.setReusePort(true);
	socket.setTcpNoDelay(true);
	socket.setKeepAlive(true);
	socket.setTcpKeepAlive(7200, 30, 9);
	socket.setRecvBuffer(65536);
	socket.setSendBuffer(65536);

	std::cout << "✓ Socket options passed\n";
}


void test_socket_connection()
{
	std::cout << "Testing socket connection...\n";

	auto server_socket = reactor::Socket::createTcp();
	reactor::InetAddress server_addr(0);

	server_socket.setReuseAddr(true);
	server_socket.bind(server_addr);
	server_socket.listen();

	// FIX: Call getLocalAddr as a member function on the object.
	auto actual_addr = server_socket.getLocalAddr();
	std::cout << "Server listening on: " << actual_addr.toIpPort() << "\n";

	std::atomic<bool> server_done{false};
	std::thread server_thread([&server_socket, &server_done]() {
		reactor::InetAddress peer_addr;

		if (waitForSocketReady(server_socket.fd(), POLLIN, 1000)) {
			// FIX: Handle the std::optional<Socket> returned by accept().
			auto client_sock_opt = server_socket.accept(peer_addr);
			if (client_sock_opt) {
				auto &client_sock = *client_sock_opt;
				std::cout << "Server accepted connection from: " << peer_addr.toIpPort() << "\n";

				if (waitForSocketReady(client_sock.fd(), POLLIN, 1000)) {
					char buffer[1024];
					// FIX: Use the Socket object's read/write methods.
					ssize_t n = client_sock.read(buffer, sizeof(buffer));
					if (n > 0) {
						ssize_t written = client_sock.write(buffer, n);
						(void)written;
					}
				}
				// FIX: No need to call close(), Socket's destructor handles it.
			}
		}
		server_done = true;
	});

	std::this_thread::sleep_for(std::chrono::milliseconds(10));

	auto client_socket = reactor::Socket::createTcp();
	reactor::InetAddress connect_addr("127.0.0.1", actual_addr.port());

	int result = client_socket.connect(connect_addr);

	if (result < 0 && errno == EINPROGRESS) {
		if (waitForSocketReady(client_socket.fd(), POLLOUT, 1000)) {
			int error = client_socket.getSocketError();
			assert(error == 0);
		}
	}

	const char* message = "Hello Server";
	ssize_t sent = client_socket.write(message, strlen(message));
	assert(sent > 0);

	if (waitForSocketReady(client_socket.fd(), POLLIN, 1000)) {
		char response[1024];
		ssize_t received = client_socket.read(response, sizeof(response));
		assert(received > 0);
		assert(strncmp(message, response, sent) == 0);
	}

	server_thread.join();
	assert(server_done);
	std::cout << "✓ Socket connection passed\n";
}

void test_udp_socket()
{
	std::cout << "Testing UDP socket operations...\n";

	auto server_socket = reactor::Socket::createUdp();
	reactor::InetAddress server_addr(0);

	server_socket.setReuseAddr(true);
	server_socket.bind(server_addr);

	// FIX: Call getLocalAddr as a member function on the object.
	auto actual_addr = server_socket.getLocalAddr();
	std::cout << "UDP server bound to: " << actual_addr.toIpPort() << "\n";

	std::thread server_thread([&server_socket]() {
		if (waitForSocketReady(server_socket.fd(), POLLIN, 1000)) {
			char buffer[1024];
			reactor::InetAddress client_addr;
			ssize_t n = server_socket.recvFrom(buffer, sizeof(buffer), client_addr);
			if (n > 0) {
				server_socket.sendTo(buffer, n, client_addr);
			}
		}
	});

	std::this_thread::sleep_for(std::chrono::milliseconds(10));

	auto client_socket = reactor::Socket::createUdp();
	reactor::InetAddress target_addr("127.0.0.1", actual_addr.port());

	const char* message = "UDP Hello";
	ssize_t sent = client_socket.sendTo(message, strlen(message), target_addr);
	assert(sent > 0);

	if (waitForSocketReady(client_socket.fd(), POLLIN, 1000)) {
		char response[1024];
		reactor::InetAddress from_addr;
		ssize_t received = client_socket.recvFrom(response, sizeof(response), from_addr);
		assert(received > 0);
		assert(strncmp(message, response, sent) == 0);
	}

	server_thread.join();
	std::cout << "✓ UDP socket operations passed\n";
}

void test_socket_shutdown()
{
	std::cout << "Testing socket shutdown...\n";

	auto server_socket = reactor::Socket::createTcp();
	reactor::InetAddress server_addr(0);

	server_socket.setReuseAddr(true);
	server_socket.bind(server_addr);
	server_socket.listen();

	// FIX: Call getLocalAddr as a member function on the object.
	auto actual_addr = server_socket.getLocalAddr();

	std::thread server_thread([&server_socket]() {
		if (waitForSocketReady(server_socket.fd(), POLLIN, 1000)) {
			reactor::InetAddress peer_addr;
			// FIX: Handle the std::optional<Socket> returned by accept().
			auto client_sock_opt = server_socket.accept(peer_addr);
			if (client_sock_opt) {
				std::this_thread::sleep_for(std::chrono::milliseconds(20));
				client_sock_opt->shutdownWrite();
			}
		}
	});

	std::this_thread::sleep_for(std::chrono::milliseconds(10));

	auto client_socket = reactor::Socket::createTcp();
	reactor::InetAddress connect_addr("127.0.0.1", actual_addr.port());

	int result = client_socket.connect(connect_addr);
	if (result < 0 && errno == EINPROGRESS) {
		waitForSocketReady(client_socket.fd(), POLLOUT, 1000);
	}

	std::this_thread::sleep_for(std::chrono::milliseconds(50));

	char buffer[1024];
	ssize_t n = client_socket.read(buffer, sizeof(buffer));
	assert(n == 0);

	server_thread.join();
	std::cout << "✓ Socket shutdown passed\n";
}

void test_socket_error_handling()
{
	std::cout << "Testing socket error handling...\n";

	auto socket = reactor::Socket::createTcp();

	int error = socket.getSocketError();
	assert(error == 0);

	reactor::InetAddress invalid_addr("192.0.2.1", 12345);
	socket.connect(invalid_addr);

	std::this_thread::sleep_for(std::chrono::milliseconds(100));

	std::cout << "✓ Socket error handling passed\n";
}

void test_address_retrieval()
{
	std::cout << "Testing address retrieval...\n";

	auto server_socket = reactor::Socket::createTcp();
	reactor::InetAddress server_addr("127.0.0.1", 0);

	server_socket.setReuseAddr(true);
	server_socket.bind(server_addr);
	server_socket.listen();

	// FIX: Call getLocalAddr as a member function on the object.
	auto local_addr = server_socket.getLocalAddr();
	assert(local_addr.toIp() == "127.0.0.1");
	assert(local_addr.port() > 0);

	std::cout << "✓ Address retrieval passed\n";
}

void test_self_connection_detection()
{
	std::cout << "Testing self connection detection...\n";

	auto socket = reactor::Socket::createTcp();
	reactor::InetAddress addr("127.0.0.1", 0);

	socket.setReuseAddr(true);
	socket.bind(addr);

	// This is not a true self-connection test, as the socket isn't connected.
	// We call it just to ensure it doesn't crash.
	bool is_self = socket.isSelfConnected();
	std::cout << "Self connected (on non-connected socket): " << std::boolalpha << is_self << "\n";

	std::cout << "✓ Self connection detection passed\n";
}

int main()
{
	try {
		std::cout << "=== Socket Tests ===\n";

		test_socket_creation();
		test_socket_move();
		test_socket_bind_listen();
		test_socket_options();
		test_socket_connection();
		test_udp_socket();
		test_socket_shutdown();
		test_socket_error_handling();
		test_address_retrieval();
		test_self_connection_detection();

		std::cout << "\nAll socket tests passed! ✓\n";
		return 0;
	} catch (const std::exception& e) {
		std::cerr << "A test failed with an exception: " << e.what() << std::endl;
		return 1;
	}
}