#include <chrono>
#include <sstream>
#include <string>
#include <thread>

#include <gtest/gtest.h>

#include <Socket.h>
#include <SocketListener.h>
#include <SocketAddress.h>

using namespace std::literals::chrono_literals;

using namespace address;

/* --------------------------------------------------------
 * Miscellaneous
 * -------------------------------------------------------- */

TEST(Misc, set)
{
	try {
		Socket s(AF_INET6, SOCK_STREAM, 0);

		s.set(IPPROTO_IPV6, IPV6_V6ONLY, false);
		ASSERT_FALSE(s.get<bool>(IPPROTO_IPV6, IPV6_V6ONLY));

		s.set(IPPROTO_IPV6, IPV6_V6ONLY, true);
		ASSERT_TRUE(s.get<bool>(IPPROTO_IPV6, IPV6_V6ONLY));
	} catch (const std::exception &ex) {
	}
}

/* --------------------------------------------------------
 * Select tests
 * -------------------------------------------------------- */

TEST(ListenerMethodSelect, add)
{
	try {
		Socket s(AF_INET, SOCK_STREAM, 0);
		Socket s2(AF_INET, SOCK_STREAM, 0);
		SocketListener listener(SocketMethod::Select);

		listener.add(s, SocketDirection::Read);
		listener.add(s2, SocketDirection::Read);

		ASSERT_EQ(2UL, listener.size());
	} catch (const std::exception &ex) {
	}
}

TEST(ListenerMethodSelect, remove)
{
	try {
		Socket s(AF_INET, SOCK_STREAM, 0);
		Socket s2(AF_INET, SOCK_STREAM, 0);
		SocketListener listener(SocketMethod::Select);

		listener.add(s, SocketDirection::Read);
		listener.add(s2, SocketDirection::Read);
		listener.remove(s, SocketDirection::Read);
		listener.remove(s2, SocketDirection::Read);

		ASSERT_EQ(0UL, listener.size());
	} catch (const std::exception &ex) {
	}
}

/*
 * Add two sockets for both reading and writing, them remove only reading and then
 * move only writing.
 */
TEST(ListenerMethodSelect, inOut)
{
	try {
		Socket s(AF_INET, SOCK_STREAM, 0);
		Socket s2(AF_INET, SOCK_STREAM, 0);
		SocketListener listener(SocketMethod::Select);

		listener.add(s, SocketDirection::Read | SocketDirection::Write);
		listener.add(s2, SocketDirection::Read | SocketDirection::Write);

		listener.list([&] (Socket &si, SocketDirection dir) {
			ASSERT_TRUE(si == s || si == s2);
			ASSERT_EQ(0x03, static_cast<int>(dir));
		});

		listener.remove(s, SocketDirection::Write);
		listener.remove(s2, SocketDirection::Write);

		ASSERT_EQ(2UL, listener.size());

		listener.list([&] (Socket &si, SocketDirection dir) {
			ASSERT_TRUE(si == s || si == s2);
			ASSERT_EQ(SocketDirection::Read, dir);
		});

		listener.remove(s, SocketDirection::Read);
		listener.remove(s2, SocketDirection::Read);

		ASSERT_EQ(0UL, listener.size());
	} catch (const std::exception &ex) {
	}
}

TEST(ListenerMethodSelect, addSame)
{
	try {
		Socket s(AF_INET, SOCK_STREAM, 0);
		SocketListener listener(SocketMethod::Select);

		listener.add(s, SocketDirection::Read);
		ASSERT_EQ(1UL, listener.size());
		listener.list([&] (const Socket &si, SocketDirection dir) {
			ASSERT_TRUE(si == s);
			ASSERT_EQ(SocketDirection::Read, dir);
		});

		listener.add(s, SocketDirection::Write);
		ASSERT_EQ(1UL, listener.size());
		listener.list([&] (const Socket &si, SocketDirection dir) {
			ASSERT_TRUE(si == s);
			ASSERT_EQ(0x03, static_cast<int>(dir));
		});

		// Oops, added the same
		listener.add(s, SocketDirection::Read);
		listener.add(s, SocketDirection::Write);
		listener.add(s, SocketDirection::Read | SocketDirection::Write);

		ASSERT_EQ(1UL, listener.size());
		listener.list([&] (const Socket &si, SocketDirection dir) {
			ASSERT_TRUE(si == s);
			ASSERT_EQ(0x03, static_cast<int>(dir));
		});
	} catch (const std::exception &ex) {
	}
}

