/*
 * main.cpp -- test sockets
 *
 * Copyright (c) 2013, 2014 David Demelier <markand@malikania.fr>
 *
 * Permission to use, copy, modify, and/or distribute this software for any
 * purpose with or without fee is hereby granted, provided that the above
 * copyright notice and this permission notice appear in all copies.
 *
 * THE SOFTWARE IS PROVIDED "AS IS" AND THE AUTHOR DISCLAIMS ALL WARRANTIES
 * WITH REGARD TO THIS SOFTWARE INCLUDING ALL IMPLIED WARRANTIES OF
 * MERCHANTABILITY AND FITNESS. IN NO EVENT SHALL THE AUTHOR BE LIABLE FOR
 * ANY SPECIAL, DIRECT, INDIRECT, OR CONSEQUENTIAL DAMAGES OR ANY DAMAGES
 * WHATSOEVER RESULTING FROM LOSS OF USE, DATA OR PROFITS, WHETHER IN AN
 * ACTION OF CONTRACT, NEGLIGENCE OR OTHER TORTIOUS ACTION, ARISING OUT OF
 * OR IN CONNECTION WITH THE USE OR PERFORMANCE OF THIS SOFTWARE.
 */

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

#include <gtest/gtest.h>

#include "Socket.h"
#include "SocketAddress.h"
#include "SocketListener.h"
#include "SocketSsl.h"
#include "SocketTcp.h"
#include "SocketUdp.h"

using namespace address;
using namespace std::literals::chrono_literals;

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

class TcpServerTest : public testing::Test {
protected:
	SocketTcp m_server{AF_INET, 0};
	SocketTcp m_client{AF_INET, 0};

	std::thread m_tserver;
	std::thread m_tclient;

public:
	TcpServerTest()
	{
		m_server.set(SOL_SOCKET, SO_REUSEADDR, 1);
	}

	~TcpServerTest()
	{
		if (m_tserver.joinable())
			m_tserver.join();
		if (m_tclient.joinable())
			m_tclient.join();
	}
};

TEST_F(TcpServerTest, connect)
{
	m_tserver = std::thread([this] () {
		m_server.bind(Internet("*", 16000, AF_INET));

		ASSERT_EQ(SocketState::Bound, m_server.state());

		m_server.listen();
		m_server.accept();
		m_server.close();
	});

	std::this_thread::sleep_for(100ms);

	m_tclient = std::thread([this] () {
		m_client.connect(Internet("127.0.0.1", 16000, AF_INET));

		ASSERT_EQ(SocketState::Connected, m_client.state());

		m_client.close();
	});
}

TEST_F(TcpServerTest, io)
{
	m_tserver = std::thread([this] () {
		m_server.bind(Internet("*", 16000, AF_INET));
		m_server.listen();

		auto client = m_server.accept();
		auto msg = client.recv(512);

		ASSERT_EQ("hello world", msg);

		client.send(msg);
		client.close();

		m_server.close();
	});

	std::this_thread::sleep_for(100ms);

	m_tclient = std::thread([this] () {
		m_client.connect(Internet("127.0.0.1", 16000, AF_INET));
		m_client.send("hello world");

		ASSERT_EQ("hello world", m_client.recv(512));

		m_client.close();
	});
}

/* --------------------------------------------------------
 * UDP tests
 * -------------------------------------------------------- */

class UdpServerTest : public testing::Test {
protected:
	SocketUdp m_server{AF_INET, 0};
	SocketUdp m_client{AF_INET, 0};

	std::thread m_tserver;
	std::thread m_tclient;

public:
	UdpServerTest()
	{
		m_server.set(SOL_SOCKET, SO_REUSEADDR, 1);
	}

	~UdpServerTest()
	{
		if (m_tserver.joinable())
			m_tserver.join();
		if (m_tclient.joinable())
			m_tclient.join();
	}
};

TEST_F(UdpServerTest, io)
{
	m_tserver = std::thread([this] () {
		SocketAddress info;

		m_server.bind(Internet("*", 16000, AF_INET));

		auto msg = m_server.recvfrom(512, info);

		ASSERT_EQ("hello world", msg);

		m_server.sendto(msg, info);
		m_server.close();
	});

	std::this_thread::sleep_for(100ms);

	m_tclient = std::thread([this] () {
		Internet info("127.0.0.1", 16000, AF_INET);

		m_client.sendto("hello world", info);

		ASSERT_EQ("hello world", m_client.recvfrom(512, info));

		m_client.close();
	});
}

