// See port_login_flow_server.h. Each step names the 40250 server function whose behaviour it plays.
#include "port_login_flow_server.h"

#include "../src/net/classic/classic_cipher.h"
#include "../src/net/classic/sequence_table.h"
#include "../src/net/classic/wire_classic.h"

#include <arpa/inet.h>
#include <netinet/in.h>
#include <poll.h>
#include <sys/socket.h>
#include <unistd.h>

#include <chrono>
#include <cstdio>
#include <cstring>

using namespace mtnet::classic;

namespace
{
constexpr std::uint32_t kHandshake = 0x2468ace0;
constexpr int kHandshakeRetryLimit = 32; // game/src/desc.h HANDSHAKE_RETRY_LIMIT
constexpr int kIdleMs = 20000;

// get_dword_time(): milliseconds on the server's clock.
std::uint32_t now_ms()
{
	using namespace std::chrono;
	static const auto start = steady_clock::now();
	return static_cast<std::uint32_t>(duration_cast<milliseconds>(steady_clock::now() - start).count());
}

int listen_loopback(int* port)
{
	int fd = socket(AF_INET, SOCK_STREAM, 0);
	int on = 1;
	setsockopt(fd, SOL_SOCKET, SO_REUSEADDR, &on, sizeof(on));
	sockaddr_in addr = {};
	addr.sin_family = AF_INET;
	addr.sin_addr.s_addr = htonl(INADDR_LOOPBACK);
	if (bind(fd, reinterpret_cast<sockaddr*>(&addr), sizeof(addr)) != 0 || listen(fd, 4) != 0)
	{
		close(fd);
		return -1;
	}
	socklen_t len = sizeof(addr);
	getsockname(fd, reinterpret_cast<sockaddr*>(&addr), &len);
	*port = ntohs(addr.sin_port);
	return fd;
}

template <class T>
T zeroed()
{
	T value;
	std::memset(&value, 0, sizeof(value));
	return value;
}
} // namespace

// One DESC: the socket, its cipher and the sequence position of the next client packet.
struct FakeLoginServer::Connection
{
	int fd = -1;
	ClassicCipher cipher;
	size_t sequence = 0;
	const std::atomic<bool>* stop = nullptr;

	~Connection()
	{
		if (fd >= 0)
			close(fd);
	}

	// DESC::Packet: encrypted once the key agreement finished.
	bool Send(const void* data, size_t size)
	{
		std::vector<std::uint8_t> bytes(static_cast<const std::uint8_t*>(data), static_cast<const std::uint8_t*>(data) + size);
		if (cipher.activated())
			cipher.encrypt(bytes.data(), bytes.size());
		size_t sent = 0;
		while (sent < bytes.size())
		{
			ssize_t n = send(fd, bytes.data() + sent, bytes.size() - sent, 0);
			if (n <= 0)
				return false;
			sent += static_cast<size_t>(n);
		}
		return true;
	}

	template <class T>
	bool Send(const T& packet)
	{
		return Send(&packet, sizeof(packet));
	}

	// DESC::SetPhase sends the phase packet.
	bool Phase(std::uint8_t phase)
	{
		Phase_ p = {HDR_GC_PHASE, phase};
		return Send(p);
	}

	// Reads exactly `size` bytes, or returns 0 on EOF / -1 on timeout, stop or error.
	int Read(void* data, size_t size)
	{
		auto* p = static_cast<std::uint8_t*>(data);
		size_t got = 0;
		for (int waited = 0; got < size;)
		{
			if (stop->load())
				return -1;
			pollfd pfd = {fd, POLLIN, 0};
			int ready = poll(&pfd, 1, 50);
			if (ready == 0)
			{
				if ((waited += 50) >= kIdleMs)
					return -1;
				continue;
			}
			ssize_t n = recv(fd, p + got, size - got, 0);
			if (n == 0)
				return got == 0 ? 0 : -1;
			if (n < 0)
				return -1;
			got += static_cast<size_t>(n);
		}
		if (cipher.activated())
			cipher.decrypt(data, size);
		return 1;
	}
};

FakeLoginServer::~FakeLoginServer()
{
	Stop();
}