/* --------------------------------------------------------
 * Poll tests
 * -------------------------------------------------------- */

TEST(ListenerMethodPoll, add)
{
	try {
		Socket s(AF_INET, SOCK_STREAM, 0);
		Socket s2(AF_INET, SOCK_STREAM, 0);
		SocketListener listener(SocketMethod::Poll);

		listener.add(s, SocketDirection::Read);
		listener.add(s2, SocketDirection::Read);

		ASSERT_EQ(2UL, listener.size());
	} catch (const std::exception &ex) {
	}
}

TEST(ListenerMethodPoll, remove)
{
	try {
		Socket s(AF_INET, SOCK_STREAM, 0);
		Socket s2(AF_INET, SOCK_STREAM, 0);
		SocketListener listener(SocketMethod::Poll);

		listener.add(s, SocketDirection::Read);
		listener.add(s2, SocketDirection::Read);
		listener.remove(s, SocketDirection::Read);
		listener.remove(s2, SocketDirection::Read);

		ASSERT_EQ(0UL, listener.size());
	} catch (const std::exception &ex) {
	}
}

TEST(ListenerMethodPoll, inOut)
{
	try {
		Socket s(AF_INET, SOCK_STREAM, 0);
		Socket s2(AF_INET, SOCK_STREAM, 0);
		SocketListener listener(SocketMethod::Poll);

		listener.add(s, SocketDirection::Read | SocketDirection::Write);
		listener.add(s2, SocketDirection::Read | SocketDirection::Write);

		listener.list([&] (Socket &si, SocketDirection dir) {
			ASSERT_TRUE(si == s || si == s2);
			ASSERT_EQ(0x03, static_cast<int>(dir));
		});

		listener.remove(s, SocketDirection::Write);
		listener.remove(s2, SocketDirection::Write);

		ASSERT_EQ(2UL, listener.size());

		listener.list([&] (Socket &si, SocketDirection dir) {
			ASSERT_TRUE(si == s || si == s2);
			ASSERT_EQ(SocketDirection::Read, dir);
		});

		listener.remove(s, SocketDirection::Read);
		listener.remove(s2, SocketDirection::Read);

		ASSERT_EQ(0UL, listener.size());
	} catch (const std::exception &ex) {
	}
}

TEST(ListenerMethodPoll, addSame)
{
	try {
		Socket s(AF_INET, SOCK_STREAM, 0);
		SocketListener listener(SocketMethod::Poll);

		listener.add(s, SocketDirection::Read);
		ASSERT_EQ(1UL, listener.size());
		listener.list([&] (const Socket &si, SocketDirection dir) {
			ASSERT_TRUE(si == s);
			ASSERT_EQ(SocketDirection::Read, dir);
		});

		listener.add(s, SocketDirection::Write);
		ASSERT_EQ(1UL, listener.size());
		listener.list([&] (const Socket &si, SocketDirection dir) {
			ASSERT_TRUE(si == s);
			ASSERT_EQ(0x03, static_cast<int>(dir));
		});

		// Oops, added the same
		listener.add(s, SocketDirection::Read);
		listener.add(s, SocketDirection::Write);
		listener.add(s, SocketDirection::Read | SocketDirection::Write);

		ASSERT_EQ(1UL, listener.size());
		listener.list([&] (const Socket &si, SocketDirection dir) {
			ASSERT_TRUE(si == s);
			ASSERT_EQ(0x03, static_cast<int>(dir));
		});
	} catch (const std::exception &ex) {
	}
}

/* --------------------------------------------------------
 * Socket listener class
 * -------------------------------------------------------- */

