// 批次 2V1-a: the 40250 EterLib network layer (CNetworkStream over the winsock shim) against a
// loopback peer: non-blocking connect, buffered send/recv, the connect-failure timeout, and the
// Crypto++ key agreement (_IMPROVED_PACKET_ENCRYPTION_) with a second Cipher playing the server.
#include "EterLib/StdAfx.h"
#include "EterLib/NetDevice.h"
#include "EterLib/NetStream.h"

#include <chrono>
#include <cstdio>
#include <cstring>
#include <thread>
#include <vector>

static int g_failures = 0;
#define CHECK(cond)                                                    \
	do {                                                               \
		if (!(cond)) {                                                 \
			std::fprintf(stderr, "%s:%d: CHECK(%s)\n", __FILE__, __LINE__, #cond); \
			++g_failures;                                              \
		}                                                              \
	} while (0)

class TestStream : public CNetworkStream
{
public:
	int successes = 0, failures = 0, remote_disconnects = 0;

	using CNetworkStream::Prepare;
	using CNetworkStream::Activate;
	using CNetworkStream::ActivateCipher;

protected:
	void OnConnectSuccess() override { ++successes; }
	void OnConnectFailure() override { ++failures; }
	void OnRemoteDisconnect() override { ++remote_disconnects; }
};

static int listen_loopback(int* port)
{
	int fd = socket(AF_INET, SOCK_STREAM, 0);
	sockaddr_in addr = {};
	addr.sin_family = AF_INET;
	addr.sin_addr.s_addr = htonl(INADDR_LOOPBACK);
	bind(fd, reinterpret_cast<sockaddr*>(&addr), sizeof(addr));
	socklen_t len = sizeof(addr);
	getsockname(fd, reinterpret_cast<sockaddr*>(&addr), &len);
	listen(fd, 1);
	*port = ntohs(addr.sin_port);
	return fd;
}

template <class F>
static bool pump(TestStream& stream, F done, int ms = 2000)
{
	for (int i = 0; i < ms; ++i)
	{
		stream.Process();
		if (done())
			return true;
		std::this_thread::sleep_for(std::chrono::milliseconds(1));
	}
	return false;
}

static bool read_exact(int fd, void* buf, size_t len)
{
	char* p = static_cast<char*>(buf);
	while (len)
	{
		ssize_t n = recv(fd, p, len, 0);
		if (n <= 0)
			return false;
		p += n;
		len -= static_cast<size_t>(n);
	}
	return true;
}

static void connect_and_exchange()
{
	int port = 0;
	int listener = listen_loopback(&port);

	TestStream stream;
	stream.SetRecvBufferSize(4096);
	stream.SetSendBufferSize(4096);
	CHECK(stream.Connect("127.0.0.1", port));
	CHECK(pump(stream, [&] { return stream.IsOnline(); }));
	CHECK(stream.successes == 1 && stream.failures == 0);
	int peer = accept(listener, nullptr, nullptr);
	CHECK(peer >= 0);

	// Plain bytes both ways; Send only queues, Process flushes.
	CHECK(stream.Send(5, "hello"));
	pump(stream, [] { return true; }, 1);
	char got[16] = {};
	CHECK(read_exact(peer, got, 5) && std::memcmp(got, "hello", 5) == 0);
	send(peer, "world!", 6, 0);
	CHECK(pump(stream, [&] { return stream.GetRecvBufferSize() >= 6; }));
	char peek[6] = {};
	CHECK(stream.Peek(6, peek) && std::memcmp(peek, "world!", 6) == 0);
	CHECK(stream.Recv(6, got) && stream.GetRecvBufferSize() == 0);
	CHECK(!stream.Recv(1, got));

	// Key agreement: each side prepares, swaps its public value and activates with the opposite
	// polarity (the 40250 game server calls Activate(false, ...)).
	Cipher server;
	std::vector<unsigned char> client_pub(1024), server_pub(1024);
	size_t client_len = client_pub.size(), server_len = server_pub.size();
	size_t client_agreed = stream.Prepare(client_pub.data(), &client_len);
	size_t server_agreed = server.Prepare(server_pub.data(), &server_len);
	CHECK(client_agreed > 0 && client_agreed == server_agreed);
	CHECK(stream.Activate(client_agreed, server_pub.data(), server_len));
	CHECK(server.Activate(false, server_agreed, client_pub.data(), client_len));
	stream.ActivateCipher();
	server.set_activated(true);
	CHECK(stream.IsSecurityMode());

	// Client → server: ciphertext on the wire, the server's decoder restores it.
	const char secret[] = "secret-login-packet";
	CHECK(stream.Send(sizeof(secret), secret));
	pump(stream, [] { return true; }, 1);
	char wire[sizeof(secret)] = {};
	CHECK(read_exact(peer, wire, sizeof(wire)));
	CHECK(std::memcmp(wire, secret, sizeof(secret)) != 0);
	server.Decrypt(wire, sizeof(wire));
	CHECK(std::memcmp(wire, secret, sizeof(secret)) == 0);

	// Server → client.
	char reply[] = "phase-select";
	server.Encrypt(reply, sizeof(reply));
	send(peer, reply, sizeof(reply), 0);
	CHECK(pump(stream, [&] { return stream.GetRecvBufferSize() >= static_cast<int>(sizeof(reply)); }));
	CHECK(stream.Recv(sizeof(reply), got) && std::strcmp(got, "phase-select") == 0);

	// The peer hangs up: recv answers 0, Process reports it once and clears the stream.
	close(peer);
	CHECK(pump(stream, [&] { return stream.remote_disconnects > 0; }));
	CHECK(stream.remote_disconnects == 1 && !stream.IsOnline() && !stream.IsSecurityMode());
	close(listener);
}

static void connect_refused()
{
	// A port nobody listens on: Winsock reports the failed connect only in exceptfds, which
	// CNetworkStream::Process does not pass, so the stream waits out the limit and then fails.
	int port = 0;
	int listener = listen_loopback(&port);
	close(listener);

	TestStream stream;
	CNetworkAddress addr;
	addr.Set("127.0.0.1", port);
	CHECK(stream.Connect(addr, 1));
	CHECK(pump(stream, [&] { return stream.failures > 0; }, 3000));
	CHECK(stream.successes == 0 && stream.failures == 1 && !stream.IsOnline());
}

int main()
{
	CNetworkDevice device;
	CHECK(device.Create());
	connect_and_exchange();
	connect_refused();

	if (g_failures)
	{
		std::fprintf(stderr, "%d check(s) failed\n", g_failures);
		return 1;
	}
	std::printf("port_net_test: ok\n");
	return 0;
}