bool FakeLoginServer::Start()
{
	m_auth_listener = listen_loopback(&m_auth_port);
	m_game_listener = listen_loopback(&m_game_port);
	if (m_auth_listener < 0 || m_game_listener < 0)
		return false;
	m_thread = std::thread([this] { Run(); });
	return true;
}

void FakeLoginServer::Stop()
{
	m_stop = true;
	if (m_thread.joinable())
		m_thread.join();
	for (int* fd : {&m_auth_listener, &m_game_listener})
		if (*fd >= 0)
		{
			close(*fd);
			*fd = -1;
		}
}

std::vector<std::string> FakeLoginServer::Events()
{
	std::lock_guard<std::mutex> lock(m_mutex);
	return m_events;
}

std::string FakeLoginServer::Error()
{
	std::lock_guard<std::mutex> lock(m_mutex);
	return m_error;
}

bool FakeLoginServer::Has(const std::string& event)
{
	std::lock_guard<std::mutex> lock(m_mutex);
	for (const auto& e : m_events)
		if (e == event)
			return true;
	return false;
}

void FakeLoginServer::Event(const std::string& event)
{
	std::lock_guard<std::mutex> lock(m_mutex);
	m_events.push_back(event);
}

bool FakeLoginServer::Fail(const std::string& error)
{
	std::lock_guard<std::mutex> lock(m_mutex);
	if (m_error.empty())
		m_error = error;
	return false;
}

bool FakeLoginServer::Accept(int listener, Connection& c)
{
	c.stop = &m_stop;
	for (int waited = 0; !m_stop; waited += 50)
	{
		pollfd pfd = {listener, POLLIN, 0};
		if (poll(&pfd, 1, 50) > 0)
		{
			c.fd = accept(listener, nullptr, nullptr);
			return c.fd >= 0;
		}
		if (waited >= kIdleMs)
			return Fail("no client connected");
	}
	return false;
}

void FakeLoginServer::Run()
{
	{
		Connection auth;
		if (!Accept(m_auth_listener, auth) || !ServeAuth(auth))
			return;
	}
	// The first game connection ends at the character list; selecting a character reconnects
	// (net.DirectEnter → CPythonNetworkStream::ConnectGameServer) and that one goes on to loading.
	for (int index = 1; !m_stop; ++index)
	{
		Connection game;
		if (!Accept(m_game_listener, game) || !ServeGame(game, index))
			return;
	}
}

// DESC::Setup → SetPhase(PHASE_HANDSHAKE) + StartHandshake; CInputProcessor::Handshake /
// DESC::HandshakeProcess; SendKeyAgreement; CInputHandshake::Analyze(HEADER_CG_KEY_AGREEMENT).
bool FakeLoginServer::Handshake(Connection& c, std::uint8_t phase_after)
{
	if (!c.Phase(PHASE_HANDSHAKE))
		return Fail("send PHASE_HANDSHAKE");
	mtnet::classic::Handshake out = {HDR_HANDSHAKE, kHandshake, now_ms(), 0};
	std::uint32_t sent_time = out.time;
	if (!c.Send(out))
		return Fail("send GC_HANDSHAKE");

	for (int retry = 0;; ++retry)
	{
		auto in = zeroed<mtnet::classic::Handshake>();
		if (c.Read(&in, sizeof(in)) != 1 || in.header != HDR_HANDSHAKE)
			return Fail("expected CG_HANDSHAKE");
		if (in.handshake != kHandshake)
			return Fail("CG_HANDSHAKE carries the wrong handshake id");
		const std::uint32_t now = now_ms();
		const int bias = static_cast<int>(now - (in.time + in.delta));
		if (in.delta >= 0 && bias >= 0 && bias <= 50)
			break;
		if (retry + 1 > kHandshakeRetryLimit)
			return Fail("handshake retry limit reached");
		std::int32_t delta = static_cast<std::int32_t>(now - in.time) / 2;
		if (delta < 0)
			delta = static_cast<std::int32_t>(now - sent_time) / 2;
		out = {HDR_HANDSHAKE, kHandshake, now, delta};
		sent_time = now;
		if (!c.Send(out))
			return Fail("send GC_HANDSHAKE retry");
	}

	auto agreement = zeroed<KeyAgreement>();
	size_t data_length = KEY_AGREEMENT_MAX_DATA_LEN;
	const size_t agreed_length = c.cipher.prepare(agreement.data, &data_length);
	if (agreed_length == 0)
		return Fail("cipher prepare");
	agreement.header = HDR_KEY_AGREEMENT;
	agreement.agreed_length = static_cast<std::uint16_t>(agreed_length);
	agreement.data_length = static_cast<std::uint16_t>(data_length);
	if (!c.Send(agreement))
		return Fail("send GC_KEY_AGREEMENT");

	auto reply = zeroed<KeyAgreement>();
	if (c.Read(&reply, sizeof(reply)) != 1 || reply.header != HDR_KEY_AGREEMENT)
		return Fail("expected CG_KEY_AGREEMENT");
	// "Send out the key agreement completion packet first to help client to enter encryption mode".
	KeyAgreementCompleted completed = {HDR_GC_KEY_AGREEMENT_COMPLETED, {}};
	if (!c.Send(completed))
		return Fail("send GC_KEY_AGREEMENT_COMPLETED");
	if (!c.cipher.activate(false, reply.agreed_length, reply.data, reply.data_length))
		return Fail("key agreement failed");
	c.cipher.set_activated(true);
	return c.Phase(phase_after) || Fail("send phase after handshake");
}