/* --------------------------------------------------------
 * Listener tests (standard)
 * -------------------------------------------------------- */

class ListenerTest : public testing::Test {
protected:
	SocketListenerBase<backend::Select> m_listener;
	SocketTcp socket1{AF_INET, 0};
	SocketUdp socket2{AF_INET, 0};

public:
	~ListenerTest()
	{
		socket1.close();
		socket2.close();
	}
};

TEST_F(ListenerTest, set)
{
	m_listener.set(socket1, SocketListener::Read);

	ASSERT_EQ(1, static_cast<int>(m_listener.size()));
	ASSERT_EQ(SocketListener::Read, m_listener.begin()->second);

	m_listener.set(socket1, SocketListener::Write);

	ASSERT_EQ(1, static_cast<int>(m_listener.size()));
	ASSERT_EQ(0x3, m_listener.begin()->second);

	// Fake a re-insert of the same socket
	m_listener.set(socket1, SocketListener::Write);

	ASSERT_EQ(1, static_cast<int>(m_listener.size()));
	ASSERT_EQ(0x3, m_listener.begin()->second);

	// Add an other socket now
	m_listener.set(socket2, SocketListener::Read | SocketListener::Write);

	ASSERT_EQ(2, static_cast<int>(m_listener.size()));

	for (auto &pair : m_listener) {
		ASSERT_EQ(0x3, pair.second);
		ASSERT_TRUE(pair.first == socket1 || pair.first == socket2);
	}
}

TEST_F(ListenerTest, unset)
{
	m_listener.set(socket1, SocketListener::Read | SocketListener::Write);
	m_listener.set(socket2, SocketListener::Read | SocketListener::Write);

	m_listener.unset(socket1, SocketListener::Read);

	ASSERT_EQ(2, static_cast<int>(m_listener.size()));

	// Use a for loop since it can be ordered differently
	for (auto &pair : m_listener) {
		if (pair.first == socket1) {
			ASSERT_EQ(0x2, pair.second);
		} else if (pair.first == socket2) {
			ASSERT_EQ(0x3, pair.second);
		}
	}

	m_listener.unset(socket1, SocketListener::Write);

	ASSERT_EQ(1, static_cast<int>(m_listener.size()));
	ASSERT_EQ(0x3, m_listener.begin()->second);
}

TEST_F(ListenerTest, remove)
{
	m_listener.set(socket1, SocketListener::Read | SocketListener::Write);
	m_listener.set(socket2, SocketListener::Read | SocketListener::Write);
	m_listener.remove(socket1);

	ASSERT_EQ(1, static_cast<int>(m_listener.size()));
	ASSERT_EQ(0x3, m_listener.begin()->second);
}

TEST_F(ListenerTest, clear)
{
	m_listener.set(socket1, SocketListener::Read | SocketListener::Write);
	m_listener.set(socket2, SocketListener::Read | SocketListener::Write);
	m_listener.clear();

	ASSERT_EQ(0, static_cast<int>(m_listener.size()));
}

/* --------------------------------------------------------
 * Listener: poll
 * -------------------------------------------------------- */

#if defined(SOCKET_HAVE_POLL)

class ListenerPollTest : public testing::Test {
protected:
	SocketListenerBase<backend::Poll> m_listener;
	SocketTcp m_masterTcp{AF_INET, 0};
	SocketTcp m_clientTcp{AF_INET, 0};

	std::thread m_tserver;
	std::thread m_tclient;

public:
	ListenerPollTest()
	{
		m_masterTcp.set(SOL_SOCKET, SO_REUSEADDR, 1);
		m_masterTcp.bind(Internet("*", 16000, AF_INET));
		m_masterTcp.listen();
	}

	~ListenerPollTest()
	{
		if (m_tserver.joinable()) {
			m_tserver.join();
		}
		if (m_tclient.joinable()) {
			m_tclient.join();
		}
	}
};

