#include "pitime.h"
#include "test_client_helper.h"

#include "gtest/gtest.h"

template<bool WithConnectPing = false, bool WithPong = false, ullong WriteSize = 64_KiB>
PIClientServer::Server * createServer() {
	auto s = new PIClientServer::Server();
	s->setClientFactory([] { return new TestServerClient<WithConnectPing, WithPong, WriteSize>(); });
	s->listenAll(12345);
	return s;
}

bool waitLoop(std::function<bool()> exit_loop, const PISystemTime & timeout) {
	PITimeMeasurer tm;
	while (tm.elapsed() < timeout) {
		if (exit_loop()) {
			return true;
		}
		piMinSleep();
	}
	return false;
}

template<typename Client>
Client * createAndConnectClient() {
	auto c = new Client();
	c->connect(PINetworkAddress::resolve("127.0.0.1", 12345));
	return c;
}

TEST(ClientServer, OneClient) {
	auto const loop_timeout = 1000_ms;
	auto s                  = createServer<false, true>();
	piMinSleep();
	auto c = createAndConnectClient<TestClient<false, false>>();

	waitLoop([s]() { return s->clientsCount() > 0; }, loop_timeout);
	EXPECT_EQ(1, s->clientsCount());

	c->ping();
	waitLoop([c]() { return c->pongCnt() > 0; }, loop_timeout);
	EXPECT_EQ(1, c->pongCnt());

	s->forEachClient([](PIClientServer::ServerClient * sc) { static_cast<TestServerClient<> *>(sc)->ping(); });
	waitLoop([c]() { return c->pongCnt() > 1; }, loop_timeout);
	EXPECT_EQ(2, c->pongCnt());

	EXPECT_TRUE(c->getTCP()->isConnected());

	delete c;
	waitLoop([s]() { return s->clientsCount() == 0; }, loop_timeout);
	EXPECT_EQ(0, s->clientsCount());
	delete s;
}


class ClientSendThread {
	using ClientType = TestClient<false, false, 10_KiB>;

public:
	ClientSendThread() {
		client = createAndConnectClient<ClientType>();
		sendThread.setName("clSend");
	}

	~ClientSendThread() {
		sendThread.stopAndWait();
		delete client;
	}

	void startSend() {
		sendThread.start([this] { client->ping(); }, 100._Hz);
	}

	void sendOnce() { client->ping(); }

	ClientType * client = nullptr;
	PIThread sendThread;
};


int getServerPongs(PIClientServer::Server * s) {
	int pongs = 0;
	s->forEachClient([&pongs](PIClientServer::ServerClient * sc) {
		const auto c = static_cast<TestServerClient<> *>(sc);
		pongs += c->pongCnt();
	});
	return pongs;
}

int getClientsPongs(const PIVector<ClientSendThread *> & clients) {
	int pongs = 0;
	clients.forEach([&pongs](ClientSendThread * c) { pongs += c->client->pongCnt(); });
	return pongs;
}

int getClientsPings(const PIVector<ClientSendThread *> & clients) {
	int pings = 0;
	clients.forEach([&pings](ClientSendThread * c) { pings += c->client->pingCnt(); });
	return pings;
}


TEST(ClientServer, ManyClients) {
	auto const loop_timeout     = 1_s;
	constexpr int clients_count = 20;
	PIVector<ClientSendThread *> clients;
	auto s = createServer<false, false, 100_KiB>();
	piMinSleep();

	piForTimes(clients_count) {
		clients.append(new ClientSendThread());
	}
	EXPECT_EQ(clients_count, clients.size_s());

	waitLoop([s]() { return s->clientsCount() == clients_count; }, loop_timeout);
	EXPECT_EQ(clients_count, s->clientsCount());

	EXPECT_EQ(0, getServerPongs(s));

	EXPECT_EQ(getClientsPings(clients), 0);

	for (const auto c: clients) {
		c->sendOnce();
	}

	EXPECT_EQ(getClientsPings(clients), clients_count);
	EXPECT_TRUE(clients.every([](ClientSendThread * c) { return c->client->pingCnt() == 1; }));
	EXPECT_TRUE(clients.every([](ClientSendThread * c) { return c->client->pongCnt() == 0; }));
	waitLoop([s]() { return getServerPongs(s) >= clients_count; }, loop_timeout);
	EXPECT_EQ(clients_count, getServerPongs(s));

	s->forEachClient([](PIClientServer::ServerClient * sc) { static_cast<TestServerClient<> *>(sc)->ping(); });
	const auto clientCheckPong = [&clients]() { return clients.every([](ClientSendThread * c) { return c->client->pongCnt() == 1; }); };
	waitLoop([&clientCheckPong]() { return clientCheckPong(); }, loop_timeout);
	EXPECT_TRUE(clientCheckPong());

	for (const auto c: clients) {
		c->startSend();
	}
	waitLoop([&clients]() { return getClientsPings(clients) > clients_count * 2; }, loop_timeout);
	EXPECT_TRUE(getClientsPings(clients) > clients_count * 2);
	waitLoop([s]() { return getServerPongs(s) > clients_count * 2; }, loop_timeout);
	EXPECT_TRUE(getServerPongs(s) > clients_count * 2);
	piDeleteAllAndClear(clients);
	waitLoop([s]() { return s->clientsCount() == 0; }, loop_timeout);
	EXPECT_EQ(0, s->clientsCount());
	delete s;
}

TEST(ClientServer, DynamicClients) {
	auto const loop_timeout     = 3_s;
	constexpr int clients_count = 20;
	PIVector<ClientSendThread *> clients;
	PIMutex clients_mutex;
	auto s                 = createServer<true, true, 10_KiB>();

	const auto spawnClient = [&clients, &clients_mutex]() {
		// if (clients.size() > 100) return;
		auto c = new ClientSendThread();
		c->startSend();
		clients_mutex.lock();
		clients << c;
		clients_mutex.unlock();
	};

	piForTimes(clients_count) {
		spawnClient();
	}

	PIThread spawnThread;
	PIThread deleteThread;
	spawnThread.setName("spawn");
	deleteThread.setName("delete");

	spawnThread.start(
		[&spawnClient]() {
			const int new_cnt = 7;
			piForTimes(new_cnt) {
				spawnClient();
			}
		},
		120_Hz);

	deleteThread.start(
		[&clients, &clients_mutex]() {
			const int rm_cnt = 8;
			piForTimes(rm_cnt) {
				ClientSendThread * c = nullptr;
				clients_mutex.lock();
				if (clients.size() > 10) {
					c = clients.take_front();
				}
				clients_mutex.unlock();
				if (c) {
					delete c;
				}
			}
		},
		130_Hz);

	waitLoop([s]() { return s->clientsCount() >= 10; }, loop_timeout);
	EXPECT_GE(s->clientsCount(), 10);

	deleteThread.stopAndWait();
	spawnThread.stopAndWait();

	piDeleteAllAndClear(clients);
	waitLoop([s]() { return s->clientsCount() == 0; }, loop_timeout);
	EXPECT_EQ(0, s->clientsCount());

	delete s;
}