TEST(Listener, connection)
{
	std::thread client([] () {
		Socket client;

		std::this_thread::sleep_for(3s);

		try {
			client = Socket(AF_INET, SOCK_STREAM, 0);
			client.connect(Internet("localhost", 10000, AF_INET));
		} catch (const std::exception &ex) {
		}

		client.close();
	});

	Socket s;
	SocketListener listener;

	try {
		s = Socket(AF_INET, SOCK_STREAM, 0);

		s.bind(Internet("localhost", 10000, AF_INET));
		s.listen(8);

		listener.add(s, SocketDirection::Read);

		auto client = listener.select(10s);

		ASSERT_TRUE(client.direction == SocketDirection::Read);
		ASSERT_TRUE(client.socket == s);
	} catch (const std::exception &ex) {
	}

	s.close();
	client.join();
}

TEST(Listener, connectionAndRead)
{
	std::thread thread([] () {
		Socket client;

		std::this_thread::sleep_for(3s);

		try {
			client = Socket(AF_INET, SOCK_STREAM, 0);
			client.connect(Internet("localhost", 10000, AF_INET));
			client.send("hello world");
		} catch (const std::exception &ex) {
		}

		client.close();
	});

	Socket s;
	SocketListener listener;

	try {
		s = Socket(AF_INET, SOCK_STREAM, 0);

		s.bind(Internet("localhost", 10000, AF_INET));
		s.listen(8);

		// Read for master
		listener.add(s, SocketDirection::Read);

		auto result = listener.select(10s);

		ASSERT_TRUE(result.direction == SocketDirection::Read);
		ASSERT_TRUE(result.socket == s);

		// Wait for client
		auto client = s.accept();
		listener.add(client, SocketDirection::Read);

		result = listener.select(10s);

		ASSERT_TRUE(result.direction == SocketDirection::Read);
		ASSERT_TRUE(result.socket == client);

		char data[512];
		auto nb = client.recv(data, sizeof (data) - 1);

		data[nb] = '\0';

		client.close();
		ASSERT_STREQ("hello world", data);
	} catch (const std::exception &ex) {
	}

	s.close();
	thread.join();
}

TEST(Listener, bigData)
{
	auto producer = [] () {
		std::string data;

		data.reserve(9000000);
		for (int i = 0; i < 9000000; ++i)
			data.push_back('a');

		data.push_back('\r');
		data.push_back('\n');

		std::this_thread::sleep_for(3s);

		Socket client;
		SocketListener listener;

		try {
			client = Socket(AF_INET, SOCK_STREAM, 0);

			client.connect(Internet("localhost", 10000, AF_INET));
			client.blockMode(false);
			listener.add(client, SocketDirection::Write);

			while (data.size() > 0) {
				auto s = listener.select(30s).socket;
				auto nb = s.send(data.data(), data.size());
				data.erase(0, nb);
			}
		} catch (const std::exception &ex) {
		}

		client.close();
	};

	auto consumer = [] () {
		std::ostringstream out;

		Socket server;
		SocketListener listener;
		bool finished(false);

		try {
			server = Socket(AF_INET, SOCK_STREAM, 0);

			server.bind(Internet("*", 10000, AF_INET));
			server.listen(10);
			listener.add(server, SocketDirection::Read);

			while (!finished) {
				auto s = listener.select(60s).socket;

				if (s == server) {
					listener.add(s.accept(), SocketDirection::Read);
				} else {
					char data[512];
					auto nb = s.recv(data, sizeof (data) - 1);

					if (nb == 0)
						finished = true;
					else {
						data[nb] = '\0';
						out << data;
					}
				}
			}

			ASSERT_EQ(9000002UL, out.str().size());
		} catch (const std::exception &ex) {
		}

		server.close();
	};

	std::thread tconsumer(consumer);
	std::thread tproducer(producer);

	tconsumer.join();
	tproducer.join();
}

/* --------------------------------------------------------
 * Multiple selection tests
 * -------------------------------------------------------- */