// Reads one client packet of a known header: the static size from packet_info.cpp plus the sequence
// byte CInputProcessor::Process checks against gc_abSequence.
static int read_packet(FakeLoginServer::Connection& c, std::uint8_t expect, void* packet, size_t size,
                       std::string* error)
{
	auto* bytes = static_cast<std::uint8_t*>(packet);
	int got = c.Read(bytes, 1);
	if (got != 1)
		return got;
	if (bytes[0] != expect)
	{
		char text[64];
		std::snprintf(text, sizeof(text), "expected CG header %u, got %u", expect, bytes[0]);
		*error = text;
		return -1;
	}
	if (static_cast<size_t>(packet_size_cg(expect)) != size)
	{
		*error = "size table disagrees with the packet struct";
		return -1;
	}
	if (size > 1 && c.Read(bytes + 1, size - 1) != 1)
		return -1;
	if (is_sequence_cg(expect))
	{
		std::uint8_t sequence = 0;
		if (c.Read(&sequence, 1) != 1)
			return -1;
		const std::uint8_t want = SEQUENCE_TABLE[c.sequence];
		if (sequence != want)
		{
			char text[80];
			std::snprintf(text, sizeof(text), "SEQUENCE mismatch 0x%x != 0x%x header %u", want, sequence, expect);
			*error = text;
			return -1;
		}
		c.sequence = (c.sequence + 1) % SEQUENCE_TABLE_SIZE;
	}
	return 1;
}

// g_bAuthServer: CInputAuth::Login → GC_AUTH_SUCCESS with the login key the game server will accept.
bool FakeLoginServer::ServeAuth(Connection& c)
{
	if (!Handshake(c, PHASE_AUTH))
		return false;
	Event("auth:handshake");

	auto login = zeroed<CGLogin3>();
	std::string error;
	if (read_packet(c, HDR_CG_LOGIN3, &login, sizeof(login), &error) != 1)
		return Fail("auth: CG_LOGIN3: " + error);
	if (std::strcmp(login.login, kLogin) != 0 || std::strcmp(login.passwd, kPassword) != 0)
		return Fail("auth: CG_LOGIN3 carries the wrong account");
	Event("auth:login3");

	GCAuthSuccess success = {HDR_GC_AUTH_SUCCESS, kLoginKey, 1};
	if (!c.Send(success))
		return Fail("auth: send GC_AUTH_SUCCESS");

	// CAccountConnector connects to the game server and then disconnects from us.
	std::uint8_t byte = 0;
	if (c.Read(&byte, 1) != 0)
		return Fail("auth: expected the client to hang up");
	Event("auth:closed");
	return true;
}