TEST_F(ListenerPollTest, accept)
{
	m_tserver = std::thread([this] () {
		try {
			m_listener.set(m_masterTcp, SocketListener::Read);
			m_listener.wait();
			m_masterTcp.accept();
			m_masterTcp.close();
		} catch (const std::exception &ex) {
			FAIL() << ex.what();
		}
	});

	std::this_thread::sleep_for(100ms);

	m_tclient = std::thread([this] () {
		m_clientTcp.connect(Internet("127.0.0.1", 16000, AF_INET));
	});
}

TEST_F(ListenerPollTest, recv)
{
	m_tserver = std::thread([this] () {
		try {
			m_listener.set(m_masterTcp, SocketListener::Read);
			m_listener.wait();

			auto sc = m_masterTcp.accept();

			ASSERT_EQ("hello", sc.recv(512));

			m_masterTcp.close();
		} catch (const std::exception &ex) {
			FAIL() << ex.what();
		}
	});

	std::this_thread::sleep_for(100ms);

	m_tclient = std::thread([this] () {
		m_clientTcp.connect(Internet("127.0.0.1", 16000, AF_INET));
		m_clientTcp.send("hello");
	});
}

#endif

/* --------------------------------------------------------
 * Listener: select
 * -------------------------------------------------------- */

class ListenerSelectTest : public testing::Test {
protected:
	SocketListenerBase<backend::Select> m_listener;
	SocketTcp m_masterTcp{AF_INET, 0};
	SocketTcp m_clientTcp{AF_INET, 0};

	std::thread m_tserver;
	std::thread m_tclient;

public:
	ListenerSelectTest()
	{
		m_masterTcp.set(SOL_SOCKET, SO_REUSEADDR, 1);
		m_masterTcp.bind(Internet("*", 16000, AF_INET));
		m_masterTcp.listen();
	}

	~ListenerSelectTest()
	{
		if (m_tserver.joinable()) {
			m_tserver.join();
		}
		if (m_tclient.joinable()) {
			m_tclient.join();
		}
	}
};

TEST_F(ListenerSelectTest, accept)
{
	m_tserver = std::thread([this] () {
		try {
			m_listener.set(m_masterTcp, SocketListener::Read);
			m_listener.wait();
			m_masterTcp.accept();
			m_masterTcp.close();
		} catch (const std::exception &ex) {
			FAIL() << ex.what();
		}
	});

	std::this_thread::sleep_for(100ms);

	m_tclient = std::thread([this] () {
		m_clientTcp.connect(Internet("127.0.0.1", 16000, AF_INET));
	});
}

TEST_F(ListenerSelectTest, recv)
{
	m_tserver = std::thread([this] () {
		try {
			m_listener.set(m_masterTcp, SocketListener::Read);
			m_listener.wait();

			auto sc = m_masterTcp.accept();

			ASSERT_EQ("hello", sc.recv(512));

			m_masterTcp.close();
		} catch (const std::exception &ex) {
			FAIL() << ex.what();
		}
	});

	std::this_thread::sleep_for(100ms);

	m_tclient = std::thread([this] () {
		m_clientTcp.connect(Internet("127.0.0.1", 16000, AF_INET));
		m_clientTcp.send("hello");
	});
}

/* --------------------------------------------------------
 * Listener: kqueue
 * -------------------------------------------------------- */

#if defined(SOCKET_HAVE_KQUEUE)

class ListenerKqueueTest : public testing::Test {
protected:
	SocketListenerBase<backend::Kqueue> m_listener;
	SocketTcp m_masterTcp{AF_INET, 0};
	SocketTcp m_clientTcp{AF_INET, 0};

	std::thread m_tserver;
	std::thread m_tclient;

public:
	ListenerKqueueTest()
	{
		m_masterTcp.set(SOL_SOCKET, SO_REUSEADDR, 1);
		m_masterTcp.bind(Internet("*", 16000, AF_INET));
		m_masterTcp.listen();
	}

	~ListenerKqueueTest()
	{
		if (m_tserver.joinable()) {
			m_tserver.join();
		}
		if (m_tclient.joinable()) {
			m_tclient.join();
		}
	}
};