TEST(MultipleSelection, select)
{
	/*
	 * Normally, 3 sockets added for writing should be marked ready immediately
	 * as there are no data being currently queued to be sent.
	 */
	std::thread tester([] () {
		try {
			SocketListener masterListener, clientListener;
			Socket master(AF_INET, SOCK_STREAM, 0);

			master.bind(Internet("*", 10000, AF_INET));
			master.listen(8);

			masterListener.add(master, SocketDirection::Read);

			while (clientListener.size() != 3) {
				masterListener.select(3s);
				clientListener.add(master.accept(), SocketDirection::Write);
			}

			// Now do the test of writing
			auto result = clientListener.selectMultiple(3s);
			ASSERT_EQ(3UL, result.size());

			clientListener.list([] (auto s, auto) {
				s.close();
			});
	
			master.close();
		} catch (const std::exception &ex) {
		}
	});
	
	try {
		Socket s1(AF_INET, SOCK_STREAM, 0);
		Socket s2(AF_INET, SOCK_STREAM, 0);
		Socket s3(AF_INET, SOCK_STREAM, 0);

		s1.connect(Internet("localhost", 10000, AF_INET));
		s2.connect(Internet("localhost", 10000, AF_INET));
		s3.connect(Internet("localhost", 10000, AF_INET));

		s1.close();
		s2.close();
		s3.close();
	} catch (const std::exception &ex) {
	}

	tester.join();
}

#if defined(SOCKET_LISTENER_HAVE_POLL)

TEST(MultipleSelection, poll)
{
	/*
	 * Normally, 3 sockets added for writing should be marked ready immediately
	 * as there are no data being currently queued to be sent.
	 */
	std::thread tester([] () {
		try {
			SocketListener masterListener(SocketMethod::Poll), clientListener(SocketMethod::Poll);
			Socket master(AF_INET, SOCK_STREAM, 0);

			master.bind(Internet("*", 10000, AF_INET));
			master.listen(8);

			masterListener.add(master, SocketDirection::Read);

			while (clientListener.size() != 3) {
				masterListener.select(3s);
				clientListener.add(master.accept(), SocketDirection::Write);
			}

			// Now do the test of writing
			auto result = clientListener.selectMultiple(3s);
			ASSERT_EQ(3UL, result.size());

			clientListener.list([] (auto s, auto) {
				s.close();
			});
	
			master.close();
		} catch (const std::exception &ex) {
		}
	});

	try {
		Socket s1(AF_INET, SOCK_STREAM, 0);
		Socket s2(AF_INET, SOCK_STREAM, 0);
		Socket s3(AF_INET, SOCK_STREAM, 0);

		s1.connect(Internet("localhost", 10000, AF_INET));
		s2.connect(Internet("localhost", 10000, AF_INET));
		s3.connect(Internet("localhost", 10000, AF_INET));

		s1.close();
		s2.close();
		s3.close();
	} catch (const std::exception &ex) {
	}

	tester.join();
}

#endif

/* --------------------------------------------------------
 * Basic TCP tests
 * -------------------------------------------------------- */

TEST(BasicTcp, sendipv4) {
	std::thread client([] () {
		Socket s;

		std::this_thread::sleep_for(3s);
		try {
			s = Socket(AF_INET, SOCK_STREAM, 0);

			s.connect(Internet("localhost", 10000, AF_INET));
			s.send("hello");
		} catch (const std::exception &ex) {
		}

		s.close();
	});

	Socket server;

	try {
		server = Socket(AF_INET, SOCK_STREAM, 0);

		server.bind(Internet("*", 10000, SOCK_STREAM));
		server.listen(8);

		auto client = server.accept();

		char data[512];
		auto nb = client.recv(data, sizeof (data) - 1);

		data[nb] = '\0';

		ASSERT_STREQ("hello", data);
	} catch (const std::exception &ex) {
	}

	server.close();
	client.join();
}