bool FakeLoginServer::ServeGame(Connection& c, int index)
{
	const std::string name = "game" + std::to_string(index);
	if (!Handshake(c, PHASE_LOGIN))
		return false;
	Event(name + ":handshake");

	// CInputLogin::LoginByKey → (DB) CInputDB::LoginSuccess: GC_EMPIRE, SetPhase(PHASE_SELECT),
	// SendLoginSuccessPacket.
	auto login = zeroed<CGLogin2>();
	std::string error;
	if (read_packet(c, HDR_CG_LOGIN2, &login, sizeof(login), &error) != 1)
		return Fail(name + ": CG_LOGIN2: " + error);
	if (std::strcmp(login.login, kLogin) != 0 || login.login_key != kLoginKey)
		return Fail(name + ": CG_LOGIN2 carries the wrong account or login key");
	Event(name + ":login2");

	GCEmpire empire = {HDR_GC_EMPIRE, 1};
	auto success = zeroed<GCLoginSuccess>();
	success.header = HDR_GC_LOGIN_SUCCESS_NEWSLOT;
	SimplePlayer& player = success.players[0];
	player.id = 1;
	std::strncpy(player.name, kCharacterName, sizeof(player.name) - 1);
	player.job = 0;
	player.level = 10;
	player.st = player.ht = player.dx = player.iq = 5;
	player.x = kX;
	player.y = kY;
	player.addr = static_cast<std::int32_t>(htonl(INADDR_LOOPBACK));
	player.port = static_cast<std::uint16_t>(m_game_port);
	success.handle = static_cast<std::uint32_t>(index);
	success.random_key = 0x5a5a5a5a;
	if (!c.Send(empire) || !c.Phase(PHASE_SELECT) || !c.Send(success))
		return Fail(name + ": send login success");

	// The character list screen: the client either reconnects (DirectEnter) or selects here.
	auto select = zeroed<CGPlayerSelect>();
	const int got = read_packet(c, HDR_CG_CHARACTER_SELECT, &select, sizeof(select), &error);
	if (got == 0)
	{
		Event(name + ":closed");
		return true;
	}
	if (got != 1)
		return Fail(name + ": CG_CHARACTER_SELECT: " + error);
	if (select.player_index != 0)
		return Fail(name + ": CG_CHARACTER_SELECT picks an empty slot");
	Event(name + ":select");

	// CInputDB::PlayerLoad: SetPhase(PHASE_LOADING), CHARACTER::MainCharacterPacket, PointsPacket.
	auto main = zeroed<GCMainCharacter>();
	main.header = HDR_GC_MAIN_CHARACTER;
	main.vid = kMainVID;
	main.race = 0;
	std::strncpy(main.name, kCharacterName, sizeof(main.name) - 1);
	main.x = kX;
	main.y = kY;
	main.empire = 1;
	auto points = zeroed<GCPoints>();
	points.header = HDR_GC_CHARACTER_POINTS;
	points.points[1] = 10; // POINT_LEVEL
	if (!c.Phase(PHASE_LOADING) || !c.Send(main) || !c.Send(points))
		return Fail(name + ": send loading packets");

	// The client answers the main character with its version (CInputProcessor::Version).
	std::uint8_t header = 0;
	if (c.Read(&header, 1) != 1)
		return Fail(name + ": expected CG_CLIENT_VERSION");
	const size_t size = header == HDR_CG_CLIENT_VERSION2 ? sizeof(CGClientVersion2)
	                  : header == HDR_CG_CLIENT_VERSION  ? sizeof(CGClientVersion)
	                                                     : 0;
	if (!size)
		return Fail(name + ": expected CG_CLIENT_VERSION, got header " + std::to_string(header));
	std::vector<std::uint8_t> rest(size - 1 + 1);
	if (c.Read(rest.data(), rest.size()) != 1 || rest.back() != SEQUENCE_TABLE[c.sequence])
		return Fail(name + ": CG_CLIENT_VERSION body or sequence");
	c.sequence = (c.sequence + 1) % SEQUENCE_TABLE_SIZE;
	Event(name + ":client_version");

	// game.GameWindow.Open (after the loading steps' net.StartGame()) sends CG_ENTERGAME.
	auto enter = zeroed<CGEnterGame>();
	if (read_packet(c, HDR_CG_ENTERGAME, &enter, sizeof(enter), &error) != 1)
		return Fail(name + ": CG_ENTERGAME: " + error);
	Event(name + ":entergame");

	// CInputLogin::Entergame order: ch->Show inserts the main actor while the client is still in
	// the loading phase (CHARACTER::EncodeInsertPacket = ADD + ADDITIONAL_INFO, then UpdatePacket),
	// then SendNPCPosition, PHASE_GAME, SendLandList, TIME, CHANNEL and the greet notice.
	GCCharacterAdd actor = zeroed<GCCharacterAdd>();
	actor.header = HDR_GC_CHARACTER_ADD;
	actor.vid = kMainVID;
	actor.x = kX;
	actor.y = kY;
	actor.type = 0; // CHAR_TYPE_PC
	actor.race = 0; // warrior male
	actor.moving_speed = 100;
	actor.attack_speed = 100;
	GCCharAddInfo info = zeroed<GCCharAddInfo>();
	info.header = HDR_GC_CHAR_ADDITIONAL_INFO;
	info.vid = kMainVID;
	std::strncpy(info.name, kCharacterName, sizeof(info.name) - 1);
	info.empire = 1;
	info.level = 10;
	GCCharacterUpdate update = zeroed<GCCharacterUpdate>();
	update.header = HDR_GC_CHARACTER_UPDATE;
	update.vid = kMainVID;
	update.moving_speed = 100;
	update.attack_speed = 100;
	if (!c.Send(actor) || !c.Send(info) || !c.Send(update))
		return Fail(name + ": send main actor");

	// packet_npc_position = DynHead + count, then TNPCPosition entries.
	NPCPosition34 npc = zeroed<NPCPosition34>();
	npc.type = 1; // CHAR_TYPE_NPC
	std::strncpy(npc.name, "Npc", sizeof(npc.name) - 1);
	npc.x = kX + 500;
	npc.y = kY + 500;
	DynHead npc_head = {HDR_GC_NPC_POSITION, static_cast<std::uint16_t>(sizeof(DynHead) + sizeof(GCNPCPositionHead) + sizeof(npc))};
	GCNPCPositionHead npc_count = {1};
	if (!c.Send(npc_head) || !c.Send(npc_count) || !c.Send(npc))
		return Fail(name + ": send NPC positions");

	if (!c.Phase(PHASE_GAME))
		return Fail(name + ": send PHASE_GAME");

	LandPacketElement24 land = zeroed<LandPacketElement24>();
	land.id = 1;
	land.x = kX + 1000;
	land.y = kY + 1000;
	land.width = 2000;
	land.height = 2000;
	DynHead land_head = {HDR_GC_LAND_LIST, static_cast<std::uint16_t>(sizeof(DynHead) + sizeof(land))};
	GCTime game_time = {HDR_GC_TIME, static_cast<std::int32_t>(now_ms())};
	GCChannel channel = {HDR_GC_CHANNEL, 1};
	if (!c.Send(land_head) || !c.Send(land) || !c.Send(game_time) || !c.Send(channel))
		return Fail(name + ": send land/time/channel");

	// SendGreetMessage: a CHAT_TYPE_NOTICE line with no speaker.
	const char greet[] = "Welcome";
	GCChatHead chat = {HDR_GC_CHAT, static_cast<std::uint16_t>(sizeof(GCChatHead) + sizeof(greet)), 2 /* CHAT_TYPE_NOTICE */, 0, 1};
	if (!c.Send(chat) || !c.Send(greet, sizeof(greet)))
		return Fail(name + ": send greet notice");
	Event(name + ":burst_sent");

	// A ping after the burst proves that GamePhase consumed every packet above and can reply.
	Blank ping = {HDR_GC_PING};
	Blank pong = {};
	if (!c.Send(ping) || read_packet(c, HDR_CG_PONG, &pong, sizeof(pong), &error) != 1)
		return Fail(name + ": CG_PONG after the entergame burst: " + error);
	Event(name + ":burst_pong");

	// Hold the connection until the test stops us.
	std::uint8_t byte = 0;
	while (!m_stop && c.Read(&byte, 1) == 1)
		;
	return true;
}