TEST_F(ListenerKqueueTest, accept)
{
	m_tserver = std::thread([this] () {
		try {
			m_listener.set(m_masterTcp, SocketListener::Read);
			m_listener.wait();
			m_masterTcp.accept();
			m_masterTcp.close();
		} catch (const std::exception &ex) {
			FAIL() << ex.what();
		}
	});

	std::this_thread::sleep_for(100ms);

	m_tclient = std::thread([this] () {
		m_clientTcp.connect(Internet("127.0.0.1", 16000, AF_INET));
	});
}

TEST_F(ListenerKqueueTest, recv)
{
	m_tserver = std::thread([this] () {
		try {
			m_listener.set(m_masterTcp, SocketListener::Read);
			m_listener.wait();

			auto sc = m_masterTcp.accept();

			ASSERT_EQ("hello", sc.recv(512));

			m_masterTcp.close();
		} catch (const std::exception &ex) {
			FAIL() << ex.what();
		}
	});

	std::this_thread::sleep_for(100ms);

	m_tclient = std::thread([this] () {
		m_clientTcp.connect(Internet("127.0.0.1", 16000, AF_INET));
		m_clientTcp.send("hello");
	});
}

#endif // !SOCKET_HAVE_KQUEUE

/* --------------------------------------------------------
 * Non-blocking connect
 * -------------------------------------------------------- */

class NonBlockingConnectTest : public testing::Test {
protected:
	SocketTcp m_server{AF_INET, 0};
	SocketTcp m_client{AF_INET, 0};

	std::thread m_tserver;
	std::thread m_tclient;

public:
	NonBlockingConnectTest()
	{
		m_client.setBlockMode(false);
	}

	~NonBlockingConnectTest()
	{
		if (m_tserver.joinable())
			m_tserver.join();
		if (m_tclient.joinable())
			m_tclient.join();
	}
};

TEST_F(NonBlockingConnectTest, success)
{
	m_server.set(SOL_SOCKET, SO_REUSEADDR, 1);
	m_server.bind(Internet("*", 16000, AF_INET));
	m_server.listen();

	m_tserver = std::thread([this] () {
		SocketTcp client = m_server.accept();

		std::this_thread::sleep_for(100ms);

		m_server.close();
		client.close();
	});

	std::this_thread::sleep_for(100ms);

	m_tclient = std::thread([this] () {
		try {
			m_client.waitConnect(Internet("127.0.0.1", 16000, AF_INET), 3000);
		} catch (const SocketError &error) {
			FAIL() << error.function() << ": " << error.what();
		}

		ASSERT_EQ(SocketState::Connected, m_client.state());

		m_client.close();
	});
}

TEST_F(NonBlockingConnectTest, fail)
{
	/*
	 * /!\ If you find a way to test this locally please tell me /!\
	 */
	m_tclient = std::thread([this] () {
		try {
			m_client.waitConnect(Internet("google.fr", 9000, AF_INET), 100);

			FAIL() << "Expected exception, got success";
		} catch (const SocketError &error) {
			ASSERT_EQ(SocketError::Timeout, error.code());
		}

		m_client.close();
	});
}

/* --------------------------------------------------------
 * TCP accept
 * -------------------------------------------------------- */

class TcpAcceptTest : public testing::Test {
protected:
	SocketTcp m_server{AF_INET, 0};
	SocketTcp m_client{AF_INET, 0};

	std::thread m_tserver;
	std::thread m_tclient;

public:
	TcpAcceptTest()
	{
		m_server.set(SOL_SOCKET, SO_REUSEADDR, 1);
		m_server.bind(Internet("*", 16000, AF_INET));
		m_server.listen();
	}

	~TcpAcceptTest()
	{
		if (m_tserver.joinable())
			m_tserver.join();
		if (m_tclient.joinable())
			m_tclient.join();
	}
};

TEST_F(TcpAcceptTest, blockingWaitSuccess)
{
	m_tserver = std::thread([this] () {
		try {
			m_server.waitAccept(3000).close();
		} catch (const SocketError &error) {
			FAIL() << error.what();
		}

		m_server.close();
	});

	std::this_thread::sleep_for(100ms);

	m_tclient = std::thread([this] () {
		m_client.connect(Internet("127.0.0.1", 16000, AF_INET));
		m_client.close();
	});
}

TEST_F(TcpAcceptTest, nonBlockingWaitSuccess)
{
	m_tserver = std::thread([this] () {
		try {
			m_server.setBlockMode(false);
			m_server.waitAccept(3000).close();
		} catch (const SocketError &error) {
			FAIL() << error.what();
		}

		m_server.close();
	});

	std::this_thread::sleep_for(100ms);

	m_tclient = std::thread([this] () {
		m_client.connect(Internet("127.0.0.1", 16000, AF_INET));
		m_client.close();
	});
}

TEST_F(TcpAcceptTest, nonBlockingWaitFail)
{
	// No client, no accept
	try {
		m_server.setBlockMode(false);
		m_server.waitAccept(100).close();

		FAIL() << "Expected exception, got success";
	} catch (const SocketError &error) {
		ASSERT_EQ(SocketError::Timeout, error.code());
	}

	m_server.close();
}

/* --------------------------------------------------------
 * TCP recv
 * -------------------------------------------------------- */

class TcpRecvTest : public testing::Test {
protected:
	SocketTcp m_server{AF_INET, 0};
	SocketTcp m_client{AF_INET, 0};

	std::thread m_tserver;
	std::thread m_tclient;

public:
	TcpRecvTest()
	{
		m_server.set(SOL_SOCKET, SO_REUSEADDR, 1);
		m_server.bind(Internet("*", 16000, AF_INET));
		m_server.listen();
	}

	~TcpRecvTest()
	{
		if (m_tserver.joinable())
			m_tserver.join();
		if (m_tclient.joinable())
			m_tclient.join();
	}
};

TEST_F(TcpRecvTest, blockingSuccess)
{
	m_tserver = std::thread([this] () {
		SocketTcp client = m_server.accept();

		ASSERT_EQ("hello", client.recv(32));

		client.close();
		m_server.close();
	});

	std::this_thread::sleep_for(100ms);

	m_tclient = std::thread([this] () {
		m_client.connect(Internet("127.0.0.1", 16000, AF_INET));
		m_client.send("hello");
		m_client.close();
	});
}

TEST_F(TcpRecvTest, blockingWaitSuccess)
{
	m_tserver = std::thread([this] () {
		SocketTcp client = m_server.accept();

		ASSERT_EQ("hello", client.waitRecv(32, 3000));

		client.close();
		m_server.close();
	});

	std::this_thread::sleep_for(100ms);

	m_tclient = std::thread([this] () {
		m_client.connect(Internet("127.0.0.1", 16000, AF_INET));
		m_client.send("hello");
		m_client.close();
	});
}

TEST_F(TcpRecvTest, nonBlockingWaitSuccess)
{
	m_tserver = std::thread([this] () {
		SocketTcp client = m_server.accept();

		client.setBlockMode(false);

		ASSERT_EQ("hello", client.waitRecv(32, 3000));

		client.close();
		m_server.close();
	});

	std::this_thread::sleep_for(100ms);

	m_tclient = std::thread([this] () {
		m_client.connect(Internet("127.0.0.1", 16000, AF_INET));
		m_client.send("hello");
		m_client.close();
	});
}

/* --------------------------------------------------------
 * Socket SSL
 * -------------------------------------------------------- */

class SslTest : public testing::Test {
protected:
	SocketSsl client{AF_INET, 0};
};

TEST_F(SslTest, connect)
{
	try {
		client.connect(Internet("google.fr", 443, AF_INET));
		client.close();
	} catch (const SocketError &error) {
		FAIL() << error.what();
	}
}

TEST_F(SslTest, recv)
{
	try {
		client.connect(Internet("google.fr", 443, AF_INET));
		client.send("GET / HTTP/1.0\r\n\r\n");

		std::string msg = client.recv(512);
		std::string content = msg.substr(0, 18);

		ASSERT_EQ("HTTP/1.0 302 Found", content);

		client.close();
	} catch (const SocketError &error) {
		FAIL() << error.what();
	}
}

int main(int argc, char **argv)
{
	testing::InitGoogleTest(&argc, argv);

	return RUN_ALL_TESTS();
}