// Minimal protocol-level login smoke test.
//
// Links directly against the server's own Src/Common/rosecrypt.cpp so the
// encryption is guaranteed bit-identical to what the real server/client
// speak (no separate reimplementation to get subtly wrong). The login
// server's crypt table is generated from a fixed constant
// (GenerateLoginServerCryptTables), so this test can compute the same
// table itself without any key exchange.
//
// Usage: rose_login_test <host> <port> <username> <password>
// Exit code 0 = login succeeded (account existed, or was auto-created and
// then logged in on retry). Non-zero = connection/protocol/login failure.

#include <cstdio>
#include <cstdlib>
#include <cstring>
#include <string>
#include <arpa/inet.h>
#include <sys/socket.h>
#include <netinet/in.h>
#include <unistd.h>

#include "../Src/Common/rosecrypt.hpp"

static bool send_packet(int sock, CCryptTable* table, unsigned short cmd,
                         const unsigned char* payload, unsigned short payload_len)
{
	unsigned char buffer[1024] = {0};
	unsigned short size = 6 + payload_len;
	*((unsigned short*)(buffer + 0)) = size;
	*((unsigned short*)(buffer + 2)) = cmd;
	if (payload_len)
		memcpy(buffer + 6, payload, payload_len);

	EncryptBuffer(table, buffer);

	ssize_t sent = send(sock, buffer, size, 0);
	return sent == (ssize_t)size;
}

// Returns command in *out_cmd, payload in out_payload (caller-sized buffer),
// payload length in *out_len. Returns false on any read/decrypt failure.
static bool recv_packet(int sock, CCryptTable* table, CCryptStatus* status,
                         unsigned short* out_cmd, unsigned char* out_payload,
                         unsigned short* out_len)
{
	unsigned char buffer[1024] = {0};
	ssize_t got = recv(sock, buffer, 6, MSG_WAITALL);
	if (got != 6) { fprintf(stderr, "short header read (%zd)\n", got); return false; }

	int fullsize = DecryptBufferHeader(status, table, buffer);
	if (fullsize < 6) { fprintf(stderr, "DecryptBufferHeader failed\n"); return false; }

	if (fullsize > 6)
	{
		got = recv(sock, buffer + 6, fullsize - 6, MSG_WAITALL);
		if (got != fullsize - 6) { fprintf(stderr, "short body read (%zd)\n", got); return false; }
	}

	if (!DecryptBufferData(table, buffer)) { fprintf(stderr, "DecryptBufferData failed (bad checksum)\n"); return false; }

	unsigned short size = *((unsigned short*)(buffer + 0));
	unsigned short cmd  = *((unsigned short*)(buffer + 2));
	*out_cmd = cmd;
	*out_len = size - 6;
	if (*out_len > 0)
		memcpy(out_payload, buffer + 6, *out_len);
	return true;
}

int main(int argc, char** argv)
{
	if (argc < 5)
	{
		fprintf(stderr, "usage: %s <host> <port> <username> <password>\n", argv[0]);
		return 2;
	}
	const char* host = argv[1];
	int port = atoi(argv[2]);
	std::string username = argv[3];
	std::string password = argv[4];

	CCryptTable* table = NULL;
	GenerateLoginServerCryptTables(table);
	CCryptStatus status = {0, 0};

	int sock = socket(AF_INET, SOCK_STREAM, 0);
	struct sockaddr_in addr = {0};
	addr.sin_family = AF_INET;
	addr.sin_port = htons(port);
	if (inet_pton(AF_INET, host, &addr.sin_addr) != 1)
	{
		fprintf(stderr, "bad host address '%s'\n", host);
		return 2;
	}
	if (connect(sock, (struct sockaddr*)&addr, sizeof(addr)) != 0)
	{
		perror("connect");
		return 2;
	}
	printf("[+] connected to %s:%d\n", host, port);

	// 1. Request encryption setup (0x703), empty payload.
	if (!send_packet(sock, table, 0x703, NULL, 0)) { fprintf(stderr, "send 0x703 failed\n"); return 1; }

	unsigned short cmd; unsigned char payload[1024]; unsigned short plen;
	if (!recv_packet(sock, table, &status, &cmd, payload, &plen)) { fprintf(stderr, "recv 0x7FF failed\n"); return 1; }
	if (cmd != 0x7FF) { fprintf(stderr, "expected 0x7FF, got %#04x\n", cmd); return 1; }
	printf("[+] encryption handshake ok (cmd=%#04x, %d byte payload)\n", cmd, plen);

	// 2. Login (0x708): 32-byte password, then up to 16-byte username.
	for (int attempt = 1; attempt <= 2; ++attempt)
	{
		unsigned char loginpayload[48] = {0};
		memcpy(loginpayload, password.c_str(), password.size() > 32 ? 32 : password.size());
		memcpy(loginpayload + 32, username.c_str(), username.size() > 16 ? 16 : username.size());

		if (!send_packet(sock, table, 0x708, loginpayload, sizeof(loginpayload)))
		{
			fprintf(stderr, "send 0x708 failed\n");
			return 1;
		}
		if (!recv_packet(sock, table, &status, &cmd, payload, &plen))
		{
			fprintf(stderr, "recv login response failed\n");
			return 1;
		}
		if (cmd != 0x708)
		{
			fprintf(stderr, "expected 0x708 response, got %#04x\n", cmd);
			return 1;
		}

		unsigned char result_byte = payload[0];
		if (result_byte == 0)
		{
			printf("[+] login SUCCESS for user '%s' (attempt %d), response payload = %d bytes\n",
			       username.c_str(), attempt, plen);
			if (plen > 5)
				printf("[+] server list present in response (channel(s) registered)\n");
			else
				printf("[!] login ok but no channel entries in response\n");
			close(sock);
			return 0;
		}
		else if (result_byte == 2 && attempt == 1)
		{
			printf("[i] account '%s' did not exist, auto-created (error code 2), retrying login...\n", username.c_str());
			continue;
		}
		else
		{
			fprintf(stderr, "[-] login FAILED, error code %d\n", result_byte);
			close(sock);
			return 1;
		}
	}

	close(sock);
	return 1;
}