TEST(BasicTcp, sendipv6) {
	std::thread client([] () {
		Socket s;

		std::this_thread::sleep_for(3s);
		try {
			s = Socket(AF_INET6, SOCK_STREAM, 0);

			s.connect(Internet("localhost", 10000, AF_INET6));
			s.send("hello");
		} catch (const std::exception &ex) {
		}

		s.close();
	});

	Socket server;

	try {
		server = Socket(AF_INET6, SOCK_STREAM, 0);

		server.bind(Internet("*", 10000, SOCK_STREAM));
		server.listen(8);

		auto client = server.accept();

		char data[512];
		auto nb = client.recv(data, sizeof (data) - 1);

		data[nb] = '\0';

		ASSERT_STREQ("hello", data);
	} catch (const std::exception &ex) {
	}

	server.close();
	client.join();
}

#if !defined(_WIN32)

TEST(BasicTcp, sendunix) {
	std::thread client([] () {
		Socket s;

		std::this_thread::sleep_for(3s);
		try {
			s = Socket(AF_UNIX, SOCK_STREAM, 0);

			s.connect(Unix("/tmp/gtest-send-tcp-unix.sock"));
			s.send("hello");
		} catch (const std::exception &ex) {
		}

		s.close();
	});

	Socket server;

	try {
		server = Socket(AF_UNIX, SOCK_STREAM, 0);

		server.bind(Unix("/tmp/gtest-send-tcp-unix.sock", true));
		server.listen(8);

		auto client = server.accept();

		char data[512];
		auto nb = client.recv(data, sizeof (data) - 1);

		data[nb] = '\0';

		ASSERT_STREQ("hello", data);
	} catch (const std::exception &ex) {
	}

	server.close();
	client.join();
}

#endif

TEST(BasicUdp, sendipv4) {
	std::thread client([] () {
		Socket s;

		std::this_thread::sleep_for(3s);
		try {
			s = Socket(AF_INET, SOCK_DGRAM, 0);

			s.sendto("hello", Internet("localhost", 10000, AF_INET));
		} catch (const std::exception &ex) {
		}

		s.close();
	});

	Socket server;

	try {
		server = Socket(AF_INET, SOCK_DGRAM, 0);

		server.bind(Internet("*", 10000, SOCK_DGRAM));

		char data[512];
		auto nb = server.recvfrom(data, sizeof (data) - 1);

		data[nb] = '\0';

		ASSERT_STREQ("hello", data);
	} catch (const std::exception &ex) {
	}

	server.close();
	client.join();
}

TEST(BasicUdp, sendipv6) {
	std::thread client([] () {
		Socket s;

		std::this_thread::sleep_for(3s);
		try {
			s = Socket(AF_INET6, SOCK_DGRAM, 0);

			s.sendto("hello", Internet("localhost", 10000, AF_INET6));
		} catch (const std::exception &ex) {
		}

		s.close();
	});

	Socket server;

	try {
		server = Socket(AF_INET6, SOCK_DGRAM, 0);

		server.bind(Internet("*", 10000, SOCK_DGRAM));

		char data[512];
		auto nb = server.recvfrom(data, sizeof (data) - 1);

		data[nb] = '\0';

		ASSERT_STREQ("hello", data);
	} catch (const std::exception &ex) {
	}

	server.close();
	client.join();
}

#if !defined(_WIN32)

TEST(BasicUdp, sendunix) {
	std::thread client([] () {
		Socket s;

		std::this_thread::sleep_for(3s);
		try {
			s = Socket(AF_UNIX, SOCK_DGRAM, 0);

			s.sendto("hello", Unix("/tmp/gtest-send-udp-unix.sock"));
		} catch (const std::exception &ex) {
		}

		s.close();
	});

	Socket server;

	try {
		server = Socket(AF_UNIX, SOCK_DGRAM, 0);

		server.bind(Unix("/tmp/gtest-send-udp-unix.sock", true));

		char data[512];
		auto nb = server.recvfrom(data, sizeof (data) - 1);

		data[nb] = '\0';

		ASSERT_STREQ("hello", data);
	} catch (const std::exception &ex) {
	}

	server.close();
	client.join();
}

#endif

int main(int argc, char **argv)
{
	Socket::init();

	testing::InitGoogleTest(&argc, argv);

	auto ret = RUN_ALL_TESTS();

	Socket::finish();

	return ret